merge to latest

Microsoft.Azure.Cosmos/src/AuthorizationHelper.cs

@@ -16,7 +16,7 @@ namespace Microsoft.Azure.Cosmos
     internal sealed class AuthorizationHelper
     {
         public const int MaxAuthorizationHeaderSize = 1024;
-
+        
         // This API is a helper method to create auth header based on client request.
         // Uri is split into resourceType/resourceId - 
         // For feed/post/put requests, resourceId = parentId,

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

@@ -152,22 +152,21 @@ public virtual bool TryAdd(ItemBatchOperation operation)
                 try
                 {
                     PartitionKeyRangeBatchExecutionResult result = await this.executor(serverRequest, cancellationToken);
-
-                    if (result.IsSplit())
-                    {
-                        foreach (ItemBatchOperation operationToRetry in result.Operations)
-                        {
-                            await this.retrier(operationToRetry, cancellationToken);
-                        }
-
-                        return;
-                    }
-
                     using (PartitionKeyRangeBatchResponse batchResponse = new PartitionKeyRangeBatchResponse(serverRequest.Operations.Count, result.ServerResponse, this.cosmosSerializer))
                     {
                         foreach (ItemBatchOperation itemBatchOperation in batchResponse.Operations)
                         {
                             BatchOperationResult response = batchResponse[itemBatchOperation.OperationIndex];
+                            if (!response.IsSuccessStatusCode)
+                            {
+                                Documents.ShouldRetryResult shouldRetry = await itemBatchOperation.Context.ShouldRetryAsync(response, cancellationToken);
+                                if (shouldRetry.ShouldRetry)
+                                {
+                                    await this.retrier(itemBatchOperation, cancellationToken);
+                                    continue;
+                                }
+                            }
+
                             itemBatchOperation.Context.Complete(this, response);
                         }
                     }

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

@@ -25,7 +25,7 @@ namespace Microsoft.Azure.Cosmos
     /// <seealso cref="BatchAsyncStreamer"/>
     internal class BatchAsyncContainerExecutor : IDisposable
     {
-        private const int DefaultDispatchTimer = 10;
+        private const int DefaultDispatchTimerInSeconds = 1;
         private const int MinimumDispatchTimerInSeconds = 1;
 
         private readonly ContainerCore cosmosContainer;
@@ -36,13 +36,21 @@ internal class BatchAsyncContainerExecutor : IDisposable
         private readonly ConcurrentDictionary<string, BatchAsyncStreamer> streamersByPartitionKeyRange = new ConcurrentDictionary<string, BatchAsyncStreamer>();
         private readonly ConcurrentDictionary<string, SemaphoreSlim> limitersByPartitionkeyRange = new ConcurrentDictionary<string, SemaphoreSlim>();
         private readonly TimerPool timerPool;
+        private readonly RetryOptions retryOptions;
+
+        /// <summary>
+        /// For unit testing.
+        /// </summary>
+        internal BatchAsyncContainerExecutor()
+        {
+        }
 
         public BatchAsyncContainerExecutor(
             ContainerCore cosmosContainer,
             CosmosClientContext cosmosClientContext,
             int maxServerRequestOperationCount,
             int maxServerRequestBodyLength,
-            int dispatchTimerInSeconds = BatchAsyncContainerExecutor.DefaultDispatchTimer)
+            int dispatchTimerInSeconds = BatchAsyncContainerExecutor.DefaultDispatchTimerInSeconds)
         {
             if (cosmosContainer == null)
             {
@@ -70,9 +78,10 @@ public BatchAsyncContainerExecutor(
             this.maxServerRequestOperationCount = maxServerRequestOperationCount;
             this.dispatchTimerInSeconds = dispatchTimerInSeconds;
             this.timerPool = new TimerPool(BatchAsyncContainerExecutor.MinimumDispatchTimerInSeconds);
+            this.retryOptions = cosmosClientContext.ClientOptions.GetConnectionPolicy().RetryOptions;
         }
 
-        public async Task<BatchOperationResult> AddAsync(
+        public virtual async Task<BatchOperationResult> AddAsync(
             ItemBatchOperation operation,
             ItemRequestOptions itemRequestOptions = null,
             CancellationToken cancellationToken = default(CancellationToken))
@@ -86,10 +95,10 @@ public async Task<BatchOperationResult> AddAsync(
 
             string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false);
             BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId);
-            ItemBatchOperationContext context = new ItemBatchOperationContext(resolvedPartitionKeyRangeId);
+            ItemBatchOperationContext context = new ItemBatchOperationContext(resolvedPartitionKeyRangeId, BatchAsyncContainerExecutor.GetRetryPolicy(this.retryOptions));
             operation.AttachContext(context);
             streamer.Add(operation);
-            return await context.Task;
+            return await context.OperationTask;
         }
 
         public void Dispose()
@@ -107,7 +116,7 @@ public void Dispose()
             this.timerPool.Dispose();
         }
 
-        internal async Task ValidateOperationAsync(
+        internal virtual async Task ValidateOperationAsync(
             ItemBatchOperation operation,
             ItemRequestOptions itemRequestOptions = null,
             CancellationToken cancellationToken = default(CancellationToken))
@@ -135,6 +144,14 @@ internal async Task ValidateOperationAsync(
             }
         }
 
+        private static IDocumentClientRetryPolicy GetRetryPolicy(RetryOptions retryOptions)
+        {
+            return new BulkPartitionKeyRangeGoneRetryPolicy(
+                new ResourceThrottleRetryPolicy(
+                retryOptions.MaxRetryAttemptsOnThrottledRequests,
+                retryOptions.MaxRetryWaitTimeInSeconds));
+        }
+
         private static bool ValidateOperationEPK(
             ItemBatchOperation operation,
             ItemRequestOptions itemRequestOptions)

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

@@ -7,9 +7,12 @@ namespace Microsoft.Azure.Cosmos
     using System;
     using System.Collections.Generic;
     using System.IO;
+    using System.Net;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Documents;
+    using Microsoft.Azure.Documents.Routing;
 
     /// <summary>
     /// Util methods for batch requests.
@@ -86,10 +89,50 @@ public static async Task<Memory<byte>> StreamToMemoryAsync(
             }
         }
 
+        public static string GetPartitionKeyRangeId(PartitionKeyInternal partitionKeyInternal, PartitionKeyDefinition partitionKeyDefinition, CollectionRoutingMap collectionRoutingMap)
+        {
+            string effectivePartitionKey = partitionKeyInternal.GetEffectivePartitionKeyString(partitionKeyDefinition);
+            return collectionRoutingMap.GetRangeByEffectivePartitionKey(effectivePartitionKey).Id;
+        }
+
+        public static void GetServerRequestLimits(out int maxServerRequestBodyLength, out int maxServerRequestOperationCount)
+        {
+            maxServerRequestBodyLength = Constants.MaxDirectModeBatchRequestBodySizeInBytes;
+            maxServerRequestOperationCount = Constants.MaxOperationsInDirectModeBatchRequest;
+        }
+
+        public static ResponseMessage EnsureValidStream(
+            IReadOnlyList<ItemBatchOperation> operations,
+            RequestOptions batchOptions,
+            int? maxOperationCount = null)
+        {
+            string errorMessage = BatchExecUtils.EnsureValidInternal(operations, batchOptions, maxOperationCount);
+
+            if (errorMessage != null)
+            {
+                return new ResponseMessage(HttpStatusCode.BadRequest, errorMessage: errorMessage);
+            }
+
+            return new ResponseMessage(HttpStatusCode.OK);
+        }
+
         public static void EnsureValid(
             IReadOnlyList<ItemBatchOperation> operations,
             RequestOptions batchOptions,
             int? maxOperationCount = null)
+        {
+            string errorMessage = BatchExecUtils.EnsureValidInternal(operations, batchOptions, maxOperationCount);
+
+            if (errorMessage != null)
+            {
+                throw new ArgumentException(errorMessage);
+            }
+        }
+
+        private static string EnsureValidInternal(
+            IReadOnlyList<ItemBatchOperation> operations,
+            RequestOptions batchOptions,
+            int? maxOperationCount = null)
         {
             string errorMessage = null;
 
@@ -130,18 +173,15 @@ public static void EnsureValid(
                                 WFConstants.BackendHeaders.EffectivePartitionKeyString);
                         }
 
-                        if (operation.PartitionKey != null)
+                        if (operation.PartitionKey != null && !operation.RequestOptions.IsEffectivePartitionKeyRouting)
                         {
                             errorMessage = ClientResources.PKAndEpkSetTogether;
                         }
                     }
                 }
             }
 
-            if (errorMessage != null)
-            {
-                throw new ArgumentException(errorMessage);
-            }
+            return errorMessage;
         }
 
         public static string GetPartitionKeyRangeId(PartitionKey partitionKey, PartitionKeyDefinition partitionKeyDefinition, Routing.CollectionRoutingMap collectionRoutingMap)

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

@@ -23,5 +23,22 @@ class BatchItemRequestOptions : RequestOptions
         /// <seealso cref="Microsoft.Azure.Cosmos.IndexingPolicy"/>
         /// <seealso cref="IndexingDirective"/>
         public IndexingDirective? IndexingDirective { get; set; }
+
+        internal static BatchItemRequestOptions FromItemRequestOptions(ItemRequestOptions itemRequestOptions)
+        {
+            if (itemRequestOptions == null)
+            {
+                return null;
+            }
+
+            RequestOptions requestOptions = itemRequestOptions as RequestOptions;
+            BatchItemRequestOptions batchItemRequestOptions = new BatchItemRequestOptions();
+            batchItemRequestOptions.IndexingDirective = itemRequestOptions.IndexingDirective;
+            batchItemRequestOptions.IfMatchEtag = requestOptions.IfMatchEtag;
+            batchItemRequestOptions.IfNoneMatchEtag = requestOptions.IfNoneMatchEtag;
+            batchItemRequestOptions.Properties = requestOptions.Properties;
+            batchItemRequestOptions.IsEffectivePartitionKeyRouting = requestOptions.IsEffectivePartitionKeyRouting;
+            return batchItemRequestOptions;
+        }
     }
 }

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

@@ -176,6 +176,16 @@ private static Result ReadOperationResult(ref RowReader reader, out BatchOperati
 
             return Result.Success;
         }
+
+        internal ResponseMessage ToResponseMessage()
+        {
+            ResponseMessage responseMessage = new ResponseMessage(this.StatusCode);
+            responseMessage.Headers.SubStatusCode = this.SubStatusCode;
+            responseMessage.Headers.ETag = this.ETag;
+            responseMessage.Headers.RetryAfter = this.RetryAfter;
+            responseMessage.Content = this.ResourceStream;
+            return responseMessage;
+        }
     }
 
     /// <summary>
@@ -194,7 +204,7 @@ class BatchOperationResult<T> : BatchOperationResult
         /// <summary>
         /// Initializes a new instance of the <see cref="BatchOperationResult{T}"/> class.
         /// </summary>
-        /// <param name="result">CosmosBatchOperationResult with stream resource.</param>
+        /// <param name="result">BatchOperationResult with stream resource.</param>
         /// <param name="resource">Deserialized resource.</param>
         internal BatchOperationResult(BatchOperationResult result, T resource)
             : base(result)

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

@@ -191,7 +191,11 @@ public virtual IEnumerator<BatchOperationResult> GetEnumerator()
         /// Gets all the Activity IDs associated with the response.
         /// </summary>
         /// <returns>An enumerable that contains the Activity IDs.</returns>
+#if INTERNAL
+        public virtual IEnumerable<string> GetActivityIds()
+#else
         internal virtual IEnumerable<string> GetActivityIds()
+#endif
         {
             yield return this.ActivityId;
         }

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

@@ -28,7 +28,7 @@ internal class ItemBatchOperation : IDisposable
         public ItemBatchOperation(
             OperationType operationType,
             int operationIndex,
-            PartitionKey? partitionKey,
+            PartitionKey partitionKey,
             string id = null,
             Stream resourceStream = null,
             BatchItemRequestOptions requestOptions = null)
@@ -55,7 +55,7 @@ public ItemBatchOperation(
             this.RequestOptions = requestOptions;
         }
 
-        public PartitionKey? PartitionKey { get; }
+        public PartitionKey? PartitionKey { get; internal set; }
 
         public string Id { get; }
 

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

@@ -5,8 +5,10 @@
 namespace Microsoft.Azure.Cosmos
 {
     using System;
+    using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Core.Trace;
+    using Microsoft.Azure.Documents;
 
     /// <summary>
     /// Context for a particular Batch operation.
@@ -17,13 +19,35 @@ internal class ItemBatchOperationContext : IDisposable
 
         public BatchAsyncBatcher CurrentBatcher { get; set; }
 
-        public Task<BatchOperationResult> Task => this.taskCompletionSource.Task;
+        public Task<BatchOperationResult> OperationTask => this.taskCompletionSource.Task;
+
+        private readonly IDocumentClientRetryPolicy retryPolicy;
 
         private TaskCompletionSource<BatchOperationResult> taskCompletionSource = new TaskCompletionSource<BatchOperationResult>();
 
-        public ItemBatchOperationContext(string partitionKeyRangeId)
+        public ItemBatchOperationContext(
+            string partitionKeyRangeId,
+            IDocumentClientRetryPolicy retryPolicy = null)
         {
             this.PartitionKeyRangeId = partitionKeyRangeId;
+            this.retryPolicy = retryPolicy;
+        }
+
+        /// <summary>
+        /// Based on the Retry Policy, if a failed response should retry.
+        /// </summary>
+        public Task<ShouldRetryResult> ShouldRetryAsync(
+            BatchOperationResult batchOperationResult,
+            CancellationToken cancellationToken)
+        {
+            if (this.retryPolicy == null
+                || batchOperationResult.IsSuccessStatusCode)
+            {
+                return Task.FromResult(ShouldRetryResult.NoRetry());
+            }
+
+            ResponseMessage responseMessage = batchOperationResult.ToResponseMessage();
+            return this.retryPolicy.ShouldRetryAsync(responseMessage, cancellationToken);
         }
 
         public void Complete(