Bulk API congestion control

Microsoft.Azure.Cosmos.Samples/Usage/BulkSupport/AppSettings.json

@@ -0,0 +1,13 @@
+{
+  "EndPointUrl": "https://localhost:8081",
+  "AuthorizationKey": "Super secret key",
+  "DatabaseName": "samples",
+  "ContainerName": "bulk-support",
+  "CollectionThroughput": "100000",
+  "ShouldCleanupOnStart": "false",
+  "ShouldCleanupOnFinish": "false",
+  "ItemSize": "400",
+  "ItemsToCreate": "10000",
+  "RuntimeInSeconds": "300",
+  "numWorkers": "1"
+}

Microsoft.Azure.Cosmos.Samples/Usage/BulkSupport/BulkSupport.csproj

@@ -1,11 +1,12 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
     <OutputType>Exe</OutputType>
     <TargetFramework>netcoreapp2.1</TargetFramework>
     <AssemblyName>Cosmos.Samples.BulkSupport</AssemblyName>
     <RootNamespace>Cosmos.Samples.BulkSupport</RootNamespace>
     <LangVersion>latest</LangVersion>
+    <ServerGarbageCollection>true</ServerGarbageCollection>
   </PropertyGroup>
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Configuration" Version="2.2.0" />
@@ -18,7 +19,7 @@
     <ProjectReference Include="..\Shared\Shared.csproj" />
   </ItemGroup>
   <ItemGroup>
-    <None Include="..\AppSettings.json">
+    <None Include="AppSettings.json">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
     </None>
   </ItemGroup>

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

@@ -6,9 +6,12 @@ namespace Microsoft.Azure.Cosmos
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
+    using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Core.Trace;
+    using Microsoft.Azure.Documents;
 
     /// <summary>
     /// Maintains a batch of operations and dispatches it as a unit of work.
@@ -117,7 +120,9 @@ public virtual bool TryAdd(ItemBatchOperation operation)
             return true;
         }
 
-        public virtual async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public virtual async Task DispatchAsync(
+            BatchPartitionMetric partitionMetric,
+            CancellationToken cancellationToken = default(CancellationToken))
         {
             this.interlockIncrementCheck.EnterLockCheck();
 
@@ -151,7 +156,16 @@ public virtual bool TryAdd(ItemBatchOperation operation)
 
                 try
                 {
+                    Stopwatch stopwatch = Stopwatch.StartNew();
+
                     PartitionKeyRangeBatchExecutionResult result = await this.executor(serverRequest, cancellationToken);
+
+                    int numThrottle = result.ServerResponse.Any(r => r.StatusCode == (System.Net.HttpStatusCode)StatusCodes.TooManyRequests) ? 1 : 0;
+                    partitionMetric.Add(
+                        numberOfDocumentsOperatedOn: result.ServerResponse.Count,
+                        timeTakenInMilliseconds: stopwatch.ElapsedMilliseconds,
+                        numberOfThrottles: numThrottle);
+
                     using (PartitionKeyRangeBatchResponse batchResponse = new PartitionKeyRangeBatchResponse(serverRequest.Operations.Count, result.ServerResponse, this.serializerCore))
                     {
                         foreach (ItemBatchOperation itemBatchOperation in batchResponse.Operations)

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

@@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos
     using System.Collections.Generic;
     using System.Diagnostics;
     using System.IO;
+    using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Diagnostics;
@@ -39,6 +40,9 @@ internal class BatchAsyncContainerExecutor : IDisposable
         private readonly TimerPool timerPool;
         private readonly RetryOptions retryOptions;
 
+        private readonly int defaultMaxDegreeOfConcurrency = 50;
+        private readonly bool adaptiveBulkExecution;
+
         /// <summary>
         /// For unit testing.
         /// </summary>
@@ -80,6 +84,7 @@ public BatchAsyncContainerExecutor(
             this.dispatchTimerInSeconds = dispatchTimerInSeconds;
             this.timerPool = new TimerPool(BatchAsyncContainerExecutor.MinimumDispatchTimerInSeconds);
             this.retryOptions = cosmosClientContext.ClientOptions.GetConnectionPolicy().RetryOptions;
+            this.adaptiveBulkExecution = cosmosClientContext.ClientOptions.AdaptiveBulkExecution;
         }
 
         public virtual async Task<TransactionalBatchOperationResult> AddAsync(
@@ -259,16 +264,18 @@ private BatchAsyncStreamer GetOrAddStreamerForPartitionKeyRange(string partition
             {
                 return streamer;
             }
-
+            SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(partitionKeyRangeId);
             BatchAsyncStreamer newStreamer = new BatchAsyncStreamer(
                 this.maxServerRequestOperationCount,
                 this.maxServerRequestBodyLength,
                 this.dispatchTimerInSeconds,
                 this.timerPool,
+                limiter,
+                this.adaptiveBulkExecution,
+                this.defaultMaxDegreeOfConcurrency,
                 this.cosmosClientContext.SerializerCore,
                 this.ExecuteAsync,
                 this.ReBatchAsync);
-
             if (!this.streamersByPartitionKeyRange.TryAdd(partitionKeyRangeId, newStreamer))
             {
                 newStreamer.Dispose();
@@ -284,7 +291,7 @@ private SemaphoreSlim GetOrAddLimiterForPartitionKeyRange(string partitionKeyRan
                 return limiter;
             }
 
-            SemaphoreSlim newLimiter = new SemaphoreSlim(1, 1);
+            SemaphoreSlim newLimiter = new SemaphoreSlim(1, this.defaultMaxDegreeOfConcurrency);
             if (!this.limitersByPartitionkeyRange.TryAdd(partitionKeyRangeId, newLimiter))
             {
                 newLimiter.Dispose();

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

@@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos
     using System;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Query.Core.Metrics;
     using Microsoft.Azure.Documents;
 
     /// <summary>
@@ -28,16 +29,34 @@ internal class BatchAsyncStreamer : IDisposable
         private readonly int dispatchTimerInSeconds;
         private readonly CosmosSerializerCore serializerCore;
         private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
+
+        private readonly int congestionIncreaseFactor = 1;
+        private readonly int congestionControllerDelayInSeconds = 1;
+        private readonly int congestionDecreaseFactor = 5;
+        private readonly int maxDegreeOfConcurrency;
+
         private volatile BatchAsyncBatcher currentBatcher;
         private TimerPool timerPool;
         private PooledTimer currentTimer;
         private Task timerTask;
 
+        private PooledTimer congestionControlTimer;
+        private Task congestionControlTask;
+        private SemaphoreSlim limiter;
+
+        private int congestionDegreeOfConcurrency = 1;
+        private long congestionWaitTimeInMilliseconds = 1000;
+        private BatchPartitionMetric oldPartitionMetric;
+        private BatchPartitionMetric partitionMetric;
+
         public BatchAsyncStreamer(
             int maxBatchOperationCount,
             int maxBatchByteSize,
             int dispatchTimerInSeconds,
             TimerPool timerPool,
+            SemaphoreSlim limiter,
+            bool adaptiveBulkExecution,
+            int maxDegreeOfConcurrency,
             CosmosSerializerCore serializerCore,
             BatchAsyncBatcherExecuteDelegate executor,
             BatchAsyncBatcherRetryDelegate retrier)
@@ -80,8 +99,17 @@ public BatchAsyncStreamer(
             this.timerPool = timerPool;
             this.serializerCore = serializerCore;
             this.currentBatcher = this.CreateBatchAsyncBatcher();
-
             this.ResetTimer();
+
+            this.limiter = limiter;
+            this.oldPartitionMetric = new BatchPartitionMetric();
+            this.partitionMetric = new BatchPartitionMetric();
+            this.maxDegreeOfConcurrency = maxDegreeOfConcurrency;
+
+            if (adaptiveBulkExecution)
+            {
+                this.StartCongestionControlTimer();
+            }
         }
 
         public void Add(ItemBatchOperation operation)
@@ -99,17 +127,25 @@ public void Add(ItemBatchOperation operation)
             if (toDispatch != null)
             {
                 // Discarded for Fire & Forget
-                _ = toDispatch.DispatchAsync(this.cancellationTokenSource.Token);
+                _ = toDispatch.DispatchAsync(this.partitionMetric, this.cancellationTokenSource.Token);
             }
         }
 
         public void Dispose()
         {
             this.cancellationTokenSource.Cancel();
             this.cancellationTokenSource.Dispose();
+
             this.currentTimer.CancelTimer();
             this.currentTimer = null;
             this.timerTask = null;
+
+            if (this.congestionControlTimer != null)
+            {
+                this.congestionControlTimer.CancelTimer();
+                this.congestionControlTimer = null;
+                this.congestionControlTask = null;
+            }
         }
 
         private void ResetTimer()
@@ -121,6 +157,15 @@ private void ResetTimer()
             }, this.cancellationTokenSource.Token);
         }
 
+        private void StartCongestionControlTimer()
+        {
+            this.congestionControlTimer = this.timerPool.GetPooledTimer(this.congestionControllerDelayInSeconds);
+            this.congestionControlTask = this.congestionControlTimer.StartTimerAsync().ContinueWith(async (task) =>
+            {
+                await this.RunCongestionControlAsync();
+            }, this.cancellationTokenSource.Token);
+        }
+
         private void DispatchTimer()
         {
             if (this.cancellationTokenSource.IsCancellationRequested)
@@ -137,7 +182,7 @@ private void DispatchTimer()
             if (toDispatch != null)
             {
                 // Discarded for Fire & Forget
-                _ = toDispatch.DispatchAsync(this.cancellationTokenSource.Token);
+                _ = toDispatch.DispatchAsync(this.partitionMetric, this.cancellationTokenSource.Token);
             }
 
             this.ResetTimer();
@@ -159,5 +204,53 @@ private BatchAsyncBatcher CreateBatchAsyncBatcher()
         {
             return new BatchAsyncBatcher(this.maxBatchOperationCount, this.maxBatchByteSize, this.serializerCore, this.executor, this.retrier);
         }
+
+        private async Task RunCongestionControlAsync()
+        {
+            while (!this.cancellationTokenSource.Token.IsCancellationRequested)
+            {
+                long elapsedTimeInMilliseconds = this.partitionMetric.TimeTakenInMilliseconds - this.oldPartitionMetric.TimeTakenInMilliseconds;
+
+                if (elapsedTimeInMilliseconds >= this.congestionWaitTimeInMilliseconds)
+                {
+                    long diffThrottle = this.partitionMetric.NumberOfThrottles - this.oldPartitionMetric.NumberOfThrottles;
+                    long changeDocCount = this.partitionMetric.NumberOfDocumentsOperatedOn - this.oldPartitionMetric.NumberOfDocumentsOperatedOn;
+                    this.oldPartitionMetric.Add(changeDocCount, elapsedTimeInMilliseconds, diffThrottle);
+
+                    if (diffThrottle > 0)
+                    {
+                        // Decrease should not lead to degreeOfConcurrency 0 as this will just block the thread here and no one would release it.
+                        int decreaseCount = Math.Min(this.congestionDecreaseFactor, this.congestionDegreeOfConcurrency / 2);
+
+                        // We got a throttle so we need to back off on the degree of concurrency.
+                        for (int i = 0; i < decreaseCount; i++)
+                        {
+                            await this.limiter.WaitAsync(this.cancellationTokenSource.Token);
+                        }
+
+                        this.congestionDegreeOfConcurrency -= decreaseCount;
+
+                        // In case of throttling increase the wait time, so as to converge max degreeOfConcurrency
+                        this.congestionWaitTimeInMilliseconds += 1000;
+                    }
+
+                    if (changeDocCount > 0 && diffThrottle == 0)
+                    {
+                        if (this.congestionDegreeOfConcurrency + this.congestionIncreaseFactor <= this.maxDegreeOfConcurrency)
+                        {
+                            // We aren't getting throttles, so we should bump up the degree of concurrency.
+                            this.limiter.Release(this.congestionIncreaseFactor);
+                            this.congestionDegreeOfConcurrency += this.congestionIncreaseFactor;
+                        }
+                    }
+                }
+                else
+                {
+                    break;
+                }
+            }
+
+            this.StartCongestionControlTimer();
+        }
     }
 }

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

@@ -0,0 +1,90 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+
+    internal sealed class BatchPartitionMetric
+    {
+        public BatchPartitionMetric()
+            : this(0, 0, 0)
+        {
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the OperationMetrics class (instance constructor).
+        /// </summary>
+        /// <param name="numberOfDocumentsOperatedOn">Number of documents operated on.</param>
+        /// <param name="timeTakenInMilliseconds">Amount of time taken to insert the documents.</param>
+        /// <param name="numberOfThrottles">The number of throttles encountered to insert the documents.</param>
+        public BatchPartitionMetric(long numberOfDocumentsOperatedOn, long timeTakenInMilliseconds, long numberOfThrottles)
+        {
+            if (numberOfDocumentsOperatedOn < 0)
+            {
+                throw new ArgumentException("numberOfDocumentsOperatedOn must be non negative");
+            }
+
+            if (timeTakenInMilliseconds < 0)
+            {
+                throw new ArgumentException("timeTakenInMilliseconds must be non negative");
+            }
+
+            if (numberOfThrottles < 0)
+            {
+                throw new ArgumentException("numberOfThrottles must be non negative");
+            }
+
+            this.NumberOfDocumentsOperatedOn = numberOfDocumentsOperatedOn;
+            this.TimeTakenInMilliseconds = timeTakenInMilliseconds;
+            this.NumberOfThrottles = numberOfThrottles;
+        }
+
+        /// <summary>
+        /// Gets the number of documents operated on.
+        /// </summary>
+        public long NumberOfDocumentsOperatedOn
+        {
+            get; private set;
+        }
+
+        /// <summary>
+        /// Gets the time taken to operate on the documents.
+        /// </summary>
+        public long TimeTakenInMilliseconds
+        {
+            get; private set;
+        }
+
+        /// <summary>
+        /// Gets the number of throttles incurred while operating on the documents.
+        /// </summary>
+        public long NumberOfThrottles
+        {
+            get; private set;
+        }
+
+        public void Add(long numberOfDocumentsOperatedOn, long timeTakenInMilliseconds, long numberOfThrottles)
+        {
+            if (numberOfDocumentsOperatedOn < 0)
+            {
+                throw new ArgumentException("numberOfDocumentsOperatedOn must be non negative");
+            }
+
+            if (timeTakenInMilliseconds < 0)
+            {
+                throw new ArgumentException("timeTakenInMilliseconds must be non negative");
+            }
+
+            if (numberOfThrottles < 0)
+            {
+                throw new ArgumentException("numberOfThrottles must be non negative");
+            }
+
+            this.NumberOfDocumentsOperatedOn += numberOfDocumentsOperatedOn;
+            this.TimeTakenInMilliseconds += timeTakenInMilliseconds;
+            this.NumberOfThrottles += numberOfThrottles;
+        }
+    }
+}