Bulk executor cache

Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutorCache.cs

@@ -0,0 +1,64 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+    using System.Collections.Concurrent;
+    using System.Collections.Generic;
+    using Microsoft.Azure.Documents;
+
+    /// <summary>
+    /// Cache to create and share Executor instances across the client's lifetime.
+    /// </summary>
+    internal class BatchAsyncContainerExecutorCache : IDisposable
+    {
+        private ConcurrentDictionary<string, BatchAsyncContainerExecutor> executorsPerContainer = new ConcurrentDictionary<string, BatchAsyncContainerExecutor>();
+
+        public BatchAsyncContainerExecutor GetExecutorForContainer(
+            ContainerCore container,
+            CosmosClientContext cosmosClientContext)
+        {
+            if (!cosmosClientContext.ClientOptions.AllowBulkExecution)
+            {
+                return null;
+            }
+
+            string containerLink = container.LinkUri.ToString();
+            if (this.executorsPerContainer.TryGetValue(containerLink, out BatchAsyncContainerExecutor executor))
+            {
+                return executor;
+            }
+
+            BatchAsyncContainerExecutor newExecutor = new BatchAsyncContainerExecutor(
+                container,
+                cosmosClientContext,
+                Constants.MaxOperationsInDirectModeBatchRequest,
+                Constants.MaxDirectModeBatchRequestBodySizeInBytes);
+            if (!this.executorsPerContainer.TryAdd(containerLink, newExecutor))
+            {
+                newExecutor.Dispose();
+            }
+
+            return this.executorsPerContainer[containerLink];
+        }
+
+        public void DisposeExecutor(ContainerCore container)
+        {
+            string containerLink = container.LinkUri.ToString();
+            if (this.executorsPerContainer.TryRemove(containerLink, out BatchAsyncContainerExecutor executor))
+            {
+                executor.Dispose();
+            }
+        }
+
+        public void Dispose()
+        {
+            foreach (KeyValuePair<string, BatchAsyncContainerExecutor> cacheEntry in this.executorsPerContainer)
+            {
+                cacheEntry.Value.Dispose();
+            }
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/CosmosClient.cs

@@ -296,6 +296,7 @@ internal CosmosClient(
         internal RequestInvokerHandler RequestHandler { get; private set; }
         internal CosmosResponseFactory ResponseFactory { get; private set; }
         internal CosmosClientContext ClientContext { get; private set; }
+        internal BatchAsyncContainerExecutorCache BatchExecutorCache { get; private set; } = new BatchAsyncContainerExecutorCache();
 
         /// <summary>
         /// Read Azure Cosmos DB account properties <see cref="Microsoft.Azure.Cosmos.AccountProperties"/>
@@ -743,6 +744,12 @@ protected virtual void Dispose(bool disposing)
                 this.DocumentClient.Dispose();
                 this.DocumentClient = null;
             }
+
+            if (this.BatchExecutorCache != null)
+            {
+                this.BatchExecutorCache.Dispose();
+                this.BatchExecutorCache = null;
+            }
         }
     }
 }

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs

@@ -157,15 +157,22 @@ internal async Task<ThroughputResponse> ReplaceThroughputIfExistsAsync(
                 cancellationToken: cancellationToken);
         }
 
-        public override Task<ResponseMessage> DeleteContainerStreamAsync(
+        public override async Task<ResponseMessage> DeleteContainerStreamAsync(
             ContainerRequestOptions requestOptions = null,
             CancellationToken cancellationToken = default(CancellationToken))
         {
-            return this.ProcessStreamAsync(
+            ResponseMessage responseMessage = await this.ProcessStreamAsync(
                streamPayload: null,
                operationType: OperationType.Delete,
                requestOptions: requestOptions,
                cancellationToken: cancellationToken);
+
+            if (responseMessage.IsSuccessStatusCode)
+            {
+                this.ClientContext.Client.BatchExecutorCache.DisposeExecutor(this);
+            }
+
+            return responseMessage;
         }
 
         public override Task<ResponseMessage> ReadContainerStreamAsync(
@@ -289,16 +296,7 @@ internal virtual Task<CollectionRoutingMap> GetRoutingMapAsync(CancellationToken
 
         internal virtual BatchAsyncContainerExecutor InitializeBatchExecutorForContainer()
         {
-            if (!this.ClientContext.ClientOptions.AllowBulkExecution)
-            {
-                return null;
-            }
-
-            return new BatchAsyncContainerExecutor(
-                this,
-                this.ClientContext,
-                Constants.MaxOperationsInDirectModeBatchRequest,
-                Constants.MaxDirectModeBatchRequestBodySizeInBytes);
+            return this.ClientContext.Client.BatchExecutorCache.GetExecutorForContainer(this, this.ClientContext);
         }
 
         private Task<ResponseMessage> ReplaceStreamInternalAsync(

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorCacheTests.cs

@@ -0,0 +1,138 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Tests
+{
+    using System;
+    using System.Collections.Generic;
+    using System.IO;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Documents;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using Moq;
+
+    [TestClass]    
+    public class BatchAsyncContainerExecutorCacheTests
+    {
+        [TestMethod]
+        public async Task ConcurrentGet_ReturnsSameExecutorInstance()
+        {
+            Mock<CosmosClient> mockClient = new Mock<CosmosClient>();
+            mockClient.Setup(x => x.Endpoint).Returns(new Uri("http://localhost"));
+
+            CosmosClientContext context = new ClientContextCore(
+                client: mockClient.Object,
+                clientOptions: new CosmosClientOptions() { AllowBulkExecution = true },
+                userJsonSerializer: null,
+                defaultJsonSerializer: null,
+                sqlQuerySpecSerializer: null,
+                cosmosResponseFactory: null,
+                requestHandler: null,
+                documentClient: null);
+
+            DatabaseCore db = new DatabaseCore(context, "test");
+
+            List<Task<ContainerCore>> tasks = new List<Task<ContainerCore>>();
+            for (int i = 0; i < 20; i++)
+            {
+                tasks.Add(Task.Run(() => Task.FromResult((ContainerCore)db.GetContainer("test"))));
+            }
+
+            await Task.WhenAll(tasks);
+
+            BatchAsyncContainerExecutor firstExecutor = tasks[0].Result.BatchExecutor;
+            Assert.IsNotNull(firstExecutor);
+            for (int i = 1; i < 20; i++)
+            {
+                BatchAsyncContainerExecutor otherExecutor = tasks[i].Result.BatchExecutor;
+                Assert.AreEqual(firstExecutor, otherExecutor);
+            }
+        }
+
+        [TestMethod]
+        public void Null_When_OptionsOff()
+        {
+            Mock<CosmosClient> mockClient = new Mock<CosmosClient>();
+            mockClient.Setup(x => x.Endpoint).Returns(new Uri("http://localhost"));
+
+            CosmosClientContext context = new ClientContextCore(
+                client: mockClient.Object,
+                clientOptions: new CosmosClientOptions() { },
+                userJsonSerializer: null,
+                defaultJsonSerializer: null,
+                sqlQuerySpecSerializer: null,
+                cosmosResponseFactory: null,
+                requestHandler: null,
+                documentClient: null);
+
+            DatabaseCore db = new DatabaseCore(context, "test");
+            ContainerCore container = (ContainerCore)db.GetContainer("test");
+            Assert.IsNull(container.BatchExecutor);
+        }
+
+        [TestMethod]
+        public void Get_And_Dispose()
+        {
+            Mock<CosmosClient> mockClient = new Mock<CosmosClient>();
+            mockClient.Setup(x => x.Endpoint).Returns(new Uri("http://localhost"));
+
+            CosmosClientContext context = new ClientContextCore(
+                client: mockClient.Object,
+                clientOptions: new CosmosClientOptions() { AllowBulkExecution = true },
+                userJsonSerializer: null,
+                defaultJsonSerializer: null,
+                sqlQuerySpecSerializer: null,
+                cosmosResponseFactory: null,
+                requestHandler: null,
+                documentClient: null);
+
+            DatabaseCore db = new DatabaseCore(context, "test");
+            ContainerCore container = (ContainerCore)db.GetContainer("test");
+
+            BatchAsyncContainerExecutorCache cache = new BatchAsyncContainerExecutorCache();
+
+            BatchAsyncContainerExecutor executor = cache.GetExecutorForContainer(container, context);
+            Assert.IsNotNull(executor);
+
+            cache.DisposeExecutor(container);
+
+            // Should return new instance
+            BatchAsyncContainerExecutor secondExecutor = cache.GetExecutorForContainer(container, context);
+            Assert.AreNotEqual(executor, secondExecutor);
+        }
+
+        [TestMethod]
+        public async Task DeleteContainerRemovesCache()
+        {
+            Mock<CosmosClient> mockClient = new Mock<CosmosClient>();
+            mockClient.Setup(x => x.Endpoint).Returns(new Uri("http://localhost"));
+
+            Mock<CosmosClientContext> mockedContext = new Mock<CosmosClientContext>();
+            mockedContext.Setup(c => c.Client).Returns(mockClient.Object);
+            mockedContext.Setup(c => c.ClientOptions).Returns(new CosmosClientOptions() { AllowBulkExecution = true });
+            mockedContext.Setup(c => c.CreateLink(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>())).Returns(new Uri("/dbs/test/colls/test", UriKind.Relative));
+            mockedContext
+                .Setup(c => c.ProcessResourceOperationStreamAsync(
+                    It.IsAny<Uri>(),
+                    It.IsAny<ResourceType>(),
+                    It.IsAny<OperationType>(),
+                    It.IsAny<RequestOptions>(),
+                    It.IsAny<ContainerCore>(),
+                    It.IsAny<Cosmos.PartitionKey?>(),
+                    It.IsAny<Stream>(),
+                    It.IsAny<Action<RequestMessage>>(),
+                    It.IsAny<CancellationToken>()))
+                .Returns(Task.FromResult(new ResponseMessage(System.Net.HttpStatusCode.NoContent)));
+
+            DatabaseCore db = new DatabaseCore(mockedContext.Object, "test");
+            ContainerCore container = (ContainerCore)db.GetContainer("test");
+            Assert.AreEqual(container.BatchExecutor, mockClient.Object.BatchExecutorCache.GetExecutorForContainer(container, mockedContext.Object));
+            await container.DeleteContainerStreamAsync();
+
+            // Asking for a new cache instance should return a new executor as the previous entry should have been deleted
+            Assert.AreNotEqual(container.BatchExecutor, mockClient.Object.BatchExecutorCache.GetExecutorForContainer(container, mockedContext.Object));
+        }
+    }
+}

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosClientResourceUnitTests.cs

@@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Core.Tests
     using System.Net.Http;
     using Microsoft.Azure.Documents;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using Moq;
 
     [TestClass]
     public class CosmosClientResourceUnitTests
@@ -19,8 +20,11 @@ public void ValidateUriGenerationForResources()
             string databaseId = "db1234";
             string crId = "cr42";
 
+            Mock<CosmosClient> mockClient = new Mock<CosmosClient>();
+            mockClient.Setup(x => x.Endpoint).Returns(new Uri("http://localhost"));
+
             CosmosClientContext context = new ClientContextCore(
-                client: null,
+                client: mockClient.Object,
                 clientOptions: new CosmosClientOptions(),
                 userJsonSerializer: null,
                 defaultJsonSerializer: null,
@@ -120,8 +124,11 @@ public void InitializeBatchExecutorForContainer_Null_WhenAllowBulk_False()
             string databaseId = "db1234";
             string crId = "cr42";
 
+            Mock<CosmosClient> mockClient = new Mock<CosmosClient>();
+            mockClient.Setup(x => x.Endpoint).Returns(new Uri("http://localhost"));
+
             CosmosClientContext context = new ClientContextCore(
-                client: null,
+                client: mockClient.Object,
                 clientOptions: new CosmosClientOptions(),
                 userJsonSerializer: null,
                 defaultJsonSerializer: null,
@@ -141,8 +148,11 @@ public void InitializeBatchExecutorForContainer_NotNull_WhenAllowBulk_True()
             string databaseId = "db1234";
             string crId = "cr42";
 
+            Mock<CosmosClient> mockClient = new Mock<CosmosClient>();
+            mockClient.Setup(x => x.Endpoint).Returns(new Uri("http://localhost"));
+
             CosmosClientContext context = new ClientContextCore(
-                client: null,
+                client: mockClient.Object,
                 clientOptions: new CosmosClientOptions() { AllowBulkExecution = true },
                 userJsonSerializer: null,
                 defaultJsonSerializer: null,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/DotNetSDKAPI.json

@@ -1192,7 +1192,7 @@
           "Attributes": [],
           "MethodInfo": "System.Threading.Tasks.Task`1[Microsoft.Azure.Cosmos.DatabaseResponse] CreateDatabaseAsync(System.String, System.Nullable`1[System.Int32], Microsoft.Azure.Cosmos.RequestOptions, System.Threading.CancellationToken)"
         },
-        "System.Threading.Tasks.Task`1[Microsoft.Azure.Cosmos.DatabaseResponse] CreateDatabaseIfNotExistsAsync(System.String, System.Nullable`1[System.Int32], Microsoft.Azure.Cosmos.RequestOptions, System.Threading.CancellationToken)[System.Runtime.CompilerServices.AsyncStateMachineAttribute(typeof(Microsoft.Azure.Cosmos.CosmosClient+<CreateDatabaseIfNotExistsAsync>d__37))]": {
+        "System.Threading.Tasks.Task`1[Microsoft.Azure.Cosmos.DatabaseResponse] CreateDatabaseIfNotExistsAsync(System.String, System.Nullable`1[System.Int32], Microsoft.Azure.Cosmos.RequestOptions, System.Threading.CancellationToken)[System.Runtime.CompilerServices.AsyncStateMachineAttribute(typeof(Microsoft.Azure.Cosmos.CosmosClient+<CreateDatabaseIfNotExistsAsync>d__41))]": {
           "Type": "Method",
           "Attributes": [
             "AsyncStateMachineAttribute"

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj

@@ -224,9 +224,6 @@
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
     </Content>
   </ItemGroup>
-  <ItemGroup>
-    <Folder Include="Batch\" />
-  </ItemGroup>
 
   <PropertyGroup>
     <SignAssembly>true</SignAssembly>