From 813f193bc0b8fc812d28ecb118fd50edbf0c0b14 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 22 Jul 2019 20:58:11 -0700 Subject: [PATCH 01/41] migrating base classes --- .../src/Batch/BatchAsyncBatcher.cs | 148 +++++++ .../src/Batch/BatchAsyncContainerExecutor.cs | 418 ++++++++++++++++++ .../src/Batch/BatchAsyncOperationContext.cs | 42 ++ .../src/Batch/BatchAsyncStreamer.cs | 187 ++++++++ .../src/Batch/BatchExecUtils.cs | 8 + .../src/Batch/BatchResponse.cs | 6 +- .../Batch/CrossPartitionKeyBatchResponse.cs | 181 ++++++++ .../CrossPartitionKeyServerBatchRequest.cs | 62 +++ .../src/Batch/ItemBatchOperation.cs | 2 +- .../PartitionKeyRangeBatchExecutionResult.cs | 24 + 10 files changed, 1074 insertions(+), 4 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs create mode 100644 Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs create mode 100644 Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs create mode 100644 Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs create mode 100644 Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyBatchResponse.cs create mode 100644 Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyServerBatchRequest.cs create mode 100644 Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs new file mode 100644 index 0000000000..0564c2a525 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -0,0 +1,148 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + + /// + /// Maintains a batch of operations and dispatches the batch through an executor. Maps results into the original operation contexts. + /// + /// + internal class BatchAsyncBatcher : IDisposable + { + private readonly SemaphoreSlim tryAddLimiter; + private readonly CosmosSerializer CosmosSerializer; + private readonly List batchOperations; + private readonly Func, CancellationToken, Task> executor; + private readonly int maxBatchByteSize; + private readonly int maxBatchOperationCount; + private long currentSize = 0; + + public bool IsEmpty => this.batchOperations.Count == 0; + + public BatchAsyncBatcher( + int maxBatchOperationCount, + int maxBatchByteSize, + CosmosSerializer cosmosSerializer, + Func, CancellationToken, Task> executor) + { + if (maxBatchOperationCount < 1) + { + throw new ArgumentOutOfRangeException(nameof(maxBatchOperationCount)); + } + + if (maxBatchByteSize < 1) + { + throw new ArgumentOutOfRangeException(nameof(maxBatchByteSize)); + } + + if (executor == null) + { + throw new ArgumentNullException(nameof(executor)); + } + + if (cosmosSerializer == null) + { + throw new ArgumentNullException(nameof(cosmosSerializer)); + } + + this.batchOperations = new List(maxBatchOperationCount); + this.tryAddLimiter = new SemaphoreSlim(1, 1); + this.executor = executor; + this.maxBatchByteSize = maxBatchByteSize; + this.maxBatchOperationCount = maxBatchOperationCount; + this.CosmosSerializer = cosmosSerializer; + } + + public bool TryAdd(BatchAsyncOperationContext batchAsyncOperation) + { + if (batchAsyncOperation == null) + { + throw new ArgumentNullException(nameof(batchAsyncOperation)); + } + + this.tryAddLimiter.Wait(); + try + { + if (this.batchOperations.Count == this.maxBatchOperationCount) + { + return false; + } + +#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits + batchAsyncOperation.Operation.MaterializeResourceAsync(this.CosmosSerializer, default(CancellationToken)).GetAwaiter().GetResult(); +#pragma warning restore VSTHRD002 // Avoid problematic synchronous waits + + int itemByteSize = batchAsyncOperation.Operation.GetApproximateSerializedLength(); + + if (itemByteSize + this.currentSize > this.maxBatchByteSize) + { + return false; + } + + this.currentSize += itemByteSize; + + // Operation index is in the scope of the current batch + batchAsyncOperation.Operation.OperationIndex = this.batchOperations.Count; + this.batchOperations.Add(batchAsyncOperation); + return true; + } + finally + { + this.tryAddLimiter.Release(); + } + } + + public async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) + { + await this.tryAddLimiter.WaitAsync(cancellationToken); + + try + { + CrossPartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken); + + // If the batch was not successful, we need to set all the responses + if (!batchResponse.IsSuccessStatusCode) + { + BatchOperationResult errorResult = new BatchOperationResult(batchResponse.StatusCode); + foreach (BatchAsyncOperationContext operation in this.batchOperations) + { + operation.Complete(errorResult); + } + + return; + } + + for (int index = 0; index < this.batchOperations.Count; index++) + { + BatchAsyncOperationContext operation = this.batchOperations[index]; + BatchOperationResult response = batchResponse[index]; + operation.Complete(response); + } + } + catch (Exception ex) + { + // Exceptions happening during execution fail all the Tasks + foreach (BatchAsyncOperationContext operation in this.batchOperations) + { + operation.Fail(ex); + } + } + finally + { + this.tryAddLimiter.Release(); + this.Dispose(); + } + } + + public void Dispose() + { + this.tryAddLimiter?.Dispose(); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs new file mode 100644 index 0000000000..582bdddf38 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -0,0 +1,418 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Collections.Concurrent; + using System.Collections.Generic; + using System.Diagnostics; + using System.IO; + using System.Linq; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Internal; + using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Documents; + + /// + /// Bulk batch executor for operations in the same container. + /// + /// + /// It groups operations by Partition Key Range and sends them to the Batch API and then unifies results as they become available. It uses as batch processor and as batch executing handler. + /// + /// + internal class BatchAsyncContainerExecutor : IDisposable + { + private const int DefaultDispatchTimer = 10; + private const int MinimumDispatchTimerInSeconds = 1; + + private readonly ContainerCore cosmosContainer; + private readonly CosmosClientContext cosmosClientContext; + private readonly int maxServerRequestBodyLength; + private readonly int maxServerRequestOperationCount; + private readonly int dispatchTimerInSeconds; + private readonly ConcurrentDictionary streamersByPartitionKeyRange = new ConcurrentDictionary(); + private readonly ConcurrentDictionary limitersByPartitionkeyRange = new ConcurrentDictionary(); + private readonly TimerPool timerPool; + + public BatchAsyncContainerExecutor( + ContainerCore cosmosContainer, + CosmosClientContext cosmosClientContext, + int maxServerRequestOperationCount, + int maxServerRequestBodyLength, + int dispatchTimerInSeconds = BatchAsyncContainerExecutor.DefaultDispatchTimer) + { + if (cosmosContainer == null) + { + throw new ArgumentNullException(nameof(cosmosContainer)); + } + + if (maxServerRequestOperationCount < 1) + { + throw new ArgumentOutOfRangeException(nameof(maxServerRequestOperationCount)); + } + + if (maxServerRequestBodyLength < 1) + { + throw new ArgumentOutOfRangeException(nameof(maxServerRequestBodyLength)); + } + + if (dispatchTimerInSeconds < 1) + { + throw new ArgumentOutOfRangeException(nameof(dispatchTimerInSeconds)); + } + + this.cosmosContainer = cosmosContainer; + this.cosmosClientContext = cosmosClientContext; + this.maxServerRequestBodyLength = maxServerRequestBodyLength; + this.maxServerRequestOperationCount = maxServerRequestOperationCount; + this.dispatchTimerInSeconds = dispatchTimerInSeconds; + this.timerPool = new TimerPool(BatchAsyncContainerExecutor.MinimumDispatchTimerInSeconds); + } + + public async Task AddAsync( + ItemBatchOperation operation, + CancellationToken cancellationToken) + { + if (operation == null) + { + throw new ArgumentNullException(nameof(operation)); + } + + string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken); + BatchAsyncStreamer streamer = this.GetStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId, operation); + Debug.Assert(streamer != null, "Could not obtain streamer"); + return await streamer.AddAsync(new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation)); + } + + public async Task ValidateOperationAsync(ItemBatchOperation operation) + { + if (operation.RequestOptions != null) + { + if (operation.RequestOptions.BaseConsistencyLevel.HasValue) + { + return false; + } + + if (operation.RequestOptions.Properties != null + && (operation.RequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKey, out object epkObj) + | operation.RequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkStrObj))) + { + byte[] epk = epkObj as byte[]; + string epkStr = epkStrObj as string; + if (epk == null || epkStr == null) + { + throw new InvalidOperationException(string.Format( + ClientResources.EpkPropertiesPairingExpected, + WFConstants.BackendHeaders.EffectivePartitionKey, + WFConstants.BackendHeaders.EffectivePartitionKeyString)); + } + + if (operation.PartitionKey != null) + { + throw new InvalidOperationException(ClientResources.PKAndEpkSetTogether); + } + } + } + + await operation.MaterializeResourceAsync(this.cosmosClientContext.CosmosSerializer, default(CancellationToken)); + + int itemByteSize = operation.GetApproximateSerializedLength(); + + if (itemByteSize > this.maxServerRequestBodyLength) + { + throw new InvalidOperationException(RMResources.RequestTooLarge); + } + + return true; + } + + public void Dispose() + { + foreach (KeyValuePair streamer in this.streamersByPartitionKeyRange) + { + streamer.Value.Dispose(); + } + + foreach (KeyValuePair limiter in this.limitersByPartitionkeyRange) + { + limiter.Value.Dispose(); + } + + this.timerPool.Dispose(); + } + + private async Task ExecuteAsync( + IReadOnlyList operations, + CancellationToken cancellationToken) + { + List> inProgressTasksByPKRangeId = new List>(); + List completedResults = new List(); + + // Initial operations are all for the same PK Range Id + string initialPartitionKeyRangeId = operations[0].PartitionKeyRangeId; + inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(initialPartitionKeyRangeId, operations, cancellationToken)); + + while (inProgressTasksByPKRangeId.Count > 0) + { + Task completedTask = await Task.WhenAny(inProgressTasksByPKRangeId); + PartitionKeyRangeBatchExecutionResult executionResult = await completedTask; + completedResults.Add(executionResult); + + // Pending operations can be due to splits, so we need to resolve the PK Range Ids + if (executionResult.PendingOperations != null) + { + IEnumerable retryOperations = BatchAsyncContainerExecutor.GetOperationsToRetry(operations, executionResult.PendingOperations.Select(o => o.OperationIndex)); + await this.GroupOperationsAndStartExecutionAsync(inProgressTasksByPKRangeId, retryOperations, cancellationToken); + executionResult.PendingOperations = null; + } + + inProgressTasksByPKRangeId.Remove(completedTask); + } + + IEnumerable serverResponses = completedResults.SelectMany(res => res.ServerResponses); + return new CrossPartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); + } + + private static IEnumerable GetOperationsToRetry( + IReadOnlyList operations, + IEnumerable indexes) + { + foreach (int index in indexes) + { + yield return operations[index]; + } + } + + private static void AddHeadersToRequestMessage(RequestMessage requestMessage, string partitionKeyRangeId) + { + requestMessage.Headers.Add(WFConstants.BackendHeaders.PartitionKeyRangeId, partitionKeyRangeId); + requestMessage.Headers.Add(HttpConstants.HttpHeaders.ShouldBatchContinueOnError, bool.TrueString); + requestMessage.Headers.Add(BatchExecUtils.IsBatchRequest, bool.TrueString); + } + + /// + /// If because of HybridRow serialization overhead, not all operations fit in the request, we send those extra operations in a separate request. + /// + private static IEnumerable GetOverflowOperations( + CrossPartitionKeyServerBatchRequest request, + IEnumerable operationsSentToRequest) + { + int totalOperations = operationsSentToRequest.Count(); + int operationsThatOverflowed = totalOperations - request.Operations.Count; + if (operationsThatOverflowed == 0) + { + return Enumerable.Empty(); + } + + return operationsSentToRequest.Skip(totalOperations - operationsThatOverflowed); + } + + private async Task GroupOperationsAndStartExecutionAsync( + List> taskContainer, + IEnumerable operations, + CancellationToken cancellation) + { + Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, cancellation); + + foreach (KeyValuePair> operationsForPKRangeId in operationsByPKRangeId) + { + Task toAdd = this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellation); + taskContainer.Add(toAdd); + } + } + + private async Task>> GroupByPartitionKeyRangeIdsAsync( + IEnumerable operations, + CancellationToken cancellationToken) + { + PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken); + CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken); + + Dictionary> operationsByPKRangeId = new Dictionary>(); + foreach (BatchAsyncOperationContext operation in operations) + { + string partitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation.Operation, cancellationToken); + + if (operationsByPKRangeId.TryGetValue(partitionKeyRangeId, out List operationsForPKRangeId)) + { + operationsForPKRangeId.Add(operation); + } + else + { + operationsByPKRangeId.Add(partitionKeyRangeId, new List() { operation }); + } + } + + return operationsByPKRangeId; + } + + private async Task ResolvePartitionKeyRangeIdAsync( + ItemBatchOperation operation, + CancellationToken cancellationToken) + { + PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken); + CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken); + + string partitionKeyRangeId; + object epkObj = null; + if (operation.RequestOptions?.Properties?.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out epkObj) ?? false) + { + string epk = epkObj as string; + Debug.Assert(epk != null, "Expected non-null string for " + WFConstants.BackendHeaders.EffectivePartitionKeyString); + partitionKeyRangeId = collectionRoutingMap.GetRangeByEffectivePartitionKey(epk).Id; + } + else + { + partitionKeyRangeId = BatchExecUtils.GetPartitionKeyRangeId(operation.ParsedPartitionKey, partitionKeyDefinition, collectionRoutingMap); + } + + return partitionKeyRangeId; + } + + private async Task StartExecutionForPKRangeIdAsync( + string pkRangeId, + IEnumerable operations, + CancellationToken cancellationToken) + { + SemaphoreSlim limiter = this.GetLimiterForPartitionKeyRange(pkRangeId); + await limiter.WaitAsync(cancellationToken); + + try + { + return await this.ExecuteServerRequestsAsync(pkRangeId, operations.Select(o => o.Operation), cancellationToken); + } + finally + { + limiter.Release(); + } + } + + private async Task CreateServerRequestAsync( + string partitionKeyRangeId, + IEnumerable operations, + CancellationToken cancellationToken) + { + ArraySegment operationsArraySegment = new ArraySegment(operations.ToArray()); + + CrossPartitionKeyServerBatchRequest request = await CrossPartitionKeyServerBatchRequest.CreateAsync( + partitionKeyRangeId, + operationsArraySegment, + this.maxServerRequestBodyLength, + this.maxServerRequestOperationCount, + ensureContinuousOperationIndexes: false, + serializer: this.cosmosClientContext.CosmosSerializer, + cancellationToken: cancellationToken); + + return request; + } + + private async Task ExecuteServerRequestsAsync( + string partitionKeyRangeId, + IEnumerable operations, + CancellationToken cancellationToken) + { + List serverResponses = new List(); + List overflowOperations = new List(); + CrossPartitionKeyServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operations, cancellationToken); + + // In case some operations overflowed + overflowOperations.AddRange(BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operations)); + do + { + BatchResponse serverResponse; + try + { + serverResponse = await this.ExecuteServerRequestAsync(serverRequest, cancellationToken); + } + catch (CosmosException ex) + { + serverResponse = new BatchResponse(ex.StatusCode, (SubStatusCodes)ex.SubStatusCode, ex.Message, serverRequest.Operations); + } + + if (serverResponse.StatusCode == HttpStatusCode.Gone + && (serverResponse.SubStatusCode == SubStatusCodes.CompletingSplit + || serverResponse.SubStatusCode == SubStatusCodes.CompletingPartitionMigration + || serverResponse.SubStatusCode == SubStatusCodes.PartitionKeyRangeGone)) + { + // lower layers would have refreshed the appropriate routing caches, but we need to retry the operations on the new PKRanges. + return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, serverResponses, new List(serverRequest.Operations)); + } + else + { + serverResponses.Add(serverResponse); + } + + serverRequest = null; + if (overflowOperations.Count > 0) + { + serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, overflowOperations, cancellationToken); + overflowOperations.Clear(); + } + } + while (serverRequest != null); + + return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, serverResponses); + } + + private async Task ExecuteServerRequestAsync(CrossPartitionKeyServerBatchRequest serverRequest, CancellationToken cancellationToken) + { + using (Stream serverRequestPayload = serverRequest.TransferBodyStream()) + { + Debug.Assert(serverRequestPayload != null, "Server request payload expected to be non-null"); + + ResponseMessage responseMessage = await ExecUtils.ProcessResourceOperationAsync( + this.cosmosClientContext.Client, + this.cosmosContainer.LinkUri, + ResourceType.Document, + OperationType.Batch, + new RequestOptions(), + cosmosContainerCore: this.cosmosContainer, + partitionKey: null, + streamPayload: serverRequestPayload, + requestEnricher: requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest.PartitionKeyRangeId), + responseCreator: cosmosResponseMessage => cosmosResponseMessage, // response creator + cancellationToken: cancellationToken); + + return await BatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.CosmosSerializer); + } + } + + private BatchAsyncStreamer GetStreamerForPartitionKeyRange( + string partitionKeyRangeId, + ItemBatchOperation operation) + { + if (this.streamersByPartitionKeyRange.TryGetValue(partitionKeyRangeId, out BatchAsyncStreamer streamer)) + { + return streamer; + } + + BatchAsyncStreamer newStreamer = new BatchAsyncStreamer(this.maxServerRequestOperationCount, this.maxServerRequestBodyLength, this.dispatchTimerInSeconds, this.timerPool, this.cosmosClientContext.CosmosSerializer, this.ExecuteAsync); + if (!this.streamersByPartitionKeyRange.TryAdd(partitionKeyRangeId, newStreamer)) + { + newStreamer.Dispose(); + } + + return this.streamersByPartitionKeyRange[partitionKeyRangeId]; + } + + private SemaphoreSlim GetLimiterForPartitionKeyRange(string partitionKeyRangeId) + { + if (this.limitersByPartitionkeyRange.TryGetValue(partitionKeyRangeId, out SemaphoreSlim limiter)) + { + return limiter; + } + + SemaphoreSlim newLimiter = new SemaphoreSlim(1, 1); + if (!this.limitersByPartitionkeyRange.TryAdd(partitionKeyRangeId, newLimiter)) + { + newLimiter.Dispose(); + } + + return this.limitersByPartitionkeyRange[partitionKeyRangeId]; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs new file mode 100644 index 0000000000..db147e5538 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs @@ -0,0 +1,42 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Threading; + using System.Threading.Tasks; + + /// + /// Context for a particular Batch operation. + /// + internal class BatchAsyncOperationContext + { + public string PartitionKeyRangeId { get; } + + public ItemBatchOperation Operation { get; } + + public Task Task => this.taskCompletionSource.Task; + + private TaskCompletionSource taskCompletionSource = new TaskCompletionSource(); + + public BatchAsyncOperationContext( + string partitionKeyRangeId, + ItemBatchOperation operation) + { + this.Operation = operation; + this.PartitionKeyRangeId = partitionKeyRangeId; + } + + public void Complete(BatchOperationResult result) + { + this.taskCompletionSource.SetResult(result); + } + + public void Fail(Exception exception) + { + this.taskCompletionSource.SetException(exception); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs new file mode 100644 index 0000000000..32cc45ae07 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -0,0 +1,187 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Collections.Generic; + using System.Diagnostics; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Internal; + using Microsoft.Azure.Documents; + + /// + /// Handles operation queueing and dispatching. + /// + /// + /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. + /// + /// + internal class BatchAsyncStreamer : IDisposable + { + private readonly List previousDispatchedTasks = new List(); + private readonly SemaphoreSlim addLimiter; + private readonly int maxBatchOperationCount; + private readonly int maxBatchByteSize; + private readonly Func, CancellationToken, Task> executor; + private readonly int dispatchTimerInSeconds; + private readonly CosmosSerializer CosmosSerializer; + private BatchAsyncBatcher currentBatcher; + private TimerPool timerPool; + private PooledTimer currentTimer; + private Task timerTask; + private bool disposed; + + public BatchAsyncStreamer( + int maxBatchOperationCount, + int maxBatchByteSize, + int dispatchTimerInSeconds, + TimerPool timerPool, + CosmosSerializer cosmosSerializer, + Func, CancellationToken, Task> executor) + { + if (maxBatchOperationCount < 1) + { + throw new ArgumentOutOfRangeException(nameof(maxBatchOperationCount)); + } + + if (maxBatchByteSize < 1) + { + throw new ArgumentOutOfRangeException(nameof(maxBatchByteSize)); + } + + if (dispatchTimerInSeconds < 1) + { + throw new ArgumentOutOfRangeException(nameof(dispatchTimerInSeconds)); + } + + if (executor == null) + { + throw new ArgumentNullException(nameof(executor)); + } + + if (CosmosSerializer == null) + { + throw new ArgumentNullException(nameof(CosmosSerializer)); + } + + this.maxBatchOperationCount = maxBatchOperationCount; + this.maxBatchByteSize = maxBatchByteSize; + this.executor = executor; + this.dispatchTimerInSeconds = dispatchTimerInSeconds; + this.timerPool = timerPool; + this.CosmosSerializer = cosmosSerializer; + this.addLimiter = new SemaphoreSlim(1, 1); + this.currentBatcher = this.GetBatchAsyncBatcher(); + + this.StartTimer(); + } + + public Task AddAsync(BatchAsyncOperationContext context) + { + BatchAsyncBatcher toDispatch = null; +#pragma warning disable VSTHRD103 // Call async methods when in an async method + this.addLimiter.Wait(); +#pragma warning restore VSTHRD103 // Call async methods when in an async method + + try + { + if (!this.currentBatcher.TryAdd(context)) + { + // Batcher is full + toDispatch = this.GetBatchToDispatch(); + Debug.Assert(this.currentBatcher.TryAdd(context), "Could not add context to batcher."); + } + } + finally + { + this.addLimiter.Release(); + } + + if (toDispatch != null) + { + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync()); + } + + return context.Task; + } + + public void Dispose() + { + this.disposed = true; + this.addLimiter?.Dispose(); + this.currentBatcher?.Dispose(); + + foreach (Task previousDispatch in this.previousDispatchedTasks) + { + try + { +#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits + previousDispatch.GetAwaiter().GetResult(); +#pragma warning restore VSTHRD002 // Avoid problematic synchronous waits + } + catch + { + // Any internal exceptions are disregarded during dispose + } + } + + this.currentTimer.CancelTimer(); + } + + private void StartTimer() + { + this.currentTimer = this.timerPool.GetPooledTimer(this.dispatchTimerInSeconds); + this.timerTask = this.currentTimer.StartTimerAsync().ContinueWith((task) => + { + this.DispatchTimer(); + }); + } + + private void DispatchTimer() + { + if (this.disposed) + { + return; + } + + this.addLimiter.Wait(); + + BatchAsyncBatcher toDispatch; + try + { + toDispatch = this.GetBatchToDispatch(); + } + finally + { + this.addLimiter.Release(); + } + + if (toDispatch != null) + { + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync()); + } + + this.StartTimer(); + } + + private BatchAsyncBatcher GetBatchToDispatch() + { + if (this.currentBatcher.IsEmpty) + { + return null; + } + + BatchAsyncBatcher previousBatcher = this.currentBatcher; + this.currentBatcher = this.GetBatchAsyncBatcher(); + return previousBatcher; + } + + private BatchAsyncBatcher GetBatchAsyncBatcher() + { + return new BatchAsyncBatcher(this.maxBatchOperationCount, this.maxBatchByteSize, this.CosmosSerializer, this.executor); + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs index 530439992e..cde64e8c76 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs @@ -17,6 +17,8 @@ namespace Microsoft.Azure.Cosmos /// internal static class BatchExecUtils { + internal const int StatusCodeMultiStatus = 207; + internal const string IsBatchRequest = "x-ms-cosmos-is-batch-request"; // Using the same buffer size as the Stream.DefaultCopyBufferSize private const int BufferSize = 81920; @@ -144,5 +146,11 @@ public static void EnsureValid( throw new ArgumentException(errorMessage); } } + + public static string GetPartitionKeyRangeId(Documents.PartitionKey partitionKey, PartitionKeyDefinition partitionKeyDefinition, Routing.CollectionRoutingMap collectionRoutingMap) + { + string effectivePartitionKey = partitionKey.InternalKey.GetEffectivePartitionKeyString(partitionKeyDefinition); + return collectionRoutingMap.GetRangeByEffectivePartitionKey(effectivePartitionKey).Id; + } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchResponse.cs index 8647df8226..608fe80d40 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchResponse.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchResponse.cs @@ -104,7 +104,7 @@ private BatchResponse( /// /// The request charge measured in request units. /// - public virtual double RequestCharge { get; } + public virtual double RequestCharge { get; internal set; } /// /// Gets the amount of time to wait before retrying this or any other request within Cosmos container or collection due to throttling. @@ -115,13 +115,13 @@ private BatchResponse( /// Gets the completion status code of the batch request. /// /// The request completion status code. - public virtual HttpStatusCode StatusCode { get; } + public virtual HttpStatusCode StatusCode { get; internal set; } /// /// Gets the reason for failure of the batch request. /// /// The reason for failure, if any. - public virtual string ErrorMessage { get; } + public virtual string ErrorMessage { get; internal set; } /// /// Gets a value indicating whether the batch was processed. diff --git a/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyBatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyBatchResponse.cs new file mode 100644 index 0000000000..973fc84d0e --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyBatchResponse.cs @@ -0,0 +1,181 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Net; + using System.Text; + using Microsoft.Azure.Documents; + + /// + /// Response of a cross partition key batch request. + /// +#pragma warning disable CA1710 // Identifiers should have correct suffix + #if PREVIEW + public +#else + internal +#endif + class CrossPartitionKeyBatchResponse : BatchResponse +#pragma warning restore CA1710 // Identifiers should have correct suffix + { + // Results sorted in the order operations had been added. + private readonly SortedList resultsByOperationIndex; + private bool isDisposed; + + /// + /// Initializes a new instance of the class. + /// + /// Completion status code of the batch request. + /// Provides further details about why the batch was not processed. + /// Operations that were supposed to be executed, but weren't. + /// The reason for failure if any. + // This constructor is expected to be used when the batch is not executed at all (if it is a bad request). + internal CrossPartitionKeyBatchResponse( + HttpStatusCode statusCode, + SubStatusCodes subStatusCode, + string errorMessage, + IReadOnlyList operations) + : base(statusCode, subStatusCode, errorMessage, operations) + { + } + + /// + /// Initializes a new instance of the class. + /// + /// Responses from the server. + /// Serializer to deserialize response resource body streams. + internal CrossPartitionKeyBatchResponse(IEnumerable serverResponses, CosmosSerializer serializer) + { + this.StatusCode = serverResponses.Any(r => r.StatusCode != HttpStatusCode.OK) + ? (HttpStatusCode)BatchExecUtils.StatusCodeMultiStatus + : HttpStatusCode.OK; + + this.ServerResponses = serverResponses; + this.resultsByOperationIndex = new SortedList(); + + StringBuilder errorMessageBuilder = new StringBuilder(); + List activityIds = new List(); + List itemBatchOperations = new List(); + foreach (BatchResponse serverResponse in serverResponses) + { + // We expect number of results == number of operations here + for (int index = 0; index < serverResponse.Operations.Count; index++) + { + int operationIndex = serverResponse.Operations[index].OperationIndex; + if (!this.resultsByOperationIndex.ContainsKey(operationIndex) + || this.resultsByOperationIndex[operationIndex].StatusCode == (HttpStatusCode)StatusCodes.TooManyRequests) + { + this.resultsByOperationIndex[operationIndex] = serverResponse[index]; + } + } + + itemBatchOperations.AddRange(serverResponse.Operations); + this.RequestCharge += serverResponse.RequestCharge; + + if (!string.IsNullOrEmpty(serverResponse.ErrorMessage)) + { + errorMessageBuilder.AppendFormat("{0}; ", serverResponse.ErrorMessage); + } + + activityIds.Add(serverResponse.ActivityId); + } + + this.ActivityIds = activityIds; + this.ErrorMessage = errorMessageBuilder.Length > 2 ? errorMessageBuilder.ToString(0, errorMessageBuilder.Length - 2) : null; + this.Operations = itemBatchOperations; + this.Serializer = serializer; + } + + /// + /// Gets the ActivityIds that identify the server requests made to execute the batch request. + /// +#pragma warning disable CA1721 // Property names should not match get methods + public virtual IEnumerable ActivityIds { get; } +#pragma warning restore CA1721 // Property names should not match get methods + + /// + public override string ActivityId => this.ActivityIds.First(); + + internal override CosmosSerializer Serializer { get; } + + // for unit testing only + internal IEnumerable ServerResponses { get; private set; } + + /// + /// Gets the number of operation results. + /// + public override int Count => this.resultsByOperationIndex.Count; + + /// + public override BatchOperationResult this[int index] => this.resultsByOperationIndex[index]; + + /// + /// Gets the result of the operation at the provided index in the batch - the returned result has a Resource of provided type. + /// + /// Type to which the Resource in the operation result needs to be deserialized to, when present. + /// 0-based index of the operation in the batch whose result needs to be returned. + /// Result of batch operation that contains a Resource deserialized to specified type. + public override BatchOperationResult GetOperationResultAtIndex(int index) + { + if (index >= this.Count) + { + throw new IndexOutOfRangeException(); + } + + BatchOperationResult result = this.resultsByOperationIndex[index]; + + T resource = default(T); + if (result.ResourceStream != null) + { + resource = this.Serializer.FromStream(result.ResourceStream); + } + + return new BatchOperationResult(result, resource); + } + + /// + /// Gets an enumerator over the operation results. + /// + /// Enumerator over the operation results. + public override IEnumerator GetEnumerator() + { + foreach (KeyValuePair pair in this.resultsByOperationIndex) + { + yield return pair.Value; + } + } + + internal override IEnumerable GetActivityIds() + { + return this.ActivityIds; + } + + /// + /// Disposes the disposable members held. + /// + /// Indicates whether to dispose managed resources or not. + protected override void Dispose(bool disposing) + { + if (disposing && !this.isDisposed) + { + this.isDisposed = true; + if (this.ServerResponses != null) + { + foreach (BatchResponse response in this.ServerResponses) + { + response.Dispose(); + } + + this.ServerResponses = null; + } + } + + base.Dispose(disposing); + } + } +} \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyServerBatchRequest.cs b/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyServerBatchRequest.cs new file mode 100644 index 0000000000..e337df17cd --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyServerBatchRequest.cs @@ -0,0 +1,62 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Threading; + using System.Threading.Tasks; + + internal sealed class CrossPartitionKeyServerBatchRequest : ServerBatchRequest + { + /// + /// Initializes a new instance of the class. + /// + /// The partition key range id associated with all requests. + /// Maximum length allowed for the request body. + /// Maximum number of operations allowed in the request. + /// Serializer to serialize user provided objects to JSON. + public CrossPartitionKeyServerBatchRequest( + string partitionKeyRangeId, + int maxBodyLength, + int maxOperationCount, + CosmosSerializer serializer) + : base(maxBodyLength, maxOperationCount, serializer) + { + this.PartitionKeyRangeId = partitionKeyRangeId; + } + + /// + /// Gets the PartitionKeyRangeId that applies to all operations in this request. + /// + public string PartitionKeyRangeId { get; } + + /// + /// Creates an instance of . + /// In case of direct mode requests, all the operations are expected to belong to the same PartitionKeyRange. + /// The body of the request is populated with operations till it reaches the provided maxBodyLength. + /// + /// The partition key range id associated with all requests. + /// Operations to be added into this batch request. + /// Desired maximum length of the request body. + /// Maximum number of operations allowed in the request. + /// Whether to stop adding operations to the request once there is non-continuity in the operation indexes. + /// Serializer to serialize user provided objects to JSON. + /// representing request cancellation. + /// A newly created instance of . + public static async Task CreateAsync( + string partitionKeyRangeId, + ArraySegment operations, + int maxBodyLength, + int maxOperationCount, + bool ensureContinuousOperationIndexes, + CosmosSerializer serializer, + CancellationToken cancellationToken) + { + CrossPartitionKeyServerBatchRequest request = new CrossPartitionKeyServerBatchRequest(partitionKeyRangeId, maxBodyLength, maxOperationCount, serializer); + await request.CreateBodyStreamAsync(operations, cancellationToken, ensureContinuousOperationIndexes); + return request; + } + } +} diff --git a/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs index 8bfbd35d93..05402387fd 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs @@ -65,7 +65,7 @@ public ItemBatchOperation( public BatchItemRequestOptions RequestOptions { get; } - public int OperationIndex { get; } + public int OperationIndex { get; internal set; } internal string PartitionKeyJson { get; set; } diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs new file mode 100644 index 0000000000..0dc19135d1 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs @@ -0,0 +1,24 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System.Collections.Generic; + + internal class PartitionKeyRangeBatchExecutionResult + { + public string PartitionKeyRangeId { get; } + + public List ServerResponses { get; } + + public List PendingOperations { get; set; } + + public PartitionKeyRangeBatchExecutionResult(string pkRangeId, List serverResponses, List pendingOperations = null) + { + this.PartitionKeyRangeId = pkRangeId; + this.ServerResponses = serverResponses; + this.PendingOperations = pendingOperations; + } + } +} \ No newline at end of file From 34f765396370dce1b5de5e9d52e76c29aa9d65b6 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 22 Jul 2019 21:09:58 -0700 Subject: [PATCH 02/41] Adding tests --- .../src/Batch/BatchAsyncStreamer.cs | 4 +- .../Batch/BatchAsyncBatcherTests.cs | 214 ++++++++++++++++++ .../Batch/BatchAsyncOperationContextTests.cs | 72 ++++++ .../Batch/BatchAsyncStreamerTests.cs | 155 +++++++++++++ .../CrossPartitionKeyBatchResponseTests.cs | 67 ++++++ 5 files changed, 510 insertions(+), 2 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 32cc45ae07..2e49794023 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -62,9 +62,9 @@ public BatchAsyncStreamer( throw new ArgumentNullException(nameof(executor)); } - if (CosmosSerializer == null) + if (cosmosSerializer == null) { - throw new ArgumentNullException(nameof(CosmosSerializer)); + throw new ArgumentNullException(nameof(cosmosSerializer)); } this.maxBatchOperationCount = maxBatchOperationCount; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs new file mode 100644 index 0000000000..0ab3462c12 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -0,0 +1,214 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Collections.Generic; + using System.IO; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] +#pragma warning disable CA1001 // Types that own disposable fields should be disposable + public class BatchAsyncBatcherTests +#pragma warning restore CA1001 // Types that own disposable fields should be disposable + { + private ItemBatchOperation ItemBatchOperation = new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); + + private static Exception expectedException = new Exception(); + + // Executor just returns a reponse matching the Id with Etag + private Func, CancellationToken, Task> Executor + = async (IReadOnlyList operations, CancellationToken cancellation) => + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operations.Count]; + int index = 0; + foreach (BatchAsyncOperationContext operation in operations) + { + results.Add( + new BatchOperationResult(HttpStatusCode.OK) + { + ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), + ETag = operation.Operation.Id + }); + + arrayOperations[index++] = operation.Operation; + } + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: (int)responseContent.Length * operations.Count, + maxOperationCount: operations.Count, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: cancellation); + + BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( + new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }, + batchRequest, + new CosmosJsonDotNetSerializer()); + + CrossPartitionKeyBatchResponse response = new CrossPartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + return response; + }; + + private Func, CancellationToken, Task> ExecutorWithFailure + = (IReadOnlyList operations, CancellationToken cancellation) => + { + throw expectedException; + }; + + [DataTestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ArgumentOutOfRangeException))] + [DataRow(0)] + [DataRow(-1)] + public void ValidatesSize(int size) + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(size, 1, new CosmosJsonDotNetSerializer(), this.Executor); + } + + [DataTestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ArgumentOutOfRangeException))] + [DataRow(0)] + [DataRow(-1)] + public void ValidatesByteSize(int size) + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, size, new CosmosJsonDotNetSerializer(), this.Executor); + } + + [TestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ArgumentNullException))] + public void ValidatesExecutor() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1, new CosmosJsonDotNetSerializer(), null); + } + + [TestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ArgumentNullException))] + public void ValidatesSerializer() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1, null, this.Executor); + } + + [TestMethod] + [Owner("maquaran")] + public void HasFixedSize() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + } + + [TestMethod] + [Owner("maquaran")] + public void HasFixedByteSize() + { + // Each operation is 2 bytes + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + } + + [TestMethod] + [Owner("maquaran")] + public void TryAddIsThreadSafe() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + Task firstOperation = Task.Run(() => batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Task secondOperation = Task.Run(() => batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Task thirdOperation = Task.Run(() => batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + + Task.WhenAll(firstOperation, secondOperation, thirdOperation).GetAwaiter().GetResult(); + + int countSucceded = (firstOperation.Result ? 1 : 0) + (secondOperation.Result ? 1 : 0) + (thirdOperation.Result ? 1 : 0); + int countFailed = (!firstOperation.Result ? 1 : 0) + (!secondOperation.Result ? 1 : 0) + (!thirdOperation.Result ? 1 : 0); + + Assert.AreEqual(2, countSucceded); + Assert.AreEqual(1, countFailed); + } + + [TestMethod] + [Owner("maquaran")] + public async Task ExceptionsFailOperationsAsync() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); + BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); + BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); + batchAsyncBatcher.TryAdd(context1); + batchAsyncBatcher.TryAdd(context2); + await batchAsyncBatcher.DispatchAsync(); + + Assert.AreEqual(TaskStatus.Faulted, context1.Task.Status); + Assert.AreEqual(TaskStatus.Faulted, context2.Task.Status); + Assert.AreEqual(expectedException, context1.Task.Exception.InnerException); + Assert.AreEqual(expectedException, context2.Task.Exception.InnerException); + } + + [TestMethod] + [Owner("maquaran")] + public async Task DispatchProcessInOrderAsync() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + List contexts = new List(10); + for (int i = 0; i < 10; i++) + { + BatchAsyncOperationContext context = new BatchAsyncOperationContext(string.Empty, new ItemBatchOperation(OperationType.Create, i, i.ToString())); + contexts.Add(context); + Assert.IsTrue(batchAsyncBatcher.TryAdd(context)); + } + + await batchAsyncBatcher.DispatchAsync(); + + for (int i = 0; i < 10; i++) + { + BatchAsyncOperationContext context = contexts[i]; + Assert.AreEqual(TaskStatus.RanToCompletion, context.Task.Status); + BatchOperationResult result = await context.Task; + Assert.AreEqual(i.ToString(), result.ETag); + } + } + + [TestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ObjectDisposedException))] + public async Task CannotDispatchTwiceAsync() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation)); + + await batchAsyncBatcher.DispatchAsync(); + await batchAsyncBatcher.DispatchAsync(); + } + + [TestMethod] + [Owner("maquaran")] + public void IsEmptyWithNoOperations() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + Assert.IsTrue(batchAsyncBatcher.IsEmpty); + } + + [TestMethod] + [Owner("maquaran")] + public void IsNotEmptyWithOperations() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsFalse(batchAsyncBatcher.IsEmpty); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs new file mode 100644 index 0000000000..3fb5e69128 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs @@ -0,0 +1,72 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Net; + using System.Threading.Tasks; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class BatchAsyncOperationContextTests + { + [TestMethod] + [Owner("maquaran")] + public void PartitionKeyRangeIdIsSetOnInitialization() + { + string expectedPkRangeId = Guid.NewGuid().ToString(); + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); + BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(expectedPkRangeId, operation); + + Assert.IsNotNull(batchAsyncOperationContext.Task); + Assert.AreEqual(operation, batchAsyncOperationContext.Operation); + Assert.AreEqual(expectedPkRangeId, batchAsyncOperationContext.PartitionKeyRangeId); + Assert.AreEqual(TaskStatus.WaitingForActivation, batchAsyncOperationContext.Task.Status); + } + + [TestMethod] + [Owner("maquaran")] + public void TaskIsCreatedOnInitialization() + { + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); + BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(string.Empty, operation); + + Assert.IsNotNull(batchAsyncOperationContext.Task); + Assert.AreEqual(operation, batchAsyncOperationContext.Operation); + Assert.AreEqual(TaskStatus.WaitingForActivation, batchAsyncOperationContext.Task.Status); + } + + [TestMethod] + [Owner("maquaran")] + public async Task TaskResultIsSetOnCompleteAsync() + { + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); + BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(string.Empty, operation); + + BatchOperationResult expected = new BatchOperationResult(HttpStatusCode.OK); + + batchAsyncOperationContext.Complete(expected); + + Assert.AreEqual(expected, await batchAsyncOperationContext.Task); + Assert.AreEqual(TaskStatus.RanToCompletion, batchAsyncOperationContext.Task.Status); + } + + [TestMethod] + [Owner("maquaran")] + public async Task ExceptionIsSetOnFailAsync() + { + Exception failure = new Exception("It failed"); + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); + BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(string.Empty, operation); + + batchAsyncOperationContext.Fail(failure); + + Exception capturedException = await Assert.ThrowsExceptionAsync(() => batchAsyncOperationContext.Task); + Assert.AreEqual(failure, capturedException); + Assert.AreEqual(TaskStatus.Faulted, batchAsyncOperationContext.Task.Status); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs new file mode 100644 index 0000000000..606ccb4d58 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -0,0 +1,155 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Collections.Generic; + using System.IO; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] +#pragma warning disable CA1001 // Types that own disposable fields should be disposable + public class BatchAsyncStreamerTests +#pragma warning restore CA1001 // Types that own disposable fields should be disposable + { + private const int DispatchTimerInSeconds = 5; + private const int MaxBatchByteSize = 100000; + private static Exception expectedException = new Exception(); + private ItemBatchOperation ItemBatchOperation = new ItemBatchOperation(OperationType.Create, 0, "0"); + private TimerPool TimerPool = new TimerPool(1); + + // Executor just returns a reponse matching the Id with Etag + private Func, CancellationToken, Task> Executor + = async (IReadOnlyList operations, CancellationToken cancellation) => + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operations.Count]; + int index = 0; + foreach (BatchAsyncOperationContext operation in operations) + { + results.Add( + new BatchOperationResult(HttpStatusCode.OK) + { + ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), + ETag = operation.Operation.Id + }); + + arrayOperations[index++] = operation.Operation; + } + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: (int)responseContent.Length * operations.Count, + maxOperationCount: operations.Count, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: cancellation); + + BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( + new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }, + batchRequest, + new CosmosJsonDotNetSerializer()); + + CrossPartitionKeyBatchResponse response = new CrossPartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + return response; + }; + + private Func, CancellationToken, Task> ExecutorWithFailure + = (IReadOnlyList operations, CancellationToken cancellation) => + { + throw expectedException; + }; + + [DataTestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ArgumentOutOfRangeException))] + [DataRow(0)] + [DataRow(-1)] + public void ValidatesSize(int size) + { + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(size, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + } + + [DataTestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ArgumentOutOfRangeException))] + [DataRow(0)] + [DataRow(-1)] + public void ValidatesDispatchTimer(int dispatchTimerInSeconds) + { + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, dispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + } + + [TestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ArgumentNullException))] + public void ValidatesExecutor() + { + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), null); + } + + [TestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(ArgumentNullException))] + public void ValidatesSerializer() + { + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, null, this.Executor); + } + + [TestMethod] + [Owner("maquaran")] + public async Task ExceptionsOnBatchBubbleUpAsync() + { + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); + + Exception capturedException = await Assert.ThrowsExceptionAsync(() => batchAsyncStreamer.AddAsync(CreateContext(this.ItemBatchOperation))); + + Assert.AreEqual(expectedException, capturedException); + } + + [TestMethod] + [Owner("maquaran")] + public async Task TimerDispatchesAsync() + { + // Bigger batch size than the amount of operations, timer should dispatch + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + + BatchOperationResult result = await batchAsyncStreamer.AddAsync(CreateContext(this.ItemBatchOperation)); + + Assert.AreEqual(this.ItemBatchOperation.Id, result.ETag); + } + + [TestMethod] + [Owner("maquaran")] + public async Task DispatchesAsync() + { + // Expect all operations to complete as their batches get dispached + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + List> contexts = new List>(10); + for (int i = 0; i < 10; i++) + { + contexts.Add(batchAsyncStreamer.AddAsync(CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())))); + } + + await Task.WhenAll(contexts); + + for (int i = 0; i < 10; i++) + { + Task context = contexts[i]; + Assert.AreEqual(TaskStatus.RanToCompletion, context.Status); + BatchOperationResult result = await context; + Assert.AreEqual(i.ToString(), result.ETag); + } + } + + private static BatchAsyncOperationContext CreateContext(ItemBatchOperation operation) => new BatchAsyncOperationContext(string.Empty, operation); + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs new file mode 100644 index 0000000000..5e64e7368a --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs @@ -0,0 +1,67 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Collections.Generic; + using System.IO; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class CrossPartitionKeyBatchResponseTests + { + [TestMethod] + [Owner("maquaran")] + public void StatusCodesAreSet() + { + const string errorMessage = "some error"; + CrossPartitionKeyBatchResponse response = new CrossPartitionKeyBatchResponse(HttpStatusCode.NotFound, SubStatusCodes.ClientTcpChannelFull, errorMessage, null); + Assert.AreEqual(HttpStatusCode.NotFound, response.StatusCode); + Assert.AreEqual(SubStatusCodes.ClientTcpChannelFull, response.SubStatusCode); + Assert.AreEqual(errorMessage, response.ErrorMessage); + } + + [TestMethod] + [Owner("maquaran")] + public async Task StatusCodesAreSetThroughResponseAsync() + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; + + ItemBatchOperation operation = new ItemBatchOperation(OperationType.AddComputeGatewayRequestCharges, 0, "0"); + + results.Add( + new BatchOperationResult(HttpStatusCode.OK) + { + ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), + ETag = operation.Id + }); + + arrayOperations[0] = operation; + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: 100, + maxOperationCount: 1, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: default(CancellationToken)); + + BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( + new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }, + batchRequest, + new CosmosJsonDotNetSerializer()); + + CrossPartitionKeyBatchResponse response = new CrossPartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); + } + } +} From 643c082722465bbd6a49cb5061f4cb202b667404 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Tue, 23 Jul 2019 07:37:43 -0700 Subject: [PATCH 03/41] Adding executor tests --- .../src/Batch/BatchAsyncContainerExecutor.cs | 27 +++-- .../src/Batch/BatchExecUtils.cs | 3 +- .../src/Batch/ItemBatchOperation.cs | 1 + .../src/Handler/RequestMessage.cs | 4 +- .../Batch/BatchAsyncContainerExecutorTests.cs | 109 ++++++++++++++++++ 5 files changed, 129 insertions(+), 15 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 582bdddf38..ee9520ea24 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -88,18 +88,23 @@ public async Task AddAsync( return await streamer.AddAsync(new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation)); } - public async Task ValidateOperationAsync(ItemBatchOperation operation) + public async Task ValidateOperationAsync( + ItemBatchOperation operation, + ItemRequestOptions itemRequestOptions = null) { - if (operation.RequestOptions != null) + if (itemRequestOptions != null) { - if (operation.RequestOptions.BaseConsistencyLevel.HasValue) + if (itemRequestOptions.BaseConsistencyLevel.HasValue + || itemRequestOptions.PreTriggers != null + || itemRequestOptions.PostTriggers != null + || itemRequestOptions.SessionToken != null) { return false; } - if (operation.RequestOptions.Properties != null - && (operation.RequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKey, out object epkObj) - | operation.RequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkStrObj))) + if (itemRequestOptions.Properties != null + && (itemRequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKey, out object epkObj) + | itemRequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkStrObj))) { byte[] epk = epkObj as byte[]; string epkStr = epkStrObj as string; @@ -189,9 +194,9 @@ private static IEnumerable GetOperationsToRetry( private static void AddHeadersToRequestMessage(RequestMessage requestMessage, string partitionKeyRangeId) { - requestMessage.Headers.Add(WFConstants.BackendHeaders.PartitionKeyRangeId, partitionKeyRangeId); + requestMessage.Headers.PartitionKeyRangeId = partitionKeyRangeId; requestMessage.Headers.Add(HttpConstants.HttpHeaders.ShouldBatchContinueOnError, bool.TrueString); - requestMessage.Headers.Add(BatchExecUtils.IsBatchRequest, bool.TrueString); + requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchRequest, bool.TrueString); } /// @@ -267,7 +272,7 @@ private async Task ResolvePartitionKeyRangeIdAsync( } else { - partitionKeyRangeId = BatchExecUtils.GetPartitionKeyRangeId(operation.ParsedPartitionKey, partitionKeyDefinition, collectionRoutingMap); + partitionKeyRangeId = BatchExecUtils.GetPartitionKeyRangeId(operation.PartitionKey.Value, partitionKeyDefinition, collectionRoutingMap); } return partitionKeyRangeId; @@ -364,8 +369,7 @@ private async Task ExecuteServerRequestAsync(CrossPartitionKeySer { Debug.Assert(serverRequestPayload != null, "Server request payload expected to be non-null"); - ResponseMessage responseMessage = await ExecUtils.ProcessResourceOperationAsync( - this.cosmosClientContext.Client, + ResponseMessage responseMessage = await this.cosmosClientContext.ProcessResourceOperationStreamAsync( this.cosmosContainer.LinkUri, ResourceType.Document, OperationType.Batch, @@ -374,7 +378,6 @@ private async Task ExecuteServerRequestAsync(CrossPartitionKeySer partitionKey: null, streamPayload: serverRequestPayload, requestEnricher: requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest.PartitionKeyRangeId), - responseCreator: cosmosResponseMessage => cosmosResponseMessage, // response creator cancellationToken: cancellationToken); return await BatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.CosmosSerializer); diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs index cde64e8c76..d4171db71c 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs @@ -18,7 +18,6 @@ namespace Microsoft.Azure.Cosmos internal static class BatchExecUtils { internal const int StatusCodeMultiStatus = 207; - internal const string IsBatchRequest = "x-ms-cosmos-is-batch-request"; // Using the same buffer size as the Stream.DefaultCopyBufferSize private const int BufferSize = 81920; @@ -147,7 +146,7 @@ public static void EnsureValid( } } - public static string GetPartitionKeyRangeId(Documents.PartitionKey partitionKey, PartitionKeyDefinition partitionKeyDefinition, Routing.CollectionRoutingMap collectionRoutingMap) + public static string GetPartitionKeyRangeId(PartitionKey partitionKey, PartitionKeyDefinition partitionKeyDefinition, Routing.CollectionRoutingMap collectionRoutingMap) { string effectivePartitionKey = partitionKey.InternalKey.GetEffectivePartitionKeyString(partitionKeyDefinition); return collectionRoutingMap.GetRangeByEffectivePartitionKey(effectivePartitionKey).Id; diff --git a/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs index 05402387fd..f576d939ff 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs @@ -39,6 +39,7 @@ public ItemBatchOperation( this.Id = id; this.ResourceStream = resourceStream; this.RequestOptions = requestOptions; + this.PartitionKeyJson = partitionKey.HasValue ? this.PartitionKey.Value.ToString() : null; } public ItemBatchOperation( diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs b/Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs index 18bc437479..7ea7a18e39 100644 --- a/Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs +++ b/Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs @@ -244,7 +244,9 @@ private bool AssertPartitioningPropertiesAndHeaders() if (partitionKeyRangeIdExists) { // Assert operation type is not write - if (this.OperationType != OperationType.Query && this.OperationType != OperationType.ReadFeed) + if (this.OperationType != OperationType.Query + && this.OperationType != OperationType.Batch + && this.OperationType != OperationType.ReadFeed) { throw new ArgumentOutOfRangeException(RMResources.UnexpectedPartitionKeyRangeId); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs new file mode 100644 index 0000000000..d3380c3e8c --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs @@ -0,0 +1,109 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests +{ + using System; + using System.Collections.Generic; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Newtonsoft.Json; + + [TestClass] +#pragma warning disable CA1001 // Types that own disposable fields should be disposable + public class BatchAsyncContainerExecutorTests +#pragma warning restore CA1001 // Types that own disposable fields should be disposable + { + private static CosmosSerializer cosmosDefaultJsonSerializer = new CosmosJsonDotNetSerializer(); + private CosmosClient cosmosClient; + private ContainerCore cosmosContainer; + + [TestInitialize] + public async Task InitializeAsync() + { + this.cosmosClient = TestCommon.CreateCosmosClient(useGateway: true); + DatabaseResponse db = await this.cosmosClient.CreateDatabaseAsync(Guid.NewGuid().ToString()); + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition(); + partitionKeyDefinition.Paths.Add("/Status"); + this.cosmosContainer = (ContainerCore)await db.Database.CreateContainerAsync(new ContainerProperties() { Id = Guid.NewGuid().ToString(), PartitionKey = partitionKeyDefinition }, 10000); + } + + [TestCleanup] + public async Task CleanupAsync() + { + await this.cosmosContainer.DeleteContainerAsync(); + } + + [TestMethod] + [Owner("maquaran")] + public async Task DoOperationsAsync() + { + BatchAsyncContainerExecutor executor = new BatchAsyncContainerExecutor(this.cosmosContainer, this.cosmosContainer.ClientContext, 20, Constants.MaxDirectModeBatchRequestBodySizeInBytes); + + List> tasks = new List>(); + for (int i = 0; i < 100; i++) + { + tasks.Add(executor.AddAsync(CreateItem(i.ToString()), default(CancellationToken))); + } + + await Task.WhenAll(tasks); + + for (int i = 0; i < 100; i++) + { + Task task = tasks[i]; + BatchOperationResult result = await task; + Assert.AreEqual(HttpStatusCode.Created, result.StatusCode); + + MyDocument document = cosmosDefaultJsonSerializer.FromStream(result.ResourceStream); + Assert.AreEqual(i.ToString(), document.id); + + ItemResponse storedDoc = await this.cosmosContainer.ReadItemAsync(i.ToString(), new Cosmos.PartitionKey(i.ToString())); + Assert.IsNotNull(storedDoc.Resource); + } + } + + [TestMethod] + [Owner("maquaran")] + public async Task ValidateInvalidRequestOptionsAsync() + { + BatchAsyncContainerExecutor executor = new BatchAsyncContainerExecutor(this.cosmosContainer, this.cosmosContainer.ClientContext, 20, Constants.MaxDirectModeBatchRequestBodySizeInBytes); + + string id = Guid.NewGuid().ToString(); + MyDocument myDocument = new MyDocument() { id = id, Status = id }; + + Assert.IsFalse(await executor.ValidateOperationAsync(new ItemBatchOperation(OperationType.Replace, 0, new Cosmos.PartitionKey(id), id, cosmosDefaultJsonSerializer.ToStream(myDocument)), new ItemRequestOptions() { SessionToken = "something" })); + } + + [TestMethod] + [Owner("maquaran")] + public async Task ValidateInvalidDocumentSizeAsync() + { + BatchAsyncContainerExecutor executor = new BatchAsyncContainerExecutor(this.cosmosContainer, this.cosmosContainer.ClientContext, 50, 2); + + string id = Guid.NewGuid().ToString(); + MyDocument myDocument = new MyDocument() { id = id, Status = id }; + + await Assert.ThrowsExceptionAsync(() => executor.ValidateOperationAsync(new ItemBatchOperation(OperationType.Replace, 0, new Cosmos.PartitionKey(id), id, cosmosDefaultJsonSerializer.ToStream(myDocument)))); + } + + private static ItemBatchOperation CreateItem(string id) + { + MyDocument myDocument = new MyDocument() { id = id, Status = id }; + return new ItemBatchOperation(OperationType.Create, 0, new Cosmos.PartitionKey(id), id, cosmosDefaultJsonSerializer.ToStream(myDocument)); + } + + private class MyDocument + { + //[JsonProperty("id")] + public string id { get; set; } + + public string Status { get; set; } + + public bool Updated { get; set; } + } + } +} \ No newline at end of file From d272a61296a5de15c0d23f4a6bbca7b1ac070e66 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Tue, 23 Jul 2019 15:35:44 -0700 Subject: [PATCH 04/41] Applying master changes --- .../src/Batch/BatchAsyncContainerExecutor.cs | 15 +++++++++++++++ .../src/Batch/ItemBatchOperation.cs | 1 - 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index ee9520ea24..aa23470286 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -272,12 +272,27 @@ private async Task ResolvePartitionKeyRangeIdAsync( } else { + await this.FillOperationPropertiesAsync(operation, cancellationToken); partitionKeyRangeId = BatchExecUtils.GetPartitionKeyRangeId(operation.PartitionKey.Value, partitionKeyDefinition, collectionRoutingMap); } return partitionKeyRangeId; } + private async Task FillOperationPropertiesAsync(ItemBatchOperation operation, CancellationToken cancellationToken) + { + // Same logic from RequestInvokerHandler to manage partition key migration + if (object.ReferenceEquals(operation.PartitionKey, PartitionKey.None)) + { + Documents.Routing.PartitionKeyInternal partitionKeyInternal = await this.cosmosContainer.GetNonePartitionKeyValueAsync(cancellationToken); + operation.PartitionKeyJson = partitionKeyInternal.ToJsonString(); + } + else + { + operation.PartitionKeyJson = operation.PartitionKey.Value.ToString(); + } + } + private async Task StartExecutionForPKRangeIdAsync( string pkRangeId, IEnumerable operations, diff --git a/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs index f576d939ff..05402387fd 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs @@ -39,7 +39,6 @@ public ItemBatchOperation( this.Id = id; this.ResourceStream = resourceStream; this.RequestOptions = requestOptions; - this.PartitionKeyJson = partitionKey.HasValue ? this.PartitionKey.Value.ToString() : null; } public ItemBatchOperation( From 03478eb364c0b11ec80f59977a3b5d83565847d5 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 29 Jul 2019 10:02:57 -0700 Subject: [PATCH 05/41] Comments --- .../src/Batch/BatchAsyncContainerExecutor.cs | 67 +++++++++---------- .../Batch/BatchAsyncContainerExecutorTests.cs | 2 +- 2 files changed, 33 insertions(+), 36 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index aa23470286..56cae927bf 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -102,25 +102,7 @@ public async Task ValidateOperationAsync( return false; } - if (itemRequestOptions.Properties != null - && (itemRequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKey, out object epkObj) - | itemRequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkStrObj))) - { - byte[] epk = epkObj as byte[]; - string epkStr = epkStrObj as string; - if (epk == null || epkStr == null) - { - throw new InvalidOperationException(string.Format( - ClientResources.EpkPropertiesPairingExpected, - WFConstants.BackendHeaders.EffectivePartitionKey, - WFConstants.BackendHeaders.EffectivePartitionKeyString)); - } - - if (operation.PartitionKey != null) - { - throw new InvalidOperationException(ClientResources.PKAndEpkSetTogether); - } - } + Debug.Assert(BatchAsyncContainerExecutor.ValidateOperationEPK(operation, itemRequestOptions)); } await operation.MaterializeResourceAsync(this.cosmosClientContext.CosmosSerializer, default(CancellationToken)); @@ -129,7 +111,7 @@ public async Task ValidateOperationAsync( if (itemByteSize > this.maxServerRequestBodyLength) { - throw new InvalidOperationException(RMResources.RequestTooLarge); + throw new ArgumentException(RMResources.RequestTooLarge); } return true; @@ -182,6 +164,33 @@ private async Task ExecuteAsync( return new CrossPartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); } + private static bool ValidateOperationEPK( + ItemBatchOperation operation, + ItemRequestOptions itemRequestOptions) + { + if (itemRequestOptions.Properties != null + && (itemRequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKey, out object epkObj) + | itemRequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkStrObj))) + { + byte[] epk = epkObj as byte[]; + string epkStr = epkStrObj as string; + if (epk == null || epkStr == null) + { + throw new InvalidOperationException(string.Format( + ClientResources.EpkPropertiesPairingExpected, + WFConstants.BackendHeaders.EffectivePartitionKey, + WFConstants.BackendHeaders.EffectivePartitionKeyString)); + } + + if (operation.PartitionKey != null) + { + throw new InvalidOperationException(ClientResources.PKAndEpkSetTogether); + } + } + + return true; + } + private static IEnumerable GetOperationsToRetry( IReadOnlyList operations, IEnumerable indexes) @@ -262,21 +271,9 @@ private async Task ResolvePartitionKeyRangeIdAsync( PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken); CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken); - string partitionKeyRangeId; - object epkObj = null; - if (operation.RequestOptions?.Properties?.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out epkObj) ?? false) - { - string epk = epkObj as string; - Debug.Assert(epk != null, "Expected non-null string for " + WFConstants.BackendHeaders.EffectivePartitionKeyString); - partitionKeyRangeId = collectionRoutingMap.GetRangeByEffectivePartitionKey(epk).Id; - } - else - { - await this.FillOperationPropertiesAsync(operation, cancellationToken); - partitionKeyRangeId = BatchExecUtils.GetPartitionKeyRangeId(operation.PartitionKey.Value, partitionKeyDefinition, collectionRoutingMap); - } - - return partitionKeyRangeId; + Debug.Assert(operation.RequestOptions?.Properties?.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkObj) == null, "EPK is not supported"); + await this.FillOperationPropertiesAsync(operation, cancellationToken); + return BatchExecUtils.GetPartitionKeyRangeId(operation.PartitionKey.Value, partitionKeyDefinition, collectionRoutingMap); } private async Task FillOperationPropertiesAsync(ItemBatchOperation operation, CancellationToken cancellationToken) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs index d3380c3e8c..0ca01ff865 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs @@ -87,7 +87,7 @@ public async Task ValidateInvalidDocumentSizeAsync() string id = Guid.NewGuid().ToString(); MyDocument myDocument = new MyDocument() { id = id, Status = id }; - await Assert.ThrowsExceptionAsync(() => executor.ValidateOperationAsync(new ItemBatchOperation(OperationType.Replace, 0, new Cosmos.PartitionKey(id), id, cosmosDefaultJsonSerializer.ToStream(myDocument)))); + await Assert.ThrowsExceptionAsync(() => executor.ValidateOperationAsync(new ItemBatchOperation(OperationType.Replace, 0, new Cosmos.PartitionKey(id), id, cosmosDefaultJsonSerializer.ToStream(myDocument)))); } private static ItemBatchOperation CreateItem(string id) From 3a84cbf70765383b707af21894dfc39d9a41570f Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 29 Jul 2019 12:46:00 -0700 Subject: [PATCH 06/41] Removing some blocking calls --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs | 4 ---- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs | 2 -- .../Batch/BatchAsyncBatcherTests.cs | 5 ++--- .../Batch/BatchAsyncStreamerTests.cs | 2 -- 4 files changed, 2 insertions(+), 11 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 0564c2a525..3a4271c712 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -74,10 +74,6 @@ public bool TryAdd(BatchAsyncOperationContext batchAsyncOperation) return false; } -#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits - batchAsyncOperation.Operation.MaterializeResourceAsync(this.CosmosSerializer, default(CancellationToken)).GetAwaiter().GetResult(); -#pragma warning restore VSTHRD002 // Avoid problematic synchronous waits - int itemByteSize = batchAsyncOperation.Operation.GetApproximateSerializedLength(); if (itemByteSize + this.currentSize > this.maxBatchByteSize) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 2e49794023..4142956a17 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -82,9 +82,7 @@ public BatchAsyncStreamer( public Task AddAsync(BatchAsyncOperationContext context) { BatchAsyncBatcher toDispatch = null; -#pragma warning disable VSTHRD103 // Call async methods when in an async method this.addLimiter.Wait(); -#pragma warning restore VSTHRD103 // Call async methods when in an async method try { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 0ab3462c12..3b3e4b7886 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -14,9 +14,7 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] -#pragma warning disable CA1001 // Types that own disposable fields should be disposable public class BatchAsyncBatcherTests -#pragma warning restore CA1001 // Types that own disposable fields should be disposable { private ItemBatchOperation ItemBatchOperation = new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); @@ -114,8 +112,9 @@ public void HasFixedSize() [TestMethod] [Owner("maquaran")] - public void HasFixedByteSize() + public async Task HasFixedByteSize() { + await ItemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), CancellationToken.None); // Each operation is 2 bytes BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor); Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 606ccb4d58..518f4d0fed 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -14,9 +14,7 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] -#pragma warning disable CA1001 // Types that own disposable fields should be disposable public class BatchAsyncStreamerTests -#pragma warning restore CA1001 // Types that own disposable fields should be disposable { private const int DispatchTimerInSeconds = 5; private const int MaxBatchByteSize = 100000; From 7dc78709ea560e5869e7169bc3bf6a1b1621acb3 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 29 Jul 2019 13:30:08 -0700 Subject: [PATCH 07/41] Removing more blocking calls --- .../src/Batch/BatchAsyncContainerExecutor.cs | 4 +++- .../src/Batch/BatchAsyncStreamer.cs | 9 +++------ .../Batch/BatchAsyncStreamerTests.cs | 15 +++++++++------ 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 56cae927bf..a6e06bcb1a 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -85,7 +85,9 @@ public async Task AddAsync( string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken); BatchAsyncStreamer streamer = this.GetStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId, operation); Debug.Assert(streamer != null, "Could not obtain streamer"); - return await streamer.AddAsync(new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation)); + BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); + await streamer.AddAsync(context); + return await context.Task; } public async Task ValidateOperationAsync( diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 4142956a17..e1474a1802 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -9,14 +9,13 @@ namespace Microsoft.Azure.Cosmos using System.Diagnostics; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Internal; using Microsoft.Azure.Documents; /// /// Handles operation queueing and dispatching. /// /// - /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. + /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. /// /// internal class BatchAsyncStreamer : IDisposable @@ -79,10 +78,10 @@ public BatchAsyncStreamer( this.StartTimer(); } - public Task AddAsync(BatchAsyncOperationContext context) + public async Task AddAsync(BatchAsyncOperationContext context) { BatchAsyncBatcher toDispatch = null; - this.addLimiter.Wait(); + await this.addLimiter.WaitAsync(); try { @@ -102,8 +101,6 @@ public Task AddAsync(BatchAsyncOperationContext context) { this.previousDispatchedTasks.Add(toDispatch.DispatchAsync()); } - - return context.Task; } public void Dispose() diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 518f4d0fed..c10e8aeb3d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -107,9 +107,9 @@ public void ValidatesSerializer() public async Task ExceptionsOnBatchBubbleUpAsync() { BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); - - Exception capturedException = await Assert.ThrowsExceptionAsync(() => batchAsyncStreamer.AddAsync(CreateContext(this.ItemBatchOperation))); - + var context = CreateContext(this.ItemBatchOperation); + await batchAsyncStreamer.AddAsync(context); + Exception capturedException = await Assert.ThrowsExceptionAsync(() => context.Task); Assert.AreEqual(expectedException, capturedException); } @@ -119,8 +119,9 @@ public async Task TimerDispatchesAsync() { // Bigger batch size than the amount of operations, timer should dispatch BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); - - BatchOperationResult result = await batchAsyncStreamer.AddAsync(CreateContext(this.ItemBatchOperation)); + var context = CreateContext(this.ItemBatchOperation); + await batchAsyncStreamer.AddAsync(context); + BatchOperationResult result = await context.Task; Assert.AreEqual(this.ItemBatchOperation.Id, result.ETag); } @@ -134,7 +135,9 @@ public async Task DispatchesAsync() List> contexts = new List>(10); for (int i = 0; i < 10; i++) { - contexts.Add(batchAsyncStreamer.AddAsync(CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())))); + var context = CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())); + await batchAsyncStreamer.AddAsync(context); + contexts.Add(context.Task); } await Task.WhenAll(contexts); From ef253817f33a53d44be7b7c41c5ff0ab1e289dbc Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 29 Jul 2019 14:06:31 -0700 Subject: [PATCH 08/41] Removing final blocking calls --- .../src/Batch/BatchAsyncStreamer.cs | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index e1474a1802..b45ab1e00b 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -27,6 +27,7 @@ internal class BatchAsyncStreamer : IDisposable private readonly Func, CancellationToken, Task> executor; private readonly int dispatchTimerInSeconds; private readonly CosmosSerializer CosmosSerializer; + private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); private BatchAsyncBatcher currentBatcher; private TimerPool timerPool; private PooledTimer currentTimer; @@ -99,7 +100,7 @@ public async Task AddAsync(BatchAsyncOperationContext context) if (toDispatch != null) { - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync()); + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); } } @@ -108,22 +109,8 @@ public void Dispose() this.disposed = true; this.addLimiter?.Dispose(); this.currentBatcher?.Dispose(); - - foreach (Task previousDispatch in this.previousDispatchedTasks) - { - try - { -#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits - previousDispatch.GetAwaiter().GetResult(); -#pragma warning restore VSTHRD002 // Avoid problematic synchronous waits - } - catch - { - // Any internal exceptions are disregarded during dispose - } - } - this.currentTimer.CancelTimer(); + this.cancellationTokenSource.Cancel(); } private void StartTimer() From 2c4332d616d0b20411401d6c6a84b47af7d88012 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 29 Jul 2019 14:12:14 -0700 Subject: [PATCH 09/41] Missing continuation token --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index b45ab1e00b..aac3d56300 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -143,7 +143,7 @@ private void DispatchTimer() if (toDispatch != null) { - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync()); + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); } this.StartTimer(); From dcbac643b57f690ead88f1a5342a3b5664295315 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 29 Jul 2019 15:05:26 -0700 Subject: [PATCH 10/41] Timer dispatch --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index aac3d56300..5789d62ce0 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -118,18 +118,18 @@ private void StartTimer() this.currentTimer = this.timerPool.GetPooledTimer(this.dispatchTimerInSeconds); this.timerTask = this.currentTimer.StartTimerAsync().ContinueWith((task) => { - this.DispatchTimer(); + return this.DispatchTimerAsync(); }); } - private void DispatchTimer() + private async Task DispatchTimerAsync() { if (this.disposed) { return; } - this.addLimiter.Wait(); + await this.addLimiter.WaitAsync(); BatchAsyncBatcher toDispatch; try From 121c4b374c27093db26ee92074d158006f651731 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 29 Jul 2019 16:34:27 -0700 Subject: [PATCH 11/41] Async Add --- .../src/Batch/BatchAsyncBatcher.cs | 6 ++- .../src/Batch/BatchAsyncContainerExecutor.cs | 2 +- .../src/Batch/BatchAsyncStreamer.cs | 10 +++-- .../Batch/BatchAsyncBatcherTests.cs | 37 +++++++++---------- .../Batch/BatchAsyncStreamerTests.cs | 6 +-- 5 files changed, 32 insertions(+), 29 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 3a4271c712..d9b209203b 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -59,14 +59,16 @@ public BatchAsyncBatcher( this.CosmosSerializer = cosmosSerializer; } - public bool TryAdd(BatchAsyncOperationContext batchAsyncOperation) + public async Task TryAddAsync( + BatchAsyncOperationContext batchAsyncOperation, + CancellationToken cancellationToken) { if (batchAsyncOperation == null) { throw new ArgumentNullException(nameof(batchAsyncOperation)); } - this.tryAddLimiter.Wait(); + await this.tryAddLimiter.WaitAsync(cancellationToken); try { if (this.batchOperations.Count == this.maxBatchOperationCount) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index a6e06bcb1a..8bee02efaf 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -86,7 +86,7 @@ public async Task AddAsync( BatchAsyncStreamer streamer = this.GetStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId, operation); Debug.Assert(streamer != null, "Could not obtain streamer"); BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); - await streamer.AddAsync(context); + await streamer.AddAsync(context, cancellationToken); return await context.Task; } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 5789d62ce0..02b76f84a6 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -15,7 +15,7 @@ namespace Microsoft.Azure.Cosmos /// Handles operation queueing and dispatching. /// /// - /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. + /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. /// /// internal class BatchAsyncStreamer : IDisposable @@ -79,18 +79,20 @@ public BatchAsyncStreamer( this.StartTimer(); } - public async Task AddAsync(BatchAsyncOperationContext context) + public async Task AddAsync( + BatchAsyncOperationContext context, + CancellationToken cancellationToken) { BatchAsyncBatcher toDispatch = null; await this.addLimiter.WaitAsync(); try { - if (!this.currentBatcher.TryAdd(context)) + if (!await this.currentBatcher.TryAddAsync(context, cancellationToken)) { // Batcher is full toDispatch = this.GetBatchToDispatch(); - Debug.Assert(this.currentBatcher.TryAdd(context), "Could not add context to batcher."); + Debug.Assert(await this.currentBatcher.TryAddAsync(context, cancellationToken), "Could not add context to batcher."); } } finally diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 3b3e4b7886..c3fe030346 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -16,11 +16,10 @@ namespace Microsoft.Azure.Cosmos.Tests [TestClass] public class BatchAsyncBatcherTests { - private ItemBatchOperation ItemBatchOperation = new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); - private static Exception expectedException = new Exception(); - // Executor just returns a reponse matching the Id with Etag + private ItemBatchOperation ItemBatchOperation = new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); + private Func, CancellationToken, Task> Executor = async (IReadOnlyList operations, CancellationToken cancellation) => { @@ -102,12 +101,12 @@ public void ValidatesSerializer() [TestMethod] [Owner("maquaran")] - public void HasFixedSize() + public async Task HasFixedSize() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); + Assert.IsFalse(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); } [TestMethod] @@ -117,9 +116,9 @@ public async Task HasFixedByteSize() await ItemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), CancellationToken.None); // Each operation is 2 bytes BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); + Assert.IsFalse(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); } [TestMethod] @@ -127,9 +126,9 @@ public async Task HasFixedByteSize() public void TryAddIsThreadSafe() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Task firstOperation = Task.Run(() => batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Task secondOperation = Task.Run(() => batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Task thirdOperation = Task.Run(() => batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Task firstOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None); + Task secondOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None); + Task thirdOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None); Task.WhenAll(firstOperation, secondOperation, thirdOperation).GetAwaiter().GetResult(); @@ -147,8 +146,8 @@ public async Task ExceptionsFailOperationsAsync() BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); - batchAsyncBatcher.TryAdd(context1); - batchAsyncBatcher.TryAdd(context2); + await batchAsyncBatcher.TryAddAsync(context1, CancellationToken.None); + await batchAsyncBatcher.TryAddAsync(context2, CancellationToken.None); await batchAsyncBatcher.DispatchAsync(); Assert.AreEqual(TaskStatus.Faulted, context1.Task.Status); @@ -167,7 +166,7 @@ public async Task DispatchProcessInOrderAsync() { BatchAsyncOperationContext context = new BatchAsyncOperationContext(string.Empty, new ItemBatchOperation(OperationType.Create, i, i.ToString())); contexts.Add(context); - Assert.IsTrue(batchAsyncBatcher.TryAdd(context)); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(context, CancellationToken.None)); } await batchAsyncBatcher.DispatchAsync(); @@ -187,7 +186,7 @@ public async Task DispatchProcessInOrderAsync() public async Task CannotDispatchTwiceAsync() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation)); + await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None); await batchAsyncBatcher.DispatchAsync(); await batchAsyncBatcher.DispatchAsync(); @@ -203,10 +202,10 @@ public void IsEmptyWithNoOperations() [TestMethod] [Owner("maquaran")] - public void IsNotEmptyWithOperations() + public async Task IsNotEmptyWithOperations() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); Assert.IsFalse(batchAsyncBatcher.IsEmpty); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index c10e8aeb3d..6c69449991 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -108,7 +108,7 @@ public async Task ExceptionsOnBatchBubbleUpAsync() { BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); var context = CreateContext(this.ItemBatchOperation); - await batchAsyncStreamer.AddAsync(context); + await batchAsyncStreamer.AddAsync(context, CancellationToken.None); Exception capturedException = await Assert.ThrowsExceptionAsync(() => context.Task); Assert.AreEqual(expectedException, capturedException); } @@ -120,7 +120,7 @@ public async Task TimerDispatchesAsync() // Bigger batch size than the amount of operations, timer should dispatch BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); var context = CreateContext(this.ItemBatchOperation); - await batchAsyncStreamer.AddAsync(context); + await batchAsyncStreamer.AddAsync(context, CancellationToken.None); BatchOperationResult result = await context.Task; Assert.AreEqual(this.ItemBatchOperation.Id, result.ETag); @@ -136,7 +136,7 @@ public async Task DispatchesAsync() for (int i = 0; i < 10; i++) { var context = CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())); - await batchAsyncStreamer.AddAsync(context); + await batchAsyncStreamer.AddAsync(context, CancellationToken.None); contexts.Add(context.Task); } From f966abb86c1157b49882e6afd20c412d966a5d4f Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Tue, 30 Jul 2019 10:26:37 -0700 Subject: [PATCH 12/41] ConfigureAwait --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs | 7 +++---- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs | 8 ++++---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index d9b209203b..43a6cded30 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -68,7 +68,7 @@ public async Task TryAddAsync( throw new ArgumentNullException(nameof(batchAsyncOperation)); } - await this.tryAddLimiter.WaitAsync(cancellationToken); + await this.tryAddLimiter.WaitAsync(cancellationToken).ConfigureAwait(false); try { if (this.batchOperations.Count == this.maxBatchOperationCount) @@ -98,11 +98,10 @@ public async Task TryAddAsync( public async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) { - await this.tryAddLimiter.WaitAsync(cancellationToken); - + await this.tryAddLimiter.WaitAsync(cancellationToken).ConfigureAwait(false); try { - CrossPartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken); + CrossPartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken).ConfigureAwait(false); // If the batch was not successful, we need to set all the responses if (!batchResponse.IsSuccessStatusCode) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 02b76f84a6..98bd24e16f 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -84,15 +84,15 @@ public async Task AddAsync( CancellationToken cancellationToken) { BatchAsyncBatcher toDispatch = null; - await this.addLimiter.WaitAsync(); + await this.addLimiter.WaitAsync(cancellationToken).ConfigureAwait(false); try { - if (!await this.currentBatcher.TryAddAsync(context, cancellationToken)) + if (!await this.currentBatcher.TryAddAsync(context, cancellationToken).ConfigureAwait(false)) { // Batcher is full toDispatch = this.GetBatchToDispatch(); - Debug.Assert(await this.currentBatcher.TryAddAsync(context, cancellationToken), "Could not add context to batcher."); + Debug.Assert(await this.currentBatcher.TryAddAsync(context, cancellationToken).ConfigureAwait(false), "Could not add context to batcher."); } } finally @@ -131,7 +131,7 @@ private async Task DispatchTimerAsync() return; } - await this.addLimiter.WaitAsync(); + await this.addLimiter.WaitAsync().ConfigureAwait(false); BatchAsyncBatcher toDispatch; try From f28cd51b48fb4bfe3e9d1a78006b4b64c477ec4c Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Tue, 30 Jul 2019 12:06:20 -0700 Subject: [PATCH 13/41] Refactoring limiters --- .../src/Batch/BatchAsyncBatcher.cs | 15 ++-- .../src/Batch/BatchAsyncContainerExecutor.cs | 42 +++++------ .../src/Batch/BatchAsyncStreamer.cs | 69 +++++++------------ .../Batch/BatchAsyncBatcherTests.cs | 38 ++++------ .../Batch/BatchAsyncStreamerTests.cs | 19 ++++- 5 files changed, 75 insertions(+), 108 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 43a6cded30..ba0868cf43 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -21,6 +21,7 @@ internal class BatchAsyncBatcher : IDisposable private readonly Func, CancellationToken, Task> executor; private readonly int maxBatchByteSize; private readonly int maxBatchOperationCount; + private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); private long currentSize = 0; public bool IsEmpty => this.batchOperations.Count == 0; @@ -59,16 +60,14 @@ public BatchAsyncBatcher( this.CosmosSerializer = cosmosSerializer; } - public async Task TryAddAsync( - BatchAsyncOperationContext batchAsyncOperation, - CancellationToken cancellationToken) + public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperation) { if (batchAsyncOperation == null) { throw new ArgumentNullException(nameof(batchAsyncOperation)); } - await this.tryAddLimiter.WaitAsync(cancellationToken).ConfigureAwait(false); + await this.tryAddLimiter.WaitAsync(this.cancellationTokenSource.Token).ConfigureAwait(false); try { if (this.batchOperations.Count == this.maxBatchOperationCount) @@ -98,7 +97,6 @@ public async Task TryAddAsync( public async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) { - await this.tryAddLimiter.WaitAsync(cancellationToken).ConfigureAwait(false); try { CrossPartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken).ConfigureAwait(false); @@ -130,16 +128,11 @@ public async Task TryAddAsync( operation.Fail(ex); } } - finally - { - this.tryAddLimiter.Release(); - this.Dispose(); - } } public void Dispose() { - this.tryAddLimiter?.Dispose(); + this.cancellationTokenSource.Cancel(); } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 8bee02efaf..1df882e5cc 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -13,7 +13,6 @@ namespace Microsoft.Azure.Cosmos using System.Net; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Internal; using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; @@ -83,10 +82,9 @@ public async Task AddAsync( } string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken); - BatchAsyncStreamer streamer = this.GetStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId, operation); - Debug.Assert(streamer != null, "Could not obtain streamer"); + BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); - await streamer.AddAsync(context, cancellationToken); + await streamer.AddAsync(context); return await context.Task; } @@ -126,11 +124,6 @@ public void Dispose() streamer.Value.Dispose(); } - foreach (KeyValuePair limiter in this.limitersByPartitionkeyRange) - { - limiter.Value.Dispose(); - } - this.timerPool.Dispose(); } @@ -142,8 +135,8 @@ private async Task ExecuteAsync( List completedResults = new List(); // Initial operations are all for the same PK Range Id - string initialPartitionKeyRangeId = operations[0].PartitionKeyRangeId; - inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(initialPartitionKeyRangeId, operations, cancellationToken)); + string pkRangeId = operations[0].PartitionKeyRangeId; + inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(pkRangeId, operations, cancellationToken)); while (inProgressTasksByPKRangeId.Count > 0) { @@ -230,13 +223,13 @@ private static IEnumerable GetOverflowOperations( private async Task GroupOperationsAndStartExecutionAsync( List> taskContainer, IEnumerable operations, - CancellationToken cancellation) + CancellationToken cancellationToken) { - Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, cancellation); + Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, cancellationToken); foreach (KeyValuePair> operationsForPKRangeId in operationsByPKRangeId) { - Task toAdd = this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellation); + Task toAdd = this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellationToken); taskContainer.Add(toAdd); } } @@ -245,8 +238,8 @@ private async Task>> GroupBy IEnumerable operations, CancellationToken cancellationToken) { - PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken); - CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken); + PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken).ConfigureAwait(false); + CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken).ConfigureAwait(false); Dictionary> operationsByPKRangeId = new Dictionary>(); foreach (BatchAsyncOperationContext operation in operations) @@ -270,6 +263,7 @@ private async Task ResolvePartitionKeyRangeIdAsync( ItemBatchOperation operation, CancellationToken cancellationToken) { + cancellationToken.ThrowIfCancellationRequested(); PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken); CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken); @@ -283,7 +277,7 @@ private async Task FillOperationPropertiesAsync(ItemBatchOperation operation, Ca // Same logic from RequestInvokerHandler to manage partition key migration if (object.ReferenceEquals(operation.PartitionKey, PartitionKey.None)) { - Documents.Routing.PartitionKeyInternal partitionKeyInternal = await this.cosmosContainer.GetNonePartitionKeyValueAsync(cancellationToken); + Documents.Routing.PartitionKeyInternal partitionKeyInternal = await this.cosmosContainer.GetNonePartitionKeyValueAsync(cancellationToken).ConfigureAwait(false); operation.PartitionKeyJson = partitionKeyInternal.ToJsonString(); } else @@ -298,7 +292,7 @@ private async Task StartExecutionForPKRan CancellationToken cancellationToken) { SemaphoreSlim limiter = this.GetLimiterForPartitionKeyRange(pkRangeId); - await limiter.WaitAsync(cancellationToken); + await limiter.WaitAsync(cancellationToken).ConfigureAwait(false); try { @@ -324,7 +318,7 @@ private async Task CreateServerRequestAsync this.maxServerRequestOperationCount, ensureContinuousOperationIndexes: false, serializer: this.cosmosClientContext.CosmosSerializer, - cancellationToken: cancellationToken); + cancellationToken: cancellationToken).ConfigureAwait(false); return request; } @@ -358,7 +352,7 @@ private async Task ExecuteServerRequestsA || serverResponse.SubStatusCode == SubStatusCodes.PartitionKeyRangeGone)) { // lower layers would have refreshed the appropriate routing caches, but we need to retry the operations on the new PKRanges. - return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, serverResponses, new List(serverRequest.Operations)); + return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, serverResponses, new List(operations)); } else { @@ -392,15 +386,13 @@ private async Task ExecuteServerRequestAsync(CrossPartitionKeySer partitionKey: null, streamPayload: serverRequestPayload, requestEnricher: requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest.PartitionKeyRangeId), - cancellationToken: cancellationToken); + cancellationToken: cancellationToken).ConfigureAwait(false); - return await BatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.CosmosSerializer); + return await BatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.CosmosSerializer).ConfigureAwait(false); } } - private BatchAsyncStreamer GetStreamerForPartitionKeyRange( - string partitionKeyRangeId, - ItemBatchOperation operation) + private BatchAsyncStreamer GetOrAddStreamerForPartitionKeyRange(string partitionKeyRangeId) { if (this.streamersByPartitionKeyRange.TryGetValue(partitionKeyRangeId, out BatchAsyncStreamer streamer)) { diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 98bd24e16f..b6aee9d753 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -15,13 +15,13 @@ namespace Microsoft.Azure.Cosmos /// Handles operation queueing and dispatching. /// /// - /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. + /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. /// /// internal class BatchAsyncStreamer : IDisposable { private readonly List previousDispatchedTasks = new List(); - private readonly SemaphoreSlim addLimiter; + private readonly SemaphoreSlim dispatchLimiter; private readonly int maxBatchOperationCount; private readonly int maxBatchByteSize; private readonly Func, CancellationToken, Task> executor; @@ -73,46 +73,30 @@ public BatchAsyncStreamer( this.dispatchTimerInSeconds = dispatchTimerInSeconds; this.timerPool = timerPool; this.CosmosSerializer = cosmosSerializer; - this.addLimiter = new SemaphoreSlim(1, 1); - this.currentBatcher = this.GetBatchAsyncBatcher(); + this.dispatchLimiter = new SemaphoreSlim(1, 1); + this.currentBatcher = this.CreateBatchAsyncBatcher(); this.StartTimer(); } - public async Task AddAsync( - BatchAsyncOperationContext context, - CancellationToken cancellationToken) + public async Task AddAsync(BatchAsyncOperationContext context) { - BatchAsyncBatcher toDispatch = null; - await this.addLimiter.WaitAsync(cancellationToken).ConfigureAwait(false); - - try - { - if (!await this.currentBatcher.TryAddAsync(context, cancellationToken).ConfigureAwait(false)) - { - // Batcher is full - toDispatch = this.GetBatchToDispatch(); - Debug.Assert(await this.currentBatcher.TryAddAsync(context, cancellationToken).ConfigureAwait(false), "Could not add context to batcher."); - } - } - finally - { - this.addLimiter.Release(); - } - - if (toDispatch != null) + if (!await this.currentBatcher.TryAddAsync(context).ConfigureAwait(false)) { + // Batcher is full + BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAsync(); this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + bool addedContext = await this.currentBatcher.TryAddAsync(context).ConfigureAwait(false); + Debug.Assert(addedContext, "Could not add context to batcher."); } } public void Dispose() { this.disposed = true; - this.addLimiter?.Dispose(); + this.cancellationTokenSource.Cancel(); this.currentBatcher?.Dispose(); this.currentTimer.CancelTimer(); - this.cancellationTokenSource.Cancel(); } private void StartTimer() @@ -131,18 +115,7 @@ private async Task DispatchTimerAsync() return; } - await this.addLimiter.WaitAsync().ConfigureAwait(false); - - BatchAsyncBatcher toDispatch; - try - { - toDispatch = this.GetBatchToDispatch(); - } - finally - { - this.addLimiter.Release(); - } - + BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAsync(); if (toDispatch != null) { this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); @@ -151,19 +124,27 @@ private async Task DispatchTimerAsync() this.StartTimer(); } - private BatchAsyncBatcher GetBatchToDispatch() + private async Task GetBatchToDispatchAsync() { if (this.currentBatcher.IsEmpty) { return null; } - BatchAsyncBatcher previousBatcher = this.currentBatcher; - this.currentBatcher = this.GetBatchAsyncBatcher(); - return previousBatcher; + await this.dispatchLimiter.WaitAsync(this.cancellationTokenSource.Token).ConfigureAwait(false); + try + { + BatchAsyncBatcher previousBatcher = this.currentBatcher; + this.currentBatcher = this.CreateBatchAsyncBatcher(); + return previousBatcher; + } + finally + { + this.dispatchLimiter.Release(); + } } - private BatchAsyncBatcher GetBatchAsyncBatcher() + private BatchAsyncBatcher CreateBatchAsyncBatcher() { return new BatchAsyncBatcher(this.maxBatchOperationCount, this.maxBatchByteSize, this.CosmosSerializer, this.executor); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index c3fe030346..4ee2436594 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -104,9 +104,9 @@ public void ValidatesSerializer() public async Task HasFixedSize() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); - Assert.IsFalse(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsFalse(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); } [TestMethod] @@ -116,9 +116,9 @@ public async Task HasFixedByteSize() await ItemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), CancellationToken.None); // Each operation is 2 bytes BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); - Assert.IsFalse(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsFalse(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); } [TestMethod] @@ -126,9 +126,9 @@ public async Task HasFixedByteSize() public void TryAddIsThreadSafe() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Task firstOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None); - Task secondOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None); - Task thirdOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None); + Task firstOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation)); + Task secondOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation)); + Task thirdOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation)); Task.WhenAll(firstOperation, secondOperation, thirdOperation).GetAwaiter().GetResult(); @@ -146,8 +146,8 @@ public async Task ExceptionsFailOperationsAsync() BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); - await batchAsyncBatcher.TryAddAsync(context1, CancellationToken.None); - await batchAsyncBatcher.TryAddAsync(context2, CancellationToken.None); + await batchAsyncBatcher.TryAddAsync(context1); + await batchAsyncBatcher.TryAddAsync(context2); await batchAsyncBatcher.DispatchAsync(); Assert.AreEqual(TaskStatus.Faulted, context1.Task.Status); @@ -166,7 +166,7 @@ public async Task DispatchProcessInOrderAsync() { BatchAsyncOperationContext context = new BatchAsyncOperationContext(string.Empty, new ItemBatchOperation(OperationType.Create, i, i.ToString())); contexts.Add(context); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(context, CancellationToken.None)); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(context)); } await batchAsyncBatcher.DispatchAsync(); @@ -180,18 +180,6 @@ public async Task DispatchProcessInOrderAsync() } } - [TestMethod] - [Owner("maquaran")] - [ExpectedException(typeof(ObjectDisposedException))] - public async Task CannotDispatchTwiceAsync() - { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None); - - await batchAsyncBatcher.DispatchAsync(); - await batchAsyncBatcher.DispatchAsync(); - } - [TestMethod] [Owner("maquaran")] public void IsEmptyWithNoOperations() @@ -205,7 +193,7 @@ public void IsEmptyWithNoOperations() public async Task IsNotEmptyWithOperations() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation), CancellationToken.None)); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); Assert.IsFalse(batchAsyncBatcher.IsEmpty); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 6c69449991..7c25ad38e5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -108,7 +108,7 @@ public async Task ExceptionsOnBatchBubbleUpAsync() { BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); var context = CreateContext(this.ItemBatchOperation); - await batchAsyncStreamer.AddAsync(context, CancellationToken.None); + await batchAsyncStreamer.AddAsync(context); Exception capturedException = await Assert.ThrowsExceptionAsync(() => context.Task); Assert.AreEqual(expectedException, capturedException); } @@ -120,7 +120,7 @@ public async Task TimerDispatchesAsync() // Bigger batch size than the amount of operations, timer should dispatch BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); var context = CreateContext(this.ItemBatchOperation); - await batchAsyncStreamer.AddAsync(context, CancellationToken.None); + await batchAsyncStreamer.AddAsync(context); BatchOperationResult result = await context.Task; Assert.AreEqual(this.ItemBatchOperation.Id, result.ETag); @@ -136,7 +136,7 @@ public async Task DispatchesAsync() for (int i = 0; i < 10; i++) { var context = CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())); - await batchAsyncStreamer.AddAsync(context, CancellationToken.None); + await batchAsyncStreamer.AddAsync(context); contexts.Add(context.Task); } @@ -151,6 +151,19 @@ public async Task DispatchesAsync() } } + [TestMethod] + [Owner("maquaran")] + [ExpectedException(typeof(TaskCanceledException))] + public async Task DisposeAndAdd() + { + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + List operations = new List(); + var context = CreateContext(new ItemBatchOperation(OperationType.Create, 0, "0")); + operations.Add(Task.Run(() => batchAsyncStreamer.Dispose())); + operations.Add(batchAsyncStreamer.AddAsync(context)); + await Task.WhenAll(operations); + } + private static BatchAsyncOperationContext CreateContext(ItemBatchOperation operation) => new BatchAsyncOperationContext(string.Empty, operation); } } From 5ed005c80a5e5848cc40b9c4c2ca3b7307fe6ce3 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Tue, 30 Jul 2019 12:29:26 -0700 Subject: [PATCH 14/41] Fixing test --- .../Batch/BatchAsyncStreamerTests.cs | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 7c25ad38e5..c10e8aeb3d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -151,19 +151,6 @@ public async Task DispatchesAsync() } } - [TestMethod] - [Owner("maquaran")] - [ExpectedException(typeof(TaskCanceledException))] - public async Task DisposeAndAdd() - { - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); - List operations = new List(); - var context = CreateContext(new ItemBatchOperation(OperationType.Create, 0, "0")); - operations.Add(Task.Run(() => batchAsyncStreamer.Dispose())); - operations.Add(batchAsyncStreamer.AddAsync(context)); - await Task.WhenAll(operations); - } - private static BatchAsyncOperationContext CreateContext(ItemBatchOperation operation) => new BatchAsyncOperationContext(string.Empty, operation); } } From e696741f48f25dae2d0cd4441fbb99b21e3fb866 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 1 Aug 2019 09:33:40 -0700 Subject: [PATCH 15/41] Addressing comments --- .../src/Batch/BatchAsyncBatcher.cs | 8 ++-- .../src/Batch/BatchAsyncContainerExecutor.cs | 48 ++++++++----------- .../src/Batch/BatchAsyncStreamer.cs | 16 +++---- ...sponse.cs => PartitionKeyBatchResponse.cs} | 10 ++-- ...t.cs => PartitionKeyServerBatchRequest.cs} | 14 +++--- .../src/ClientResources.Designer.cs | 11 ++++- .../src/ClientResources.resx | 3 ++ .../Batch/BatchAsyncContainerExecutorTests.cs | 4 +- .../Batch/BatchAsyncBatcherTests.cs | 6 +-- .../Batch/BatchAsyncStreamerTests.cs | 6 +-- .../CrossPartitionKeyBatchResponseTests.cs | 4 +- 11 files changed, 67 insertions(+), 63 deletions(-) rename Microsoft.Azure.Cosmos/src/Batch/{CrossPartitionKeyBatchResponse.cs => PartitionKeyBatchResponse.cs} (94%) rename Microsoft.Azure.Cosmos/src/Batch/{CrossPartitionKeyServerBatchRequest.cs => PartitionKeyServerBatchRequest.cs} (80%) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index ba0868cf43..116b03363f 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -18,7 +18,7 @@ internal class BatchAsyncBatcher : IDisposable private readonly SemaphoreSlim tryAddLimiter; private readonly CosmosSerializer CosmosSerializer; private readonly List batchOperations; - private readonly Func, CancellationToken, Task> executor; + private readonly Func, CancellationToken, Task> executor; private readonly int maxBatchByteSize; private readonly int maxBatchOperationCount; private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); @@ -30,7 +30,7 @@ public BatchAsyncBatcher( int maxBatchOperationCount, int maxBatchByteSize, CosmosSerializer cosmosSerializer, - Func, CancellationToken, Task> executor) + Func, CancellationToken, Task> executor) { if (maxBatchOperationCount < 1) { @@ -67,7 +67,7 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati throw new ArgumentNullException(nameof(batchAsyncOperation)); } - await this.tryAddLimiter.WaitAsync(this.cancellationTokenSource.Token).ConfigureAwait(false); + await this.tryAddLimiter.WaitAsync(this.cancellationTokenSource.Token); try { if (this.batchOperations.Count == this.maxBatchOperationCount) @@ -99,7 +99,7 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati { try { - CrossPartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken).ConfigureAwait(false); + PartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken); // If the batch was not successful, we need to set all the responses if (!batchResponse.IsSuccessStatusCode) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 1df882e5cc..efd0d8d111 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -74,23 +74,27 @@ public BatchAsyncContainerExecutor( public async Task AddAsync( ItemBatchOperation operation, - CancellationToken cancellationToken) + ItemRequestOptions itemRequestOptions = null, + CancellationToken cancellationToken = default(CancellationToken)) { if (operation == null) { throw new ArgumentNullException(nameof(operation)); } - string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken); + await this.ValidateOperationAsync(operation, itemRequestOptions, cancellationToken); + + string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false); BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); - await streamer.AddAsync(context); + await streamer.AddAsync(context).ConfigureAwait(false); return await context.Task; } - public async Task ValidateOperationAsync( + internal async Task ValidateOperationAsync( ItemBatchOperation operation, - ItemRequestOptions itemRequestOptions = null) + ItemRequestOptions itemRequestOptions = null, + CancellationToken cancellationToken = default(CancellationToken)) { if (itemRequestOptions != null) { @@ -99,13 +103,13 @@ public async Task ValidateOperationAsync( || itemRequestOptions.PostTriggers != null || itemRequestOptions.SessionToken != null) { - return false; + throw new InvalidOperationException(ClientResources.UnsupportedBatchRequestOptions); } Debug.Assert(BatchAsyncContainerExecutor.ValidateOperationEPK(operation, itemRequestOptions)); } - await operation.MaterializeResourceAsync(this.cosmosClientContext.CosmosSerializer, default(CancellationToken)); + await operation.MaterializeResourceAsync(this.cosmosClientContext.CosmosSerializer, cancellationToken).ConfigureAwait(false); int itemByteSize = operation.GetApproximateSerializedLength(); @@ -113,8 +117,6 @@ public async Task ValidateOperationAsync( { throw new ArgumentException(RMResources.RequestTooLarge); } - - return true; } public void Dispose() @@ -127,7 +129,7 @@ public void Dispose() this.timerPool.Dispose(); } - private async Task ExecuteAsync( + private async Task ExecuteAsync( IReadOnlyList operations, CancellationToken cancellationToken) { @@ -149,14 +151,13 @@ private async Task ExecuteAsync( { IEnumerable retryOperations = BatchAsyncContainerExecutor.GetOperationsToRetry(operations, executionResult.PendingOperations.Select(o => o.OperationIndex)); await this.GroupOperationsAndStartExecutionAsync(inProgressTasksByPKRangeId, retryOperations, cancellationToken); - executionResult.PendingOperations = null; } inProgressTasksByPKRangeId.Remove(completedTask); } IEnumerable serverResponses = completedResults.SelectMany(res => res.ServerResponses); - return new CrossPartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); + return new PartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); } private static bool ValidateOperationEPK( @@ -207,7 +208,7 @@ private static void AddHeadersToRequestMessage(RequestMessage requestMessage, st /// If because of HybridRow serialization overhead, not all operations fit in the request, we send those extra operations in a separate request. /// private static IEnumerable GetOverflowOperations( - CrossPartitionKeyServerBatchRequest request, + PartitionKeyServerBatchRequest request, IEnumerable operationsSentToRequest) { int totalOperations = operationsSentToRequest.Count(); @@ -244,7 +245,7 @@ private async Task>> GroupBy Dictionary> operationsByPKRangeId = new Dictionary>(); foreach (BatchAsyncOperationContext operation in operations) { - string partitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation.Operation, cancellationToken); + string partitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation.Operation, cancellationToken).ConfigureAwait(false); if (operationsByPKRangeId.TryGetValue(partitionKeyRangeId, out List operationsForPKRangeId)) { @@ -304,14 +305,14 @@ private async Task StartExecutionForPKRan } } - private async Task CreateServerRequestAsync( + private async Task CreateServerRequestAsync( string partitionKeyRangeId, IEnumerable operations, CancellationToken cancellationToken) { ArraySegment operationsArraySegment = new ArraySegment(operations.ToArray()); - CrossPartitionKeyServerBatchRequest request = await CrossPartitionKeyServerBatchRequest.CreateAsync( + PartitionKeyServerBatchRequest request = await PartitionKeyServerBatchRequest.CreateAsync( partitionKeyRangeId, operationsArraySegment, this.maxServerRequestBodyLength, @@ -330,22 +331,13 @@ private async Task ExecuteServerRequestsA { List serverResponses = new List(); List overflowOperations = new List(); - CrossPartitionKeyServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operations, cancellationToken); + PartitionKeyServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operations, cancellationToken); // In case some operations overflowed overflowOperations.AddRange(BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operations)); do { - BatchResponse serverResponse; - try - { - serverResponse = await this.ExecuteServerRequestAsync(serverRequest, cancellationToken); - } - catch (CosmosException ex) - { - serverResponse = new BatchResponse(ex.StatusCode, (SubStatusCodes)ex.SubStatusCode, ex.Message, serverRequest.Operations); - } - + BatchResponse serverResponse = await this.ExecuteServerRequestAsync(serverRequest, cancellationToken); if (serverResponse.StatusCode == HttpStatusCode.Gone && (serverResponse.SubStatusCode == SubStatusCodes.CompletingSplit || serverResponse.SubStatusCode == SubStatusCodes.CompletingPartitionMigration @@ -371,7 +363,7 @@ private async Task ExecuteServerRequestsA return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, serverResponses); } - private async Task ExecuteServerRequestAsync(CrossPartitionKeyServerBatchRequest serverRequest, CancellationToken cancellationToken) + private async Task ExecuteServerRequestAsync(PartitionKeyServerBatchRequest serverRequest, CancellationToken cancellationToken) { using (Stream serverRequestPayload = serverRequest.TransferBodyStream()) { diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index b6aee9d753..01ebc65a36 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -24,7 +24,7 @@ internal class BatchAsyncStreamer : IDisposable private readonly SemaphoreSlim dispatchLimiter; private readonly int maxBatchOperationCount; private readonly int maxBatchByteSize; - private readonly Func, CancellationToken, Task> executor; + private readonly Func, CancellationToken, Task> executor; private readonly int dispatchTimerInSeconds; private readonly CosmosSerializer CosmosSerializer; private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); @@ -40,7 +40,7 @@ public BatchAsyncStreamer( int dispatchTimerInSeconds, TimerPool timerPool, CosmosSerializer cosmosSerializer, - Func, CancellationToken, Task> executor) + Func, CancellationToken, Task> executor) { if (maxBatchOperationCount < 1) { @@ -81,12 +81,12 @@ public BatchAsyncStreamer( public async Task AddAsync(BatchAsyncOperationContext context) { - if (!await this.currentBatcher.TryAddAsync(context).ConfigureAwait(false)) + if (!await this.currentBatcher.TryAddAsync(context)) { // Batcher is full - BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAsync(); + BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAndCreateAsync(); this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); - bool addedContext = await this.currentBatcher.TryAddAsync(context).ConfigureAwait(false); + bool addedContext = await this.currentBatcher.TryAddAsync(context); Debug.Assert(addedContext, "Could not add context to batcher."); } } @@ -115,7 +115,7 @@ private async Task DispatchTimerAsync() return; } - BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAsync(); + BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAndCreateAsync(); if (toDispatch != null) { this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); @@ -124,14 +124,14 @@ private async Task DispatchTimerAsync() this.StartTimer(); } - private async Task GetBatchToDispatchAsync() + private async Task GetBatchToDispatchAndCreateAsync() { if (this.currentBatcher.IsEmpty) { return null; } - await this.dispatchLimiter.WaitAsync(this.cancellationTokenSource.Token).ConfigureAwait(false); + await this.dispatchLimiter.WaitAsync(this.cancellationTokenSource.Token); try { BatchAsyncBatcher previousBatcher = this.currentBatcher; diff --git a/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyBatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs similarity index 94% rename from Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyBatchResponse.cs rename to Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs index 973fc84d0e..e7b3655a75 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyBatchResponse.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs @@ -20,7 +20,7 @@ namespace Microsoft.Azure.Cosmos #else internal #endif - class CrossPartitionKeyBatchResponse : BatchResponse + class PartitionKeyBatchResponse : BatchResponse #pragma warning restore CA1710 // Identifiers should have correct suffix { // Results sorted in the order operations had been added. @@ -28,14 +28,14 @@ class CrossPartitionKeyBatchResponse : BatchResponse private bool isDisposed; /// - /// Initializes a new instance of the class. + /// Initializes a new instance of the class. /// /// Completion status code of the batch request. /// Provides further details about why the batch was not processed. /// Operations that were supposed to be executed, but weren't. /// The reason for failure if any. // This constructor is expected to be used when the batch is not executed at all (if it is a bad request). - internal CrossPartitionKeyBatchResponse( + internal PartitionKeyBatchResponse( HttpStatusCode statusCode, SubStatusCodes subStatusCode, string errorMessage, @@ -45,11 +45,11 @@ internal CrossPartitionKeyBatchResponse( } /// - /// Initializes a new instance of the class. + /// Initializes a new instance of the class. /// /// Responses from the server. /// Serializer to deserialize response resource body streams. - internal CrossPartitionKeyBatchResponse(IEnumerable serverResponses, CosmosSerializer serializer) + internal PartitionKeyBatchResponse(IEnumerable serverResponses, CosmosSerializer serializer) { this.StatusCode = serverResponses.Any(r => r.StatusCode != HttpStatusCode.OK) ? (HttpStatusCode)BatchExecUtils.StatusCodeMultiStatus diff --git a/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyServerBatchRequest.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyServerBatchRequest.cs similarity index 80% rename from Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyServerBatchRequest.cs rename to Microsoft.Azure.Cosmos/src/Batch/PartitionKeyServerBatchRequest.cs index e337df17cd..98e73c5a0b 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/CrossPartitionKeyServerBatchRequest.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyServerBatchRequest.cs @@ -8,16 +8,16 @@ namespace Microsoft.Azure.Cosmos using System.Threading; using System.Threading.Tasks; - internal sealed class CrossPartitionKeyServerBatchRequest : ServerBatchRequest + internal sealed class PartitionKeyServerBatchRequest : ServerBatchRequest { /// - /// Initializes a new instance of the class. + /// Initializes a new instance of the class. /// /// The partition key range id associated with all requests. /// Maximum length allowed for the request body. /// Maximum number of operations allowed in the request. /// Serializer to serialize user provided objects to JSON. - public CrossPartitionKeyServerBatchRequest( + public PartitionKeyServerBatchRequest( string partitionKeyRangeId, int maxBodyLength, int maxOperationCount, @@ -33,7 +33,7 @@ public CrossPartitionKeyServerBatchRequest( public string PartitionKeyRangeId { get; } /// - /// Creates an instance of . + /// Creates an instance of . /// In case of direct mode requests, all the operations are expected to belong to the same PartitionKeyRange. /// The body of the request is populated with operations till it reaches the provided maxBodyLength. /// @@ -44,8 +44,8 @@ public CrossPartitionKeyServerBatchRequest( /// Whether to stop adding operations to the request once there is non-continuity in the operation indexes. /// Serializer to serialize user provided objects to JSON. /// representing request cancellation. - /// A newly created instance of . - public static async Task CreateAsync( + /// A newly created instance of . + public static async Task CreateAsync( string partitionKeyRangeId, ArraySegment operations, int maxBodyLength, @@ -54,7 +54,7 @@ public static async Task CreateAsync( CosmosSerializer serializer, CancellationToken cancellationToken) { - CrossPartitionKeyServerBatchRequest request = new CrossPartitionKeyServerBatchRequest(partitionKeyRangeId, maxBodyLength, maxOperationCount, serializer); + PartitionKeyServerBatchRequest request = new PartitionKeyServerBatchRequest(partitionKeyRangeId, maxBodyLength, maxOperationCount, serializer); await request.CreateBodyStreamAsync(operations, cancellationToken, ensureContinuousOperationIndexes); return request; } diff --git a/Microsoft.Azure.Cosmos/src/ClientResources.Designer.cs b/Microsoft.Azure.Cosmos/src/ClientResources.Designer.cs index 3fc843adb5..57dc4df2a3 100644 --- a/Microsoft.Azure.Cosmos/src/ClientResources.Designer.cs +++ b/Microsoft.Azure.Cosmos/src/ClientResources.Designer.cs @@ -19,7 +19,7 @@ namespace Microsoft.Azure.Cosmos { // class via a tool like ResGen or Visual Studio. // To add or remove a member, edit your .ResX file then rerun ResGen // with the /str option, or rebuild your VS project. - [global::System.CodeDom.Compiler.GeneratedCodeAttribute("System.Resources.Tools.StronglyTypedResourceBuilder", "16.0.0.0")] + [global::System.CodeDom.Compiler.GeneratedCodeAttribute("System.Resources.Tools.StronglyTypedResourceBuilder", "15.0.0.0")] [global::System.Diagnostics.DebuggerNonUserCodeAttribute()] [global::System.Runtime.CompilerServices.CompilerGeneratedAttribute()] internal class ClientResources { @@ -564,6 +564,15 @@ internal static string UnexpectedTokenType { } } + /// + /// Looks up a localized string similar to Consistency, Session, and Triggers are not allowed when using the Batch streaming feature.. + /// + internal static string UnsupportedBatchRequestOptions { + get { + return ResourceManager.GetString("UnsupportedBatchRequestOptions", resourceCulture); + } + } + /// /// Looks up a localized string similar to Unsupported type {0} for partitionKey.. /// diff --git a/Microsoft.Azure.Cosmos/src/ClientResources.resx b/Microsoft.Azure.Cosmos/src/ClientResources.resx index ceb48e974a..d781f3822d 100644 --- a/Microsoft.Azure.Cosmos/src/ClientResources.resx +++ b/Microsoft.Azure.Cosmos/src/ClientResources.resx @@ -291,4 +291,7 @@ Instantiation of only value types, anonymous types and spatial types are supported. + + Consistency, Session, and Triggers are not allowed when using the Batch streaming feature. + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs index 0ca01ff865..720ab61481 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs @@ -47,7 +47,7 @@ public async Task DoOperationsAsync() List> tasks = new List>(); for (int i = 0; i < 100; i++) { - tasks.Add(executor.AddAsync(CreateItem(i.ToString()), default(CancellationToken))); + tasks.Add(executor.AddAsync(CreateItem(i.ToString()), null, default(CancellationToken))); } await Task.WhenAll(tasks); @@ -75,7 +75,7 @@ public async Task ValidateInvalidRequestOptionsAsync() string id = Guid.NewGuid().ToString(); MyDocument myDocument = new MyDocument() { id = id, Status = id }; - Assert.IsFalse(await executor.ValidateOperationAsync(new ItemBatchOperation(OperationType.Replace, 0, new Cosmos.PartitionKey(id), id, cosmosDefaultJsonSerializer.ToStream(myDocument)), new ItemRequestOptions() { SessionToken = "something" })); + await Assert.ThrowsExceptionAsync(() => executor.ValidateOperationAsync(new ItemBatchOperation(OperationType.Replace, 0, new Cosmos.PartitionKey(id), id, cosmosDefaultJsonSerializer.ToStream(myDocument)), new ItemRequestOptions() { SessionToken = "something" })); } [TestMethod] diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 4ee2436594..549197fc0a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -20,7 +20,7 @@ public class BatchAsyncBatcherTests private ItemBatchOperation ItemBatchOperation = new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); - private Func, CancellationToken, Task> Executor + private Func, CancellationToken, Task> Executor = async (IReadOnlyList operations, CancellationToken cancellation) => { List results = new List(); @@ -53,11 +53,11 @@ private Func, CancellationToken, Task< batchRequest, new CosmosJsonDotNetSerializer()); - CrossPartitionKeyBatchResponse response = new CrossPartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); return response; }; - private Func, CancellationToken, Task> ExecutorWithFailure + private Func, CancellationToken, Task> ExecutorWithFailure = (IReadOnlyList operations, CancellationToken cancellation) => { throw expectedException; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index c10e8aeb3d..e2ebea66d0 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -23,7 +23,7 @@ public class BatchAsyncStreamerTests private TimerPool TimerPool = new TimerPool(1); // Executor just returns a reponse matching the Id with Etag - private Func, CancellationToken, Task> Executor + private Func, CancellationToken, Task> Executor = async (IReadOnlyList operations, CancellationToken cancellation) => { List results = new List(); @@ -56,11 +56,11 @@ private Func, CancellationToken, Task< batchRequest, new CosmosJsonDotNetSerializer()); - CrossPartitionKeyBatchResponse response = new CrossPartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); return response; }; - private Func, CancellationToken, Task> ExecutorWithFailure + private Func, CancellationToken, Task> ExecutorWithFailure = (IReadOnlyList operations, CancellationToken cancellation) => { throw expectedException; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs index 5e64e7368a..baa02b69e4 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs @@ -21,7 +21,7 @@ public class CrossPartitionKeyBatchResponseTests public void StatusCodesAreSet() { const string errorMessage = "some error"; - CrossPartitionKeyBatchResponse response = new CrossPartitionKeyBatchResponse(HttpStatusCode.NotFound, SubStatusCodes.ClientTcpChannelFull, errorMessage, null); + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(HttpStatusCode.NotFound, SubStatusCodes.ClientTcpChannelFull, errorMessage, null); Assert.AreEqual(HttpStatusCode.NotFound, response.StatusCode); Assert.AreEqual(SubStatusCodes.ClientTcpChannelFull, response.SubStatusCode); Assert.AreEqual(errorMessage, response.ErrorMessage); @@ -60,7 +60,7 @@ public async Task StatusCodesAreSetThroughResponseAsync() batchRequest, new CosmosJsonDotNetSerializer()); - CrossPartitionKeyBatchResponse response = new CrossPartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } } From 9e75160644b3488e548d4485de3a6dd968d9ff27 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 1 Aug 2019 09:53:56 -0700 Subject: [PATCH 16/41] pending tasks --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 01ebc65a36..b9c4729c54 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -97,6 +97,12 @@ public void Dispose() this.cancellationTokenSource.Cancel(); this.currentBatcher?.Dispose(); this.currentTimer.CancelTimer(); + foreach (Task previousDispatchedTask in this.previousDispatchedTasks) + { +#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits + previousDispatchedTask.GetAwaiter().GetResult(); +#pragma warning restore VSTHRD002 // Avoid problematic synchronous waits + } } private void StartTimer() From 9e1a0df764979002586117e00a04830d84e39fa3 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 1 Aug 2019 10:22:51 -0700 Subject: [PATCH 17/41] Removing lock --- .../src/Batch/BatchAsyncContainerExecutor.cs | 12 ++++++------ .../src/Batch/BatchAsyncStreamer.cs | 8 +++----- .../Batch/BatchAsyncContainerExecutorTests.cs | 2 ++ 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index efd0d8d111..9a518b311d 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -23,7 +23,7 @@ namespace Microsoft.Azure.Cosmos /// It groups operations by Partition Key Range and sends them to the Batch API and then unifies results as they become available. It uses as batch processor and as batch executing handler. /// /// - internal class BatchAsyncContainerExecutor : IDisposable + internal class BatchAsyncContainerExecutor { private const int DefaultDispatchTimer = 10; private const int MinimumDispatchTimerInSeconds = 1; @@ -85,7 +85,7 @@ public async Task AddAsync( await this.ValidateOperationAsync(operation, itemRequestOptions, cancellationToken); string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false); - BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); + BatchAsyncStreamer streamer = await this.GetOrAddStreamerForPartitionKeyRangeAsync(resolvedPartitionKeyRangeId).ConfigureAwait(false); BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); await streamer.AddAsync(context).ConfigureAwait(false); return await context.Task; @@ -119,11 +119,11 @@ internal async Task ValidateOperationAsync( } } - public void Dispose() + public async Task DisposeAsync() { foreach (KeyValuePair streamer in this.streamersByPartitionKeyRange) { - streamer.Value.Dispose(); + await streamer.Value.DisposeAsync(); } this.timerPool.Dispose(); @@ -384,7 +384,7 @@ private async Task ExecuteServerRequestAsync(PartitionKeyServerBa } } - private BatchAsyncStreamer GetOrAddStreamerForPartitionKeyRange(string partitionKeyRangeId) + private async Task GetOrAddStreamerForPartitionKeyRangeAsync(string partitionKeyRangeId) { if (this.streamersByPartitionKeyRange.TryGetValue(partitionKeyRangeId, out BatchAsyncStreamer streamer)) { @@ -394,7 +394,7 @@ private BatchAsyncStreamer GetOrAddStreamerForPartitionKeyRange(string partition BatchAsyncStreamer newStreamer = new BatchAsyncStreamer(this.maxServerRequestOperationCount, this.maxServerRequestBodyLength, this.dispatchTimerInSeconds, this.timerPool, this.cosmosClientContext.CosmosSerializer, this.ExecuteAsync); if (!this.streamersByPartitionKeyRange.TryAdd(partitionKeyRangeId, newStreamer)) { - newStreamer.Dispose(); + await newStreamer.DisposeAsync(); } return this.streamersByPartitionKeyRange[partitionKeyRangeId]; diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index b9c4729c54..b7025a1bc1 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -18,7 +18,7 @@ namespace Microsoft.Azure.Cosmos /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. /// /// - internal class BatchAsyncStreamer : IDisposable + internal class BatchAsyncStreamer { private readonly List previousDispatchedTasks = new List(); private readonly SemaphoreSlim dispatchLimiter; @@ -91,7 +91,7 @@ public async Task AddAsync(BatchAsyncOperationContext context) } } - public void Dispose() + public async Task DisposeAsync() { this.disposed = true; this.cancellationTokenSource.Cancel(); @@ -99,9 +99,7 @@ public void Dispose() this.currentTimer.CancelTimer(); foreach (Task previousDispatchedTask in this.previousDispatchedTasks) { -#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits - previousDispatchedTask.GetAwaiter().GetResult(); -#pragma warning restore VSTHRD002 // Avoid problematic synchronous waits + await previousDispatchedTask; } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs index 720ab61481..4d1c882441 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs @@ -64,6 +64,8 @@ public async Task DoOperationsAsync() ItemResponse storedDoc = await this.cosmosContainer.ReadItemAsync(i.ToString(), new Cosmos.PartitionKey(i.ToString())); Assert.IsNotNull(storedDoc.Resource); } + + await executor.DisposeAsync(); } [TestMethod] From 0122a62d5a6a6eaf459d5b0d6c41b3c39f884fce Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 1 Aug 2019 10:56:42 -0700 Subject: [PATCH 18/41] Adding Dispose UTs --- .../Batch/BatchAsyncBatcherTests.cs | 10 +++++++++ .../Batch/BatchAsyncStreamerTests.cs | 22 +++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 549197fc0a..270b753bad 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -196,5 +196,15 @@ public async Task IsNotEmptyWithOperations() Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); Assert.IsFalse(batchAsyncBatcher.IsEmpty); } + + [TestMethod] + [Owner("maquaran")] + public async Task CannotAddToDisposedBatch() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + batchAsyncBatcher.Dispose(); + await Assert.ThrowsExceptionAsync(() => batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index e2ebea66d0..9269afb2a2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -151,6 +151,28 @@ public async Task DispatchesAsync() } } + [TestMethod] + [Owner("maquaran")] + public async Task DisposeAsyncShouldDisposeBatcher() + { + // Expect all operations to complete as their batches get dispached + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + List> contexts = new List>(10); + for (int i = 0; i < 10; i++) + { + var context = CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())); + await batchAsyncStreamer.AddAsync(context); + contexts.Add(context.Task); + } + + await Task.WhenAll(contexts); + + await batchAsyncStreamer.DisposeAsync(); + var newContext = CreateContext(new ItemBatchOperation(OperationType.Create, 0, "0")); + // Disposed batcher's internal cancellation was signaled + await Assert.ThrowsExceptionAsync(() => batchAsyncStreamer.AddAsync(newContext)); + } + private static BatchAsyncOperationContext CreateContext(ItemBatchOperation operation) => new BatchAsyncOperationContext(string.Empty, operation); } } From 30e71de650f5a7de69fd451e9228424c88302ea1 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 1 Aug 2019 16:04:14 -0700 Subject: [PATCH 19/41] Pushing retry to batcher --- .../src/Batch/BatchAsyncBatcher.cs | 16 ++- .../src/Batch/BatchAsyncContainerExecutor.cs | 99 +++++-------------- .../src/Batch/BatchAsyncStreamer.cs | 3 + .../src/Batch/PartitionKeyBatchResponse.cs | 6 ++ .../Batch/BatchAsyncBatcherTests.cs | 64 +++++++++++- .../Batch/BatchAsyncStreamerTests.cs | 2 +- ...s.cs => PartitionKeyBatchResponseTests.cs} | 58 ++++++++++- 7 files changed, 172 insertions(+), 76 deletions(-) rename Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/{CrossPartitionKeyBatchResponseTests.cs => PartitionKeyBatchResponseTests.cs} (51%) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 116b03363f..c11546a28a 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -99,7 +99,7 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati { try { - PartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken); + PartitionKeyBatchResponse batchResponse = await this.InvokeExecutorAsync(cancellationToken); // If the batch was not successful, we need to set all the responses if (!batchResponse.IsSuccessStatusCode) @@ -133,6 +133,20 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati public void Dispose() { this.cancellationTokenSource.Cancel(); + this.cancellationTokenSource.Dispose(); + this.tryAddLimiter.Dispose(); + } + + private async Task InvokeExecutorAsync(CancellationToken cancellationToken = default(CancellationToken)) + { + PartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken); + if (batchResponse.ContainsSplit()) + { + // Retry batch as it contained a split + return await this.executor(this.batchOperations, cancellationToken); + } + + return batchResponse; } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 9a518b311d..7895687683 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -126,38 +126,12 @@ public async Task DisposeAsync() await streamer.Value.DisposeAsync(); } - this.timerPool.Dispose(); - } - - private async Task ExecuteAsync( - IReadOnlyList operations, - CancellationToken cancellationToken) - { - List> inProgressTasksByPKRangeId = new List>(); - List completedResults = new List(); - - // Initial operations are all for the same PK Range Id - string pkRangeId = operations[0].PartitionKeyRangeId; - inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(pkRangeId, operations, cancellationToken)); - - while (inProgressTasksByPKRangeId.Count > 0) + foreach (KeyValuePair limiter in this.limitersByPartitionkeyRange) { - Task completedTask = await Task.WhenAny(inProgressTasksByPKRangeId); - PartitionKeyRangeBatchExecutionResult executionResult = await completedTask; - completedResults.Add(executionResult); - - // Pending operations can be due to splits, so we need to resolve the PK Range Ids - if (executionResult.PendingOperations != null) - { - IEnumerable retryOperations = BatchAsyncContainerExecutor.GetOperationsToRetry(operations, executionResult.PendingOperations.Select(o => o.OperationIndex)); - await this.GroupOperationsAndStartExecutionAsync(inProgressTasksByPKRangeId, retryOperations, cancellationToken); - } - - inProgressTasksByPKRangeId.Remove(completedTask); + limiter.Value.Dispose(); } - IEnumerable serverResponses = completedResults.SelectMany(res => res.ServerResponses); - return new PartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); + this.timerPool.Dispose(); } private static bool ValidateOperationEPK( @@ -187,16 +161,6 @@ private static bool ValidateOperationEPK( return true; } - private static IEnumerable GetOperationsToRetry( - IReadOnlyList operations, - IEnumerable indexes) - { - foreach (int index in indexes) - { - yield return operations[index]; - } - } - private static void AddHeadersToRequestMessage(RequestMessage requestMessage, string partitionKeyRangeId) { requestMessage.Headers.PartitionKeyRangeId = partitionKeyRangeId; @@ -221,18 +185,28 @@ private static IEnumerable GetOverflowOperations( return operationsSentToRequest.Skip(totalOperations - operationsThatOverflowed); } - private async Task GroupOperationsAndStartExecutionAsync( - List> taskContainer, - IEnumerable operations, + private async Task ExecuteAsync( + IReadOnlyList operations, CancellationToken cancellationToken) { + // Initially, all operations should be for the same PKRange, but this might be a retry after a split, and it might contain operations from + List> inProgressTasksByPKRangeId = new List>(); + List completedResults = new List(); Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, cancellationToken); - foreach (KeyValuePair> operationsForPKRangeId in operationsByPKRangeId) { - Task toAdd = this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellationToken); - taskContainer.Add(toAdd); + inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellationToken)); } + + while (inProgressTasksByPKRangeId.Count > 0) + { + Task completedTask = await Task.WhenAny(inProgressTasksByPKRangeId); + completedResults.Add(await completedTask); + inProgressTasksByPKRangeId.Remove(completedTask); + } + + IEnumerable serverResponses = completedResults.SelectMany(res => res.ServerResponses); + return new PartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); } private async Task>> GroupByPartitionKeyRangeIdsAsync( @@ -292,7 +266,7 @@ private async Task StartExecutionForPKRan IEnumerable operations, CancellationToken cancellationToken) { - SemaphoreSlim limiter = this.GetLimiterForPartitionKeyRange(pkRangeId); + SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(pkRangeId); await limiter.WaitAsync(cancellationToken).ConfigureAwait(false); try @@ -330,35 +304,16 @@ private async Task ExecuteServerRequestsA CancellationToken cancellationToken) { List serverResponses = new List(); - List overflowOperations = new List(); PartitionKeyServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operations, cancellationToken); - // In case some operations overflowed - overflowOperations.AddRange(BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operations)); - do + // In case some operations overflowed due to HybridRow serialization + IEnumerable overFlowOperations = BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operations); + serverResponses.Add(await this.ExecuteServerRequestAsync(serverRequest, cancellationToken)); + if (overFlowOperations.Any()) { - BatchResponse serverResponse = await this.ExecuteServerRequestAsync(serverRequest, cancellationToken); - if (serverResponse.StatusCode == HttpStatusCode.Gone - && (serverResponse.SubStatusCode == SubStatusCodes.CompletingSplit - || serverResponse.SubStatusCode == SubStatusCodes.CompletingPartitionMigration - || serverResponse.SubStatusCode == SubStatusCodes.PartitionKeyRangeGone)) - { - // lower layers would have refreshed the appropriate routing caches, but we need to retry the operations on the new PKRanges. - return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, serverResponses, new List(operations)); - } - else - { - serverResponses.Add(serverResponse); - } - - serverRequest = null; - if (overflowOperations.Count > 0) - { - serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, overflowOperations, cancellationToken); - overflowOperations.Clear(); - } + serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, overFlowOperations, cancellationToken); + serverResponses.Add(await this.ExecuteServerRequestAsync(serverRequest, cancellationToken)); } - while (serverRequest != null); return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, serverResponses); } @@ -400,7 +355,7 @@ private async Task GetOrAddStreamerForPartitionKeyRangeAsync return this.streamersByPartitionKeyRange[partitionKeyRangeId]; } - private SemaphoreSlim GetLimiterForPartitionKeyRange(string partitionKeyRangeId) + private SemaphoreSlim GetOrAddLimiterForPartitionKeyRange(string partitionKeyRangeId) { if (this.limitersByPartitionkeyRange.TryGetValue(partitionKeyRangeId, out SemaphoreSlim limiter)) { diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index b7025a1bc1..f53e9747a4 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -95,12 +95,15 @@ public async Task DisposeAsync() { this.disposed = true; this.cancellationTokenSource.Cancel(); + this.cancellationTokenSource.Dispose(); this.currentBatcher?.Dispose(); this.currentTimer.CancelTimer(); foreach (Task previousDispatchedTask in this.previousDispatchedTasks) { await previousDispatchedTask; } + + this.dispatchLimiter.Dispose(); } private void StartTimer() diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs index e7b3655a75..d6d2554cb2 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs @@ -155,6 +155,12 @@ internal override IEnumerable GetActivityIds() return this.ActivityIds; } + internal bool ContainsSplit() => this.ServerResponses != null && this.ServerResponses.Any(serverResponse => + (serverResponse.StatusCode == HttpStatusCode.Gone + && (serverResponse.SubStatusCode == SubStatusCodes.CompletingSplit + || serverResponse.SubStatusCode == SubStatusCodes.CompletingPartitionMigration + || serverResponse.SubStatusCode == SubStatusCodes.PartitionKeyRangeGone))); + /// /// Disposes the disposable members held. /// diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 270b753bad..f9903c1c04 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -12,6 +12,7 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Threading.Tasks; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; [TestClass] public class BatchAsyncBatcherTests @@ -63,6 +64,46 @@ private Func, CancellationToken, Task< throw expectedException; }; + private Func, CancellationToken, Task> ExecutorWithSplit + = async (IReadOnlyList operations, CancellationToken cancellation) => + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operations.Count]; + int index = 0; + foreach (BatchAsyncOperationContext operation in operations) + { + results.Add( + new BatchOperationResult(HttpStatusCode.Gone) + { + ETag = operation.Operation.Id, + SubStatusCode = SubStatusCodes.PartitionKeyRangeGone + }); + + arrayOperations[index++] = operation.Operation; + } + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: (int)responseContent.Length * operations.Count, + maxOperationCount: operations.Count, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: cancellation); + + ResponseMessage responseMessage = new ResponseMessage(HttpStatusCode.Gone) { Content = responseContent }; + responseMessage.Headers.SubStatusCode = SubStatusCodes.PartitionKeyRangeGone; + + BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( + responseMessage, + batchRequest, + new CosmosJsonDotNetSerializer()); + + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + return response; + }; + [DataTestMethod] [Owner("maquaran")] [ExpectedException(typeof(ArgumentOutOfRangeException))] @@ -204,7 +245,28 @@ public async Task CannotAddToDisposedBatch() BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); batchAsyncBatcher.Dispose(); - await Assert.ThrowsExceptionAsync(() => batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + await Assert.ThrowsExceptionAsync(() => batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + } + + [TestMethod] + [Owner("maquaran")] + public async Task RetryOnSplit() + { + int executeCount = 0; + Func, CancellationToken, Task> executor = (IReadOnlyList operations, CancellationToken cancellationToken) => + { + if (executeCount ++ == 0) + { + return this.ExecutorWithSplit(operations, cancellationToken); + } + + return this.Executor(operations, cancellationToken); + }; + + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), executor); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + await batchAsyncBatcher.DispatchAsync(); + Assert.AreEqual(2, executeCount); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 9269afb2a2..6168d70886 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -170,7 +170,7 @@ public async Task DisposeAsyncShouldDisposeBatcher() await batchAsyncStreamer.DisposeAsync(); var newContext = CreateContext(new ItemBatchOperation(OperationType.Create, 0, "0")); // Disposed batcher's internal cancellation was signaled - await Assert.ThrowsExceptionAsync(() => batchAsyncStreamer.AddAsync(newContext)); + await Assert.ThrowsExceptionAsync(() => batchAsyncStreamer.AddAsync(newContext)); } private static BatchAsyncOperationContext CreateContext(ItemBatchOperation operation) => new BatchAsyncOperationContext(string.Empty, operation); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs similarity index 51% rename from Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs rename to Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs index baa02b69e4..4d9e5067e5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/CrossPartitionKeyBatchResponseTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs @@ -14,7 +14,7 @@ namespace Microsoft.Azure.Cosmos.Tests using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] - public class CrossPartitionKeyBatchResponseTests + public class PartitionKeyBatchResponseTests { [TestMethod] [Owner("maquaran")] @@ -27,6 +27,23 @@ public void StatusCodesAreSet() Assert.AreEqual(errorMessage, response.ErrorMessage); } + [TestMethod] + [Owner("maquaran")] + public async Task ConstainsSplitIsTrue() + { + Assert.IsTrue(await ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.CompletingSplit)); + Assert.IsTrue(await ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.CompletingPartitionMigration)); + Assert.IsTrue(await ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.PartitionKeyRangeGone)); + } + + [TestMethod] + [Owner("maquaran")] + public async Task ConstainsSplitIsFalse() + { + Assert.IsFalse(await ConstainsSplitIsTrueInternal(HttpStatusCode.OK, SubStatusCodes.Unknown)); + Assert.IsFalse(await ConstainsSplitIsTrueInternal((HttpStatusCode)429, SubStatusCodes.Unknown)); + } + [TestMethod] [Owner("maquaran")] public async Task StatusCodesAreSetThroughResponseAsync() @@ -63,5 +80,44 @@ public async Task StatusCodesAreSetThroughResponseAsync() PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } + + private async Task ConstainsSplitIsTrueInternal(HttpStatusCode statusCode, SubStatusCodes subStatusCode) + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; + + ItemBatchOperation operation = new ItemBatchOperation(OperationType.AddComputeGatewayRequestCharges, 0, "0"); + + results.Add( + new BatchOperationResult(HttpStatusCode.OK) + { + ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), + ETag = operation.Id + }); + + arrayOperations[0] = operation; + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: 100, + maxOperationCount: 1, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: default(CancellationToken)); + + ResponseMessage response = new ResponseMessage(statusCode) { Content = responseContent }; + response.Headers.SubStatusCode = subStatusCode; + + BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( + response, + batchRequest, + new CosmosJsonDotNetSerializer()); + + PartitionKeyBatchResponse pkResponse = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + + return pkResponse.ContainsSplit(); + } } } From 8e968fd67fd8edca447c566a20b9c131dd76c09a Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 1 Aug 2019 16:08:13 -0700 Subject: [PATCH 20/41] Missing using --- .../src/Batch/BatchAsyncBatcher.cs | 32 ++++++++++--------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index c11546a28a..c9d0a28051 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -99,25 +99,26 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati { try { - PartitionKeyBatchResponse batchResponse = await this.InvokeExecutorAsync(cancellationToken); - - // If the batch was not successful, we need to set all the responses - if (!batchResponse.IsSuccessStatusCode) + using (PartitionKeyBatchResponse batchResponse = await this.InvokeExecutorAsync(cancellationToken)) { - BatchOperationResult errorResult = new BatchOperationResult(batchResponse.StatusCode); - foreach (BatchAsyncOperationContext operation in this.batchOperations) + // If the batch was not successful, we need to set all the responses + if (!batchResponse.IsSuccessStatusCode) { - operation.Complete(errorResult); - } + BatchOperationResult errorResult = new BatchOperationResult(batchResponse.StatusCode); + foreach (BatchAsyncOperationContext operation in this.batchOperations) + { + operation.Complete(errorResult); + } - return; - } + return; + } - for (int index = 0; index < this.batchOperations.Count; index++) - { - BatchAsyncOperationContext operation = this.batchOperations[index]; - BatchOperationResult response = batchResponse[index]; - operation.Complete(response); + for (int index = 0; index < this.batchOperations.Count; index++) + { + BatchAsyncOperationContext operation = this.batchOperations[index]; + BatchOperationResult response = batchResponse[index]; + operation.Complete(response); + } } } catch (Exception ex) @@ -142,6 +143,7 @@ public void Dispose() PartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken); if (batchResponse.ContainsSplit()) { + batchResponse.Dispose(); // Retry batch as it contained a split return await this.executor(this.batchOperations, cancellationToken); } From 7eed1c229f7beb6f2b95059f6d3a5e723f675141 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 5 Aug 2019 11:38:51 -0700 Subject: [PATCH 21/41] Removing unnecessary reprocesing of PKRange --- .../src/Batch/BatchAsyncBatcher.cs | 15 +----- .../src/Batch/BatchAsyncContainerExecutor.cs | 49 ++++++++++++++----- .../src/Batch/BatchAsyncStreamer.cs | 16 +++--- .../Batch/BatchAsyncBatcherTests.cs | 21 -------- 4 files changed, 47 insertions(+), 54 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index c9d0a28051..9103932d87 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -99,7 +99,7 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati { try { - using (PartitionKeyBatchResponse batchResponse = await this.InvokeExecutorAsync(cancellationToken)) + using (PartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken)) { // If the batch was not successful, we need to set all the responses if (!batchResponse.IsSuccessStatusCode) @@ -137,18 +137,5 @@ public void Dispose() this.cancellationTokenSource.Dispose(); this.tryAddLimiter.Dispose(); } - - private async Task InvokeExecutorAsync(CancellationToken cancellationToken = default(CancellationToken)) - { - PartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken); - if (batchResponse.ContainsSplit()) - { - batchResponse.Dispose(); - // Retry batch as it contained a split - return await this.executor(this.batchOperations, cancellationToken); - } - - return batchResponse; - } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 7895687683..3fec07d2fc 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos using System.Diagnostics; using System.IO; using System.Linq; - using System.Net; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Routing; @@ -185,17 +184,34 @@ private static IEnumerable GetOverflowOperations( return operationsSentToRequest.Skip(totalOperations - operationsThatOverflowed); } - private async Task ExecuteAsync( + private Task ExecuteAsync( IReadOnlyList operations, CancellationToken cancellationToken) { - // Initially, all operations should be for the same PKRange, but this might be a retry after a split, and it might contain operations from + return this.ExecuteInternalAsync(operations, false, cancellationToken); + } + + private async Task ExecuteInternalAsync( + IReadOnlyList operations, + bool refreshPartitionKeyRanges, + CancellationToken cancellationToken) + { List> inProgressTasksByPKRangeId = new List>(); List completedResults = new List(); - Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, cancellationToken); - foreach (KeyValuePair> operationsForPKRangeId in operationsByPKRangeId) + + // Initially, all operations should be for the same PKRange + if (!refreshPartitionKeyRanges) { - inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellationToken)); + inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(operations[0].PartitionKeyRangeId, operations, cancellationToken)); + } + else + { + // If a refresh was requested, recalculate PKRange + Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, refreshPartitionKeyRanges, cancellationToken); + foreach (KeyValuePair> operationsForPKRangeId in operationsByPKRangeId) + { + inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellationToken)); + } } while (inProgressTasksByPKRangeId.Count > 0) @@ -206,28 +222,37 @@ private async Task ExecuteAsync( } IEnumerable serverResponses = completedResults.SelectMany(res => res.ServerResponses); - return new PartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); + PartitionKeyBatchResponse batchResponse = new PartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); + if (batchResponse.ContainsSplit()) + { + batchResponse.Dispose(); + // Retry batch as it contained a split + return await this.ExecuteInternalAsync(operations, true, cancellationToken); + } + + return batchResponse; } private async Task>> GroupByPartitionKeyRangeIdsAsync( - IEnumerable operations, + IEnumerable contexts, + bool refreshPartitionKeyRanges, CancellationToken cancellationToken) { PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken).ConfigureAwait(false); CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken).ConfigureAwait(false); Dictionary> operationsByPKRangeId = new Dictionary>(); - foreach (BatchAsyncOperationContext operation in operations) + foreach (BatchAsyncOperationContext context in contexts) { - string partitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation.Operation, cancellationToken).ConfigureAwait(false); + string partitionKeyRangeId = BatchExecUtils.GetPartitionKeyRangeId(context.Operation.PartitionKey.Value, partitionKeyDefinition, collectionRoutingMap); if (operationsByPKRangeId.TryGetValue(partitionKeyRangeId, out List operationsForPKRangeId)) { - operationsForPKRangeId.Add(operation); + operationsForPKRangeId.Add(context); } else { - operationsByPKRangeId.Add(partitionKeyRangeId, new List() { operation }); + operationsByPKRangeId.Add(partitionKeyRangeId, new List() { context }); } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index f53e9747a4..418e5a5127 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -26,13 +26,12 @@ internal class BatchAsyncStreamer private readonly int maxBatchByteSize; private readonly Func, CancellationToken, Task> executor; private readonly int dispatchTimerInSeconds; - private readonly CosmosSerializer CosmosSerializer; + private readonly CosmosSerializer cosmosSerializer; private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); private BatchAsyncBatcher currentBatcher; private TimerPool timerPool; private PooledTimer currentTimer; private Task timerTask; - private bool disposed; public BatchAsyncStreamer( int maxBatchOperationCount, @@ -72,7 +71,7 @@ public BatchAsyncStreamer( this.executor = executor; this.dispatchTimerInSeconds = dispatchTimerInSeconds; this.timerPool = timerPool; - this.CosmosSerializer = cosmosSerializer; + this.cosmosSerializer = cosmosSerializer; this.dispatchLimiter = new SemaphoreSlim(1, 1); this.currentBatcher = this.CreateBatchAsyncBatcher(); @@ -85,7 +84,11 @@ public async Task AddAsync(BatchAsyncOperationContext context) { // Batcher is full BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAndCreateAsync(); - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + if (toDispatch != null) + { + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + } + bool addedContext = await this.currentBatcher.TryAddAsync(context); Debug.Assert(addedContext, "Could not add context to batcher."); } @@ -93,7 +96,6 @@ public async Task AddAsync(BatchAsyncOperationContext context) public async Task DisposeAsync() { - this.disposed = true; this.cancellationTokenSource.Cancel(); this.cancellationTokenSource.Dispose(); this.currentBatcher?.Dispose(); @@ -117,7 +119,7 @@ private void StartTimer() private async Task DispatchTimerAsync() { - if (this.disposed) + if (this.cancellationTokenSource.IsCancellationRequested) { return; } @@ -153,7 +155,7 @@ private async Task GetBatchToDispatchAndCreateAsync() private BatchAsyncBatcher CreateBatchAsyncBatcher() { - return new BatchAsyncBatcher(this.maxBatchOperationCount, this.maxBatchByteSize, this.CosmosSerializer, this.executor); + return new BatchAsyncBatcher(this.maxBatchOperationCount, this.maxBatchByteSize, this.cosmosSerializer, this.executor); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index f9903c1c04..511dc1f738 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -247,26 +247,5 @@ public async Task CannotAddToDisposedBatch() batchAsyncBatcher.Dispose(); await Assert.ThrowsExceptionAsync(() => batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); } - - [TestMethod] - [Owner("maquaran")] - public async Task RetryOnSplit() - { - int executeCount = 0; - Func, CancellationToken, Task> executor = (IReadOnlyList operations, CancellationToken cancellationToken) => - { - if (executeCount ++ == 0) - { - return this.ExecutorWithSplit(operations, cancellationToken); - } - - return this.Executor(operations, cancellationToken); - }; - - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), executor); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - await batchAsyncBatcher.DispatchAsync(); - Assert.AreEqual(2, executeCount); - } } } From 80616bcd025e1bb5ef9fa245648f994a4ded2c36 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 5 Aug 2019 12:29:41 -0700 Subject: [PATCH 22/41] New UTs for Executor --- .../src/Resource/Container/ContainerCore.cs | 4 +- .../Batch/BatchAsyncContainerExecutorTests.cs | 198 ++++++++++++++++++ .../Microsoft.Azure.Cosmos.Tests.csproj | 3 - 3 files changed, 200 insertions(+), 5 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs index 1e632efa2b..39d1108812 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs @@ -196,7 +196,7 @@ internal async Task GetRIDAsync(CancellationToken cancellationToken) return containerProperties?.ResourceId; } - internal Task GetPartitionKeyDefinitionAsync(CancellationToken cancellationToken = default(CancellationToken)) + internal virtual Task GetPartitionKeyDefinitionAsync(CancellationToken cancellationToken = default(CancellationToken)) { return this.GetCachedContainerPropertiesAsync(cancellationToken) .ContinueWith(containerPropertiesTask => containerPropertiesTask.Result?.PartitionKey, cancellationToken); @@ -239,7 +239,7 @@ internal async Task GetRIDAsync(CancellationToken cancellationToken) return containerProperties.GetNoneValue(); } - internal Task GetRoutingMapAsync(CancellationToken cancellationToken) + internal virtual Task GetRoutingMapAsync(CancellationToken cancellationToken) { string collectionRID = null; return this.GetRIDAsync(cancellationToken) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs new file mode 100644 index 0000000000..cc6801fc0b --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs @@ -0,0 +1,198 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Collections.Generic; + using System.Collections.ObjectModel; + using System.IO; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; + + [TestClass] + public class BatchAsyncContainerExecutorTests + { + private static CosmosSerializer cosmosDefaultJsonSerializer = new CosmosJsonDotNetSerializer(); + + [TestMethod] + public async Task RetryOnSplit() + { + ItemBatchOperation itemBatchOperation = CreateItem("test"); + + Mock mockedContext = new Mock(); + mockedContext + .SetupSequence(c => c.ProcessResourceOperationStreamAsync( + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny())) + .Returns(GenerateSplitResponseAsync(itemBatchOperation)) + .Returns(GenerateOkResponseAsync(itemBatchOperation)); + + mockedContext.Setup(c => c.CosmosSerializer).Returns(new CosmosJsonDotNetSerializer()); + + Uri link = new Uri($"/dbs/db/colls/colls", UriKind.Relative); + Mock mockContainer = new Mock(); + mockContainer.Setup(x => x.LinkUri).Returns(link); + mockContainer.Setup(x => x.GetPartitionKeyDefinitionAsync(It.IsAny())).Returns(Task.FromResult(new PartitionKeyDefinition() { Paths = new Collection() { "/id" } })); + + CollectionRoutingMap routingMap = CollectionRoutingMap.TryCreateCompleteRoutingMap( + new[] + { + Tuple.Create(new PartitionKeyRange{ Id = "0", MinInclusive = "", MaxExclusive = "FF"}, (ServiceIdentity)null) + }, + string.Empty); + mockContainer.Setup(x => x.GetRoutingMapAsync(It.IsAny())).Returns(Task.FromResult(routingMap)); + BatchAsyncContainerExecutor executor = new BatchAsyncContainerExecutor(mockContainer.Object, mockedContext.Object, 20, Constants.MaxDirectModeBatchRequestBodySizeInBytes, 1); + BatchOperationResult result = await executor.AddAsync(itemBatchOperation); + + Mock.Get(mockContainer.Object) + .Verify(x => x.GetPartitionKeyDefinitionAsync(It.IsAny()), Times.Exactly(2)); + Mock.Get(mockedContext.Object) + .Verify(c => c.ProcessResourceOperationStreamAsync( + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny()), Times.Exactly(2)); + Assert.AreEqual(HttpStatusCode.OK, result.StatusCode); + } + + [TestMethod] + public async Task DoesNotRecalculatePartitionKeyRangeOnNoSplits() + { + ItemBatchOperation itemBatchOperation = CreateItem("test"); + + Mock mockedContext = new Mock(); + mockedContext + .Setup(c => c.ProcessResourceOperationStreamAsync( + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny())) + .Returns(GenerateOkResponseAsync(itemBatchOperation)); + + mockedContext.Setup(c => c.CosmosSerializer).Returns(new CosmosJsonDotNetSerializer()); + + Uri link = new Uri($"/dbs/db/colls/colls", UriKind.Relative); + Mock mockContainer = new Mock(); + mockContainer.Setup(x => x.LinkUri).Returns(link); + mockContainer.Setup(x => x.GetPartitionKeyDefinitionAsync(It.IsAny())).Returns(Task.FromResult(new PartitionKeyDefinition() { Paths = new Collection() { "/id" } })); + + CollectionRoutingMap routingMap = CollectionRoutingMap.TryCreateCompleteRoutingMap( + new[] + { + Tuple.Create(new PartitionKeyRange{ Id = "0", MinInclusive = "", MaxExclusive = "FF"}, (ServiceIdentity)null) + }, + string.Empty); + mockContainer.Setup(x => x.GetRoutingMapAsync(It.IsAny())).Returns(Task.FromResult(routingMap)); + BatchAsyncContainerExecutor executor = new BatchAsyncContainerExecutor(mockContainer.Object, mockedContext.Object, 20, Constants.MaxDirectModeBatchRequestBodySizeInBytes, 1); + BatchOperationResult result = await executor.AddAsync(itemBatchOperation); + + Mock.Get(mockContainer.Object) + .Verify(x => x.GetPartitionKeyDefinitionAsync(It.IsAny()), Times.Once); + Mock.Get(mockedContext.Object) + .Verify(c => c.ProcessResourceOperationStreamAsync( + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny()), Times.Once); + Assert.AreEqual(HttpStatusCode.OK, result.StatusCode); + } + + private async Task GenerateSplitResponseAsync(ItemBatchOperation itemBatchOperation) + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; + results.Add( + new BatchOperationResult(HttpStatusCode.Gone) + { + ETag = itemBatchOperation.Id, + SubStatusCode = SubStatusCodes.PartitionKeyRangeGone + }); + + arrayOperations[0] = itemBatchOperation; + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: 100, + maxOperationCount: 1, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: CancellationToken.None); + + ResponseMessage responseMessage = new ResponseMessage(HttpStatusCode.Gone) { Content = responseContent }; + responseMessage.Headers.SubStatusCode = SubStatusCodes.PartitionKeyRangeGone; + return responseMessage; + } + + private async Task GenerateOkResponseAsync(ItemBatchOperation itemBatchOperation) + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; + results.Add( + new BatchOperationResult(HttpStatusCode.OK) + { + ETag = itemBatchOperation.Id + }); + + arrayOperations[0] = itemBatchOperation; + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: 100, + maxOperationCount: 1, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: CancellationToken.None); + + ResponseMessage responseMessage = new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }; + return responseMessage; + } + + private static ItemBatchOperation CreateItem(string id) + { + MyDocument myDocument = new MyDocument() { id = id, Status = id }; + return new ItemBatchOperation(OperationType.Create, 0, new Cosmos.PartitionKey(id), id, cosmosDefaultJsonSerializer.ToStream(myDocument)); + } + + private class MyDocument + { + public string id { get; set; } + + public string Status { get; set; } + + public bool Updated { get; set; } + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj index a8673d5e94..de237b021e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj @@ -157,9 +157,6 @@ PreserveNewest - - - C:\Users\abpai\.nuget\packages\microsoft.azure.cosmos.direct.myget\3.0.0.33-preview\runtimes\any\lib\netstandard2.0\Microsoft.Azure.Cosmos.Core.dll From 133fb9e92a541bd1d05f02af8d01e703ca8e62f7 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Wed, 7 Aug 2019 11:18:57 -0700 Subject: [PATCH 23/41] Support partial splits --- .../src/Batch/BatchAsyncContainerExecutor.cs | 44 ++++--- .../src/Batch/BatchExecUtils.cs | 2 - .../src/Batch/PartitionKeyBatchResponse.cs | 17 +-- .../PartitionKeyRangeBatchExecutionResult.cs | 19 ++- .../Batch/BatchAsyncBatcherTests.cs | 14 +-- .../Batch/BatchAsyncContainerExecutorTests.cs | 6 +- .../Batch/BatchAsyncOperationContextTests.cs | 4 - .../Batch/BatchAsyncStreamerTests.cs | 8 -- .../Batch/PartitionKeyBatchResponseTests.cs | 58 ---------- ...titionKeyRangeBatchExecutionResultTests.cs | 109 ++++++++++++++++++ 10 files changed, 161 insertions(+), 120 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 3fec07d2fc..d5e864c504 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -184,14 +184,28 @@ private static IEnumerable GetOverflowOperations( return operationsSentToRequest.Skip(totalOperations - operationsThatOverflowed); } - private Task ExecuteAsync( + private static IReadOnlyList GetOperationsToRetry( + IReadOnlyList baseOperations, + IEnumerable indexes) + { + List operations = new List(); + foreach (int index in indexes) + { + operations.Add(baseOperations[index]); + } + + return operations; + } + + private async Task ExecuteAsync( IReadOnlyList operations, CancellationToken cancellationToken) { - return this.ExecuteInternalAsync(operations, false, cancellationToken); + IEnumerable responses = await this.ExecuteInternalAsync(operations, false, cancellationToken); + return new PartitionKeyBatchResponse(responses, this.cosmosClientContext.CosmosSerializer); } - private async Task ExecuteInternalAsync( + private async Task> ExecuteInternalAsync( IReadOnlyList operations, bool refreshPartitionKeyRanges, CancellationToken cancellationToken) @@ -207,7 +221,7 @@ private async Task ExecuteInternalAsync( else { // If a refresh was requested, recalculate PKRange - Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, refreshPartitionKeyRanges, cancellationToken); + Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, cancellationToken); foreach (KeyValuePair> operationsForPKRangeId in operationsByPKRangeId) { inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellationToken)); @@ -221,21 +235,25 @@ private async Task ExecuteInternalAsync( inProgressTasksByPKRangeId.Remove(completedTask); } - IEnumerable serverResponses = completedResults.SelectMany(res => res.ServerResponses); - PartitionKeyBatchResponse batchResponse = new PartitionKeyBatchResponse(serverResponses, this.cosmosClientContext.CosmosSerializer); - if (batchResponse.ContainsSplit()) + List responses = new List(operations.Count); + foreach (PartitionKeyRangeBatchExecutionResult result in completedResults) { - batchResponse.Dispose(); - // Retry batch as it contained a split - return await this.ExecuteInternalAsync(operations, true, cancellationToken); + if (!result.ContainsSplit()) + { + responses.AddRange(result.ServerResponses); + } + else + { + IReadOnlyList retryOperations = BatchAsyncContainerExecutor.GetOperationsToRetry(operations, result.Operations.Select(o => o.OperationIndex)); + responses.AddRange(await this.ExecuteInternalAsync(retryOperations, true, cancellationToken)); + } } - return batchResponse; + return responses; } private async Task>> GroupByPartitionKeyRangeIdsAsync( IEnumerable contexts, - bool refreshPartitionKeyRanges, CancellationToken cancellationToken) { PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken).ConfigureAwait(false); @@ -340,7 +358,7 @@ private async Task ExecuteServerRequestsA serverResponses.Add(await this.ExecuteServerRequestAsync(serverRequest, cancellationToken)); } - return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, serverResponses); + return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, operations, serverResponses); } private async Task ExecuteServerRequestAsync(PartitionKeyServerBatchRequest serverRequest, CancellationToken cancellationToken) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs index d4171db71c..655674208d 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchExecUtils.cs @@ -7,7 +7,6 @@ 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.Documents; @@ -17,7 +16,6 @@ namespace Microsoft.Azure.Cosmos /// internal static class BatchExecUtils { - internal const int StatusCodeMultiStatus = 207; // Using the same buffer size as the Stream.DefaultCopyBufferSize private const int BufferSize = 81920; diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs index d6d2554cb2..b95fba3fbd 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs @@ -14,14 +14,7 @@ namespace Microsoft.Azure.Cosmos /// /// Response of a cross partition key batch request. /// -#pragma warning disable CA1710 // Identifiers should have correct suffix - #if PREVIEW - public -#else - internal -#endif - class PartitionKeyBatchResponse : BatchResponse -#pragma warning restore CA1710 // Identifiers should have correct suffix + internal class PartitionKeyBatchResponse : BatchResponse { // Results sorted in the order operations had been added. private readonly SortedList resultsByOperationIndex; @@ -52,7 +45,7 @@ internal PartitionKeyBatchResponse( internal PartitionKeyBatchResponse(IEnumerable serverResponses, CosmosSerializer serializer) { this.StatusCode = serverResponses.Any(r => r.StatusCode != HttpStatusCode.OK) - ? (HttpStatusCode)BatchExecUtils.StatusCodeMultiStatus + ? (HttpStatusCode)StatusCodes.MultiStatus : HttpStatusCode.OK; this.ServerResponses = serverResponses; @@ -155,12 +148,6 @@ internal override IEnumerable GetActivityIds() return this.ActivityIds; } - internal bool ContainsSplit() => this.ServerResponses != null && this.ServerResponses.Any(serverResponse => - (serverResponse.StatusCode == HttpStatusCode.Gone - && (serverResponse.SubStatusCode == SubStatusCodes.CompletingSplit - || serverResponse.SubStatusCode == SubStatusCodes.CompletingPartitionMigration - || serverResponse.SubStatusCode == SubStatusCodes.PartitionKeyRangeGone))); - /// /// Disposes the disposable members held. /// diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs index 0dc19135d1..cfe12e1a78 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs @@ -5,20 +5,31 @@ namespace Microsoft.Azure.Cosmos { using System.Collections.Generic; + using System.Linq; + using System.Net; internal class PartitionKeyRangeBatchExecutionResult { public string PartitionKeyRangeId { get; } - public List ServerResponses { get; } + public IReadOnlyList ServerResponses { get; } - public List PendingOperations { get; set; } + public IEnumerable Operations { get; } - public PartitionKeyRangeBatchExecutionResult(string pkRangeId, List serverResponses, List pendingOperations = null) + public PartitionKeyRangeBatchExecutionResult( + string pkRangeId, + IEnumerable operations, + List serverResponses) { this.PartitionKeyRangeId = pkRangeId; this.ServerResponses = serverResponses; - this.PendingOperations = pendingOperations; + this.Operations = operations; } + + internal bool ContainsSplit() => this.ServerResponses != null && this.ServerResponses.Any(serverResponse => + serverResponse.StatusCode == HttpStatusCode.Gone + && (serverResponse.SubStatusCode == Documents.SubStatusCodes.CompletingSplit + || serverResponse.SubStatusCode == Documents.SubStatusCodes.CompletingPartitionMigration + || serverResponse.SubStatusCode == Documents.SubStatusCodes.PartitionKeyRangeGone)); } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 511dc1f738..03a0e25a23 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -105,7 +105,6 @@ private Func, CancellationToken, Task< }; [DataTestMethod] - [Owner("maquaran")] [ExpectedException(typeof(ArgumentOutOfRangeException))] [DataRow(0)] [DataRow(-1)] @@ -115,7 +114,6 @@ public void ValidatesSize(int size) } [DataTestMethod] - [Owner("maquaran")] [ExpectedException(typeof(ArgumentOutOfRangeException))] [DataRow(0)] [DataRow(-1)] @@ -125,7 +123,6 @@ public void ValidatesByteSize(int size) } [TestMethod] - [Owner("maquaran")] [ExpectedException(typeof(ArgumentNullException))] public void ValidatesExecutor() { @@ -133,7 +130,6 @@ public void ValidatesExecutor() } [TestMethod] - [Owner("maquaran")] [ExpectedException(typeof(ArgumentNullException))] public void ValidatesSerializer() { @@ -141,7 +137,6 @@ public void ValidatesSerializer() } [TestMethod] - [Owner("maquaran")] public async Task HasFixedSize() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); @@ -151,10 +146,9 @@ public async Task HasFixedSize() } [TestMethod] - [Owner("maquaran")] public async Task HasFixedByteSize() { - await ItemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), CancellationToken.None); + await this.ItemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), CancellationToken.None); // Each operation is 2 bytes BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor); Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); @@ -163,7 +157,6 @@ public async Task HasFixedByteSize() } [TestMethod] - [Owner("maquaran")] public void TryAddIsThreadSafe() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); @@ -181,7 +174,6 @@ public void TryAddIsThreadSafe() } [TestMethod] - [Owner("maquaran")] public async Task ExceptionsFailOperationsAsync() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); @@ -198,7 +190,6 @@ public async Task ExceptionsFailOperationsAsync() } [TestMethod] - [Owner("maquaran")] public async Task DispatchProcessInOrderAsync() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor); @@ -222,7 +213,6 @@ public async Task DispatchProcessInOrderAsync() } [TestMethod] - [Owner("maquaran")] public void IsEmptyWithNoOperations() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor); @@ -230,7 +220,6 @@ public void IsEmptyWithNoOperations() } [TestMethod] - [Owner("maquaran")] public async Task IsNotEmptyWithOperations() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); @@ -239,7 +228,6 @@ public async Task IsNotEmptyWithOperations() } [TestMethod] - [Owner("maquaran")] public async Task CannotAddToDisposedBatch() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs index cc6801fc0b..ab6b342f88 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs @@ -38,8 +38,8 @@ public async Task RetryOnSplit() It.IsAny(), It.IsAny>(), It.IsAny())) - .Returns(GenerateSplitResponseAsync(itemBatchOperation)) - .Returns(GenerateOkResponseAsync(itemBatchOperation)); + .Returns(this.GenerateSplitResponseAsync(itemBatchOperation)) + .Returns(this.GenerateOkResponseAsync(itemBatchOperation)); mockedContext.Setup(c => c.CosmosSerializer).Returns(new CosmosJsonDotNetSerializer()); @@ -91,7 +91,7 @@ public async Task DoesNotRecalculatePartitionKeyRangeOnNoSplits() It.IsAny(), It.IsAny>(), It.IsAny())) - .Returns(GenerateOkResponseAsync(itemBatchOperation)); + .Returns(this.GenerateOkResponseAsync(itemBatchOperation)); mockedContext.Setup(c => c.CosmosSerializer).Returns(new CosmosJsonDotNetSerializer()); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs index 3fb5e69128..8017fa3587 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs @@ -14,7 +14,6 @@ namespace Microsoft.Azure.Cosmos.Tests public class BatchAsyncOperationContextTests { [TestMethod] - [Owner("maquaran")] public void PartitionKeyRangeIdIsSetOnInitialization() { string expectedPkRangeId = Guid.NewGuid().ToString(); @@ -28,7 +27,6 @@ public void PartitionKeyRangeIdIsSetOnInitialization() } [TestMethod] - [Owner("maquaran")] public void TaskIsCreatedOnInitialization() { ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); @@ -40,7 +38,6 @@ public void TaskIsCreatedOnInitialization() } [TestMethod] - [Owner("maquaran")] public async Task TaskResultIsSetOnCompleteAsync() { ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); @@ -55,7 +52,6 @@ public async Task TaskResultIsSetOnCompleteAsync() } [TestMethod] - [Owner("maquaran")] public async Task ExceptionIsSetOnFailAsync() { Exception failure = new Exception("It failed"); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 6168d70886..f94cd222ae 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -67,7 +67,6 @@ private Func, CancellationToken, Task< }; [DataTestMethod] - [Owner("maquaran")] [ExpectedException(typeof(ArgumentOutOfRangeException))] [DataRow(0)] [DataRow(-1)] @@ -77,7 +76,6 @@ public void ValidatesSize(int size) } [DataTestMethod] - [Owner("maquaran")] [ExpectedException(typeof(ArgumentOutOfRangeException))] [DataRow(0)] [DataRow(-1)] @@ -87,7 +85,6 @@ public void ValidatesDispatchTimer(int dispatchTimerInSeconds) } [TestMethod] - [Owner("maquaran")] [ExpectedException(typeof(ArgumentNullException))] public void ValidatesExecutor() { @@ -95,7 +92,6 @@ public void ValidatesExecutor() } [TestMethod] - [Owner("maquaran")] [ExpectedException(typeof(ArgumentNullException))] public void ValidatesSerializer() { @@ -103,7 +99,6 @@ public void ValidatesSerializer() } [TestMethod] - [Owner("maquaran")] public async Task ExceptionsOnBatchBubbleUpAsync() { BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); @@ -114,7 +109,6 @@ public async Task ExceptionsOnBatchBubbleUpAsync() } [TestMethod] - [Owner("maquaran")] public async Task TimerDispatchesAsync() { // Bigger batch size than the amount of operations, timer should dispatch @@ -127,7 +121,6 @@ public async Task TimerDispatchesAsync() } [TestMethod] - [Owner("maquaran")] public async Task DispatchesAsync() { // Expect all operations to complete as their batches get dispached @@ -152,7 +145,6 @@ public async Task DispatchesAsync() } [TestMethod] - [Owner("maquaran")] public async Task DisposeAsyncShouldDisposeBatcher() { // Expect all operations to complete as their batches get dispached diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs index 4d9e5067e5..f3b39ac891 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs @@ -17,7 +17,6 @@ namespace Microsoft.Azure.Cosmos.Tests public class PartitionKeyBatchResponseTests { [TestMethod] - [Owner("maquaran")] public void StatusCodesAreSet() { const string errorMessage = "some error"; @@ -28,24 +27,6 @@ public void StatusCodesAreSet() } [TestMethod] - [Owner("maquaran")] - public async Task ConstainsSplitIsTrue() - { - Assert.IsTrue(await ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.CompletingSplit)); - Assert.IsTrue(await ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.CompletingPartitionMigration)); - Assert.IsTrue(await ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.PartitionKeyRangeGone)); - } - - [TestMethod] - [Owner("maquaran")] - public async Task ConstainsSplitIsFalse() - { - Assert.IsFalse(await ConstainsSplitIsTrueInternal(HttpStatusCode.OK, SubStatusCodes.Unknown)); - Assert.IsFalse(await ConstainsSplitIsTrueInternal((HttpStatusCode)429, SubStatusCodes.Unknown)); - } - - [TestMethod] - [Owner("maquaran")] public async Task StatusCodesAreSetThroughResponseAsync() { List results = new List(); @@ -80,44 +61,5 @@ public async Task StatusCodesAreSetThroughResponseAsync() PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } - - private async Task ConstainsSplitIsTrueInternal(HttpStatusCode statusCode, SubStatusCodes subStatusCode) - { - List results = new List(); - ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; - - ItemBatchOperation operation = new ItemBatchOperation(OperationType.AddComputeGatewayRequestCharges, 0, "0"); - - results.Add( - new BatchOperationResult(HttpStatusCode.OK) - { - ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), - ETag = operation.Id - }); - - arrayOperations[0] = operation; - - MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); - - SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( - partitionKey: null, - operations: new ArraySegment(arrayOperations), - maxBodyLength: 100, - maxOperationCount: 1, - serializer: new CosmosJsonDotNetSerializer(), - cancellationToken: default(CancellationToken)); - - ResponseMessage response = new ResponseMessage(statusCode) { Content = responseContent }; - response.Headers.SubStatusCode = subStatusCode; - - BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( - response, - batchRequest, - new CosmosJsonDotNetSerializer()); - - PartitionKeyBatchResponse pkResponse = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); - - return pkResponse.ContainsSplit(); - } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs new file mode 100644 index 0000000000..b7c08241f8 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs @@ -0,0 +1,109 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Collections.Generic; + using System.IO; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class PartitionKeyRangeBatchExecutionResultTests + { + [TestMethod] + public async Task ConstainsSplitIsTrue() + { + Assert.IsTrue(await this.ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.CompletingSplit)); + Assert.IsTrue(await this.ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.CompletingPartitionMigration)); + Assert.IsTrue(await this.ConstainsSplitIsTrueInternal(HttpStatusCode.Gone, SubStatusCodes.PartitionKeyRangeGone)); + } + + [TestMethod] + public async Task ConstainsSplitIsFalse() + { + Assert.IsFalse(await this.ConstainsSplitIsTrueInternal(HttpStatusCode.OK, SubStatusCodes.Unknown)); + Assert.IsFalse(await this.ConstainsSplitIsTrueInternal((HttpStatusCode)429, SubStatusCodes.Unknown)); + } + + [TestMethod] + public async Task StatusCodesAreSetThroughResponseAsync() + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; + + ItemBatchOperation operation = new ItemBatchOperation(OperationType.AddComputeGatewayRequestCharges, 0, "0"); + + results.Add( + new BatchOperationResult(HttpStatusCode.OK) + { + ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), + ETag = operation.Id + }); + + arrayOperations[0] = operation; + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: 100, + maxOperationCount: 1, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: default(CancellationToken)); + + BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( + new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }, + batchRequest, + new CosmosJsonDotNetSerializer()); + + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); + } + + private async Task ConstainsSplitIsTrueInternal(HttpStatusCode statusCode, SubStatusCodes subStatusCode) + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; + + ItemBatchOperation operation = new ItemBatchOperation(OperationType.AddComputeGatewayRequestCharges, 0, "0"); + + results.Add( + new BatchOperationResult(HttpStatusCode.OK) + { + ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), + ETag = operation.Id + }); + + arrayOperations[0] = operation; + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: 100, + maxOperationCount: 1, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: default(CancellationToken)); + + ResponseMessage response = new ResponseMessage(statusCode) { Content = responseContent }; + response.Headers.SubStatusCode = subStatusCode; + + BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( + response, + batchRequest, + new CosmosJsonDotNetSerializer()); + + PartitionKeyRangeBatchExecutionResult result = new PartitionKeyRangeBatchExecutionResult("0", arrayOperations, new List { batchresponse }); + + return result.ContainsSplit(); + } + } +} From f9519400f8620d5d788514fcf2ed14725d1647aa Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Wed, 7 Aug 2019 11:22:54 -0700 Subject: [PATCH 24/41] volatile --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs | 2 +- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 9103932d87..cf973693b6 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -22,7 +22,7 @@ internal class BatchAsyncBatcher : IDisposable private readonly int maxBatchByteSize; private readonly int maxBatchOperationCount; private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); - private long currentSize = 0; + private volatile long currentSize = 0; public bool IsEmpty => this.batchOperations.Count == 0; diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 418e5a5127..3e829565af 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -28,7 +28,7 @@ internal class BatchAsyncStreamer private readonly int dispatchTimerInSeconds; private readonly CosmosSerializer cosmosSerializer; private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); - private BatchAsyncBatcher currentBatcher; + private volatile BatchAsyncBatcher currentBatcher; private TimerPool timerPool; private PooledTimer currentTimer; private Task timerTask; From 681b64d11693c529d6ca6f77cc382a7f1720d706 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Wed, 7 Aug 2019 11:31:54 -0700 Subject: [PATCH 25/41] removing --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index cf973693b6..9103932d87 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -22,7 +22,7 @@ internal class BatchAsyncBatcher : IDisposable private readonly int maxBatchByteSize; private readonly int maxBatchOperationCount; private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); - private volatile long currentSize = 0; + private long currentSize = 0; public bool IsEmpty => this.batchOperations.Count == 0; From 4dd74b9265f6cd1afef0f85c3dfac11cdc23a4fd Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Wed, 7 Aug 2019 12:33:22 -0700 Subject: [PATCH 26/41] Semaphore wrapper --- .../src/Batch/BatchAsyncBatcher.cs | 7 +---- .../src/Batch/BatchAsyncContainerExecutor.cs | 8 +---- .../src/Batch/BatchAsyncStreamer.cs | 7 +---- Microsoft.Azure.Cosmos/src/Util/Extensions.cs | 9 ++++++ .../src/Util/UsableSemaphoreWrapper.cs | 30 +++++++++++++++++++ .../UsableSemaphoreWrapperTests.cs | 29 ++++++++++++++++++ 6 files changed, 71 insertions(+), 19 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/UsableSemaphoreWrapperTests.cs diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 9103932d87..68ca1a97a5 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -67,8 +67,7 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati throw new ArgumentNullException(nameof(batchAsyncOperation)); } - await this.tryAddLimiter.WaitAsync(this.cancellationTokenSource.Token); - try + using (await this.tryAddLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) { if (this.batchOperations.Count == this.maxBatchOperationCount) { @@ -89,10 +88,6 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati this.batchOperations.Add(batchAsyncOperation); return true; } - finally - { - this.tryAddLimiter.Release(); - } } public async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index d5e864c504..4fa64ba38f 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -310,16 +310,10 @@ private async Task StartExecutionForPKRan CancellationToken cancellationToken) { SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(pkRangeId); - await limiter.WaitAsync(cancellationToken).ConfigureAwait(false); - - try + using (await limiter.UsingWaitAsync(cancellationToken)) { return await this.ExecuteServerRequestsAsync(pkRangeId, operations.Select(o => o.Operation), cancellationToken); } - finally - { - limiter.Release(); - } } private async Task CreateServerRequestAsync( diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 3e829565af..7eaee1b929 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -140,17 +140,12 @@ private async Task GetBatchToDispatchAndCreateAsync() return null; } - await this.dispatchLimiter.WaitAsync(this.cancellationTokenSource.Token); - try + using (await this.dispatchLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) { BatchAsyncBatcher previousBatcher = this.currentBatcher; this.currentBatcher = this.CreateBatchAsyncBatcher(); return previousBatcher; } - finally - { - this.dispatchLimiter.Release(); - } } private BatchAsyncBatcher CreateBatchAsyncBatcher() diff --git a/Microsoft.Azure.Cosmos/src/Util/Extensions.cs b/Microsoft.Azure.Cosmos/src/Util/Extensions.cs index 1056c9869b..f9ed13fa65 100644 --- a/Microsoft.Azure.Cosmos/src/Util/Extensions.cs +++ b/Microsoft.Azure.Cosmos/src/Util/Extensions.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos using System.Net.Http; using System.Net.Sockets; using System.Text; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Core.Trace; using Microsoft.Azure.Cosmos.Internal; @@ -119,6 +120,14 @@ internal static void TraceException(Exception e) } } + public static async Task UsingWaitAsync( + this SemaphoreSlim semaphoreSlim, + CancellationToken cancellationToken = default(CancellationToken)) + { + await semaphoreSlim.WaitAsync(cancellationToken).ConfigureAwait(false); + return new UsableSemaphoreWrapper(semaphoreSlim); + } + private static void TraceExceptionInternal(Exception e) { while (e != null) diff --git a/Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs b/Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs new file mode 100644 index 0000000000..4fb9bef16f --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs @@ -0,0 +1,30 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Threading; + + internal class UsableSemaphoreWrapper : IDisposable + { + private readonly SemaphoreSlim semaphore; + private bool diposed; + public UsableSemaphoreWrapper(SemaphoreSlim semaphore) + { + this.semaphore = semaphore; + } + + public void Dispose() + { + if (this.diposed) + { + return; + } + + this.semaphore.Release(); + this.diposed = true; + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/UsableSemaphoreWrapperTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/UsableSemaphoreWrapperTests.cs new file mode 100644 index 0000000000..40be9cd5ea --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/UsableSemaphoreWrapperTests.cs @@ -0,0 +1,29 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ +namespace Microsoft.Azure.Cosmos.Tests +{ + using System.Threading; + using System.Threading.Tasks; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class UsableSemaphoreWrapperTests + { + [TestMethod] + public async Task NotDisposedAfterUsing() + { + SemaphoreSlim semaphore = new SemaphoreSlim(1, 1); + using(await semaphore.UsingWaitAsync()) + { + ; + } + + // Normal flow + await semaphore.WaitAsync(); + semaphore.Release(); + + semaphore.Dispose(); + } + } +} From 8d09157b5b24d151493dfbefe80fdd0c8699bf57 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Wed, 7 Aug 2019 12:34:38 -0700 Subject: [PATCH 27/41] Spelling --- Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs b/Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs index 4fb9bef16f..bc7ee5a1a8 100644 --- a/Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs +++ b/Microsoft.Azure.Cosmos/src/Util/UsableSemaphoreWrapper.cs @@ -10,7 +10,7 @@ namespace Microsoft.Azure.Cosmos internal class UsableSemaphoreWrapper : IDisposable { private readonly SemaphoreSlim semaphore; - private bool diposed; + private bool disposed; public UsableSemaphoreWrapper(SemaphoreSlim semaphore) { this.semaphore = semaphore; @@ -18,13 +18,13 @@ public UsableSemaphoreWrapper(SemaphoreSlim semaphore) public void Dispose() { - if (this.diposed) + if (this.disposed) { return; } this.semaphore.Release(); - this.diposed = true; + this.disposed = true; } } } From ec1ba7c74d4d08c06e1d5ce37dce1f866babaa74 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Wed, 7 Aug 2019 16:45:20 -0700 Subject: [PATCH 28/41] Refactoring to requeue failed operations --- .../src/Batch/BatchAsyncBatcher.cs | 10 +- .../src/Batch/BatchAsyncContainerExecutor.cs | 192 ++++++------------ .../Batch/BatchAsyncBatcherTests.cs | 79 +++++-- 3 files changed, 134 insertions(+), 147 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 68ca1a97a5..66d3b83ccf 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -111,8 +111,14 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati for (int index = 0; index < this.batchOperations.Count; index++) { BatchAsyncOperationContext operation = this.batchOperations[index]; - BatchOperationResult response = batchResponse[index]; - operation.Complete(response); + try + { + BatchOperationResult response = batchResponse[operation.Operation.OperationIndex]; + operation.Complete(response); + } + catch (KeyNotFoundException) + { + } } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 4fa64ba38f..b149113a29 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -90,8 +90,23 @@ public async Task AddAsync( return await context.Task; } + public async Task DisposeAsync() + { + foreach (KeyValuePair streamer in this.streamersByPartitionKeyRange) + { + await streamer.Value.DisposeAsync(); + } + + foreach (KeyValuePair limiter in this.limitersByPartitionkeyRange) + { + limiter.Value.Dispose(); + } + + this.timerPool.Dispose(); + } + internal async Task ValidateOperationAsync( - ItemBatchOperation operation, + ItemBatchOperation operation, ItemRequestOptions itemRequestOptions = null, CancellationToken cancellationToken = default(CancellationToken)) { @@ -118,23 +133,8 @@ internal async Task ValidateOperationAsync( } } - public async Task DisposeAsync() - { - foreach (KeyValuePair streamer in this.streamersByPartitionKeyRange) - { - await streamer.Value.DisposeAsync(); - } - - foreach (KeyValuePair limiter in this.limitersByPartitionkeyRange) - { - limiter.Value.Dispose(); - } - - this.timerPool.Dispose(); - } - private static bool ValidateOperationEPK( - ItemBatchOperation operation, + ItemBatchOperation operation, ItemRequestOptions itemRequestOptions) { if (itemRequestOptions.Properties != null @@ -170,15 +170,15 @@ private static void AddHeadersToRequestMessage(RequestMessage requestMessage, st /// /// If because of HybridRow serialization overhead, not all operations fit in the request, we send those extra operations in a separate request. /// - private static IEnumerable GetOverflowOperations( + private static IEnumerable GetOverflowOperations( PartitionKeyServerBatchRequest request, - IEnumerable operationsSentToRequest) + IEnumerable operationsSentToRequest) { int totalOperations = operationsSentToRequest.Count(); int operationsThatOverflowed = totalOperations - request.Operations.Count; if (operationsThatOverflowed == 0) { - return Enumerable.Empty(); + return Enumerable.Empty(); } return operationsSentToRequest.Skip(totalOperations - operationsThatOverflowed); @@ -197,84 +197,46 @@ private static IReadOnlyList GetOperationsToRetry( return operations; } - private async Task ExecuteAsync( - IReadOnlyList operations, + private async Task ReBatchAsync( + BatchAsyncOperationContext context, CancellationToken cancellationToken) { - IEnumerable responses = await this.ExecuteInternalAsync(operations, false, cancellationToken); - return new PartitionKeyBatchResponse(responses, this.cosmosClientContext.CosmosSerializer); + string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(context.Operation, cancellationToken).ConfigureAwait(false); + BatchAsyncStreamer streamer = await this.GetOrAddStreamerForPartitionKeyRangeAsync(resolvedPartitionKeyRangeId).ConfigureAwait(false); + await streamer.AddAsync(context).ConfigureAwait(false); } - private async Task> ExecuteInternalAsync( + private async Task ExecuteAsync( IReadOnlyList operations, - bool refreshPartitionKeyRanges, CancellationToken cancellationToken) { - List> inProgressTasksByPKRangeId = new List>(); - List completedResults = new List(); - - // Initially, all operations should be for the same PKRange - if (!refreshPartitionKeyRanges) - { - inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(operations[0].PartitionKeyRangeId, operations, cancellationToken)); - } - else + // All operations should be for the same PKRange + string partitionKeyRangeId = operations[0].PartitionKeyRangeId; + PartitionKeyServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operations.Select(o => o.Operation), cancellationToken); + // Any overflow goes to a new batch + IEnumerable overFlowOperations = BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operations); + foreach (BatchAsyncOperationContext overflowedContext in overFlowOperations) { - // If a refresh was requested, recalculate PKRange - Dictionary> operationsByPKRangeId = await this.GroupByPartitionKeyRangeIdsAsync(operations, cancellationToken); - foreach (KeyValuePair> operationsForPKRangeId in operationsByPKRangeId) - { - inProgressTasksByPKRangeId.Add(this.StartExecutionForPKRangeIdAsync(operationsForPKRangeId.Key, operationsForPKRangeId.Value, cancellationToken)); - } + await this.ReBatchAsync(overflowedContext, cancellationToken); } - while (inProgressTasksByPKRangeId.Count > 0) - { - Task completedTask = await Task.WhenAny(inProgressTasksByPKRangeId); - completedResults.Add(await completedTask); - inProgressTasksByPKRangeId.Remove(completedTask); - } + PartitionKeyRangeBatchExecutionResult result = await this.ExecuteServerRequestAsync(serverRequest, cancellationToken); - List responses = new List(operations.Count); - foreach (PartitionKeyRangeBatchExecutionResult result in completedResults) + List responses = new List(serverRequest.Operations.Count); + if (!result.ContainsSplit()) { - if (!result.ContainsSplit()) - { - responses.AddRange(result.ServerResponses); - } - else - { - IReadOnlyList retryOperations = BatchAsyncContainerExecutor.GetOperationsToRetry(operations, result.Operations.Select(o => o.OperationIndex)); - responses.AddRange(await this.ExecuteInternalAsync(retryOperations, true, cancellationToken)); - } + responses.AddRange(result.ServerResponses); } - - return responses; - } - - private async Task>> GroupByPartitionKeyRangeIdsAsync( - IEnumerable contexts, - CancellationToken cancellationToken) - { - PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken).ConfigureAwait(false); - CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken).ConfigureAwait(false); - - Dictionary> operationsByPKRangeId = new Dictionary>(); - foreach (BatchAsyncOperationContext context in contexts) + else { - string partitionKeyRangeId = BatchExecUtils.GetPartitionKeyRangeId(context.Operation.PartitionKey.Value, partitionKeyDefinition, collectionRoutingMap); - - if (operationsByPKRangeId.TryGetValue(partitionKeyRangeId, out List operationsForPKRangeId)) - { - operationsForPKRangeId.Add(context); - } - else + IReadOnlyList retryContexts = BatchAsyncContainerExecutor.GetOperationsToRetry(operations, result.Operations.Select(o => o.OperationIndex)); + foreach (BatchAsyncOperationContext retryContext in retryContexts) { - operationsByPKRangeId.Add(partitionKeyRangeId, new List() { context }); + await this.ReBatchAsync(retryContext, cancellationToken); } } - return operationsByPKRangeId; + return new PartitionKeyBatchResponse(responses, this.cosmosClientContext.CosmosSerializer); } private async Task ResolvePartitionKeyRangeIdAsync( @@ -304,18 +266,6 @@ private async Task FillOperationPropertiesAsync(ItemBatchOperation operation, Ca } } - private async Task StartExecutionForPKRangeIdAsync( - string pkRangeId, - IEnumerable operations, - CancellationToken cancellationToken) - { - SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(pkRangeId); - using (await limiter.UsingWaitAsync(cancellationToken)) - { - return await this.ExecuteServerRequestsAsync(pkRangeId, operations.Select(o => o.Operation), cancellationToken); - } - } - private async Task CreateServerRequestAsync( string partitionKeyRangeId, IEnumerable operations, @@ -335,44 +285,32 @@ private async Task CreateServerRequestAsync( return request; } - private async Task ExecuteServerRequestsAsync( - string partitionKeyRangeId, - IEnumerable operations, + private async Task ExecuteServerRequestAsync( + PartitionKeyServerBatchRequest serverRequest, CancellationToken cancellationToken) { - List serverResponses = new List(); - PartitionKeyServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operations, cancellationToken); - - // In case some operations overflowed due to HybridRow serialization - IEnumerable overFlowOperations = BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operations); - serverResponses.Add(await this.ExecuteServerRequestAsync(serverRequest, cancellationToken)); - if (overFlowOperations.Any()) - { - serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, overFlowOperations, cancellationToken); - serverResponses.Add(await this.ExecuteServerRequestAsync(serverRequest, cancellationToken)); - } - - return new PartitionKeyRangeBatchExecutionResult(partitionKeyRangeId, operations, serverResponses); - } - - private async Task ExecuteServerRequestAsync(PartitionKeyServerBatchRequest serverRequest, CancellationToken cancellationToken) - { - using (Stream serverRequestPayload = serverRequest.TransferBodyStream()) + SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(serverRequest.PartitionKeyRangeId); + using (await limiter.UsingWaitAsync(cancellationToken)) { - Debug.Assert(serverRequestPayload != null, "Server request payload expected to be non-null"); - - ResponseMessage responseMessage = await this.cosmosClientContext.ProcessResourceOperationStreamAsync( - this.cosmosContainer.LinkUri, - ResourceType.Document, - OperationType.Batch, - new RequestOptions(), - cosmosContainerCore: this.cosmosContainer, - partitionKey: null, - streamPayload: serverRequestPayload, - requestEnricher: requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest.PartitionKeyRangeId), - cancellationToken: cancellationToken).ConfigureAwait(false); - - return await BatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.CosmosSerializer).ConfigureAwait(false); + using (Stream serverRequestPayload = serverRequest.TransferBodyStream()) + { + Debug.Assert(serverRequestPayload != null, "Server request payload expected to be non-null"); + + ResponseMessage responseMessage = await this.cosmosClientContext.ProcessResourceOperationStreamAsync( + this.cosmosContainer.LinkUri, + ResourceType.Document, + OperationType.Batch, + new RequestOptions(), + cosmosContainerCore: this.cosmosContainer, + partitionKey: null, + streamPayload: serverRequestPayload, + requestEnricher: requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest.PartitionKeyRangeId), + cancellationToken: cancellationToken).ConfigureAwait(false); + + BatchResponse serverResponse = await BatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.CosmosSerializer).ConfigureAwait(false); + + return new PartitionKeyRangeBatchExecutionResult(serverRequest.PartitionKeyRangeId, serverRequest.Operations, new List() { serverResponse }); + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 03a0e25a23..db8c142a6f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Tests using System; using System.Collections.Generic; using System.IO; + using System.Linq; using System.Net; using System.Threading; using System.Threading.Tasks; @@ -58,25 +59,21 @@ private Func, CancellationToken, Task< return response; }; - private Func, CancellationToken, Task> ExecutorWithFailure - = (IReadOnlyList operations, CancellationToken cancellation) => - { - throw expectedException; - }; - - private Func, CancellationToken, Task> ExecutorWithSplit + // The response will include all but 2 operation responses + private Func, CancellationToken, Task> ExecutorWithLessResponses = async (IReadOnlyList operations, CancellationToken cancellation) => { + int operationCount = operations.Count - 2; List results = new List(); - ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operations.Count]; + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operationCount]; int index = 0; - foreach (BatchAsyncOperationContext operation in operations) + foreach (BatchAsyncOperationContext operation in operations.Take(operationCount)) { results.Add( - new BatchOperationResult(HttpStatusCode.Gone) + new BatchOperationResult(HttpStatusCode.OK) { - ETag = operation.Operation.Id, - SubStatusCode = SubStatusCodes.PartitionKeyRangeGone + ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), + ETag = operation.Operation.Id }); arrayOperations[index++] = operation.Operation; @@ -87,16 +84,13 @@ private Func, CancellationToken, Task< SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( partitionKey: null, operations: new ArraySegment(arrayOperations), - maxBodyLength: (int)responseContent.Length * operations.Count, - maxOperationCount: operations.Count, + maxBodyLength: (int)responseContent.Length * operationCount, + maxOperationCount: operationCount, serializer: new CosmosJsonDotNetSerializer(), cancellationToken: cancellation); - ResponseMessage responseMessage = new ResponseMessage(HttpStatusCode.Gone) { Content = responseContent }; - responseMessage.Headers.SubStatusCode = SubStatusCodes.PartitionKeyRangeGone; - BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( - responseMessage, + new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }, batchRequest, new CosmosJsonDotNetSerializer()); @@ -104,6 +98,12 @@ private Func, CancellationToken, Task< return response; }; + private Func, CancellationToken, Task> ExecutorWithFailure + = (IReadOnlyList operations, CancellationToken cancellation) => + { + throw expectedException; + }; + [DataTestMethod] [ExpectedException(typeof(ArgumentOutOfRangeException))] [DataRow(0)] @@ -212,6 +212,49 @@ public async Task DispatchProcessInOrderAsync() } } + [TestMethod] + public async Task DispatchWithLessResponses() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithLessResponses); + BatchAsyncBatcher secondAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + List contexts = new List(10); + for (int i = 0; i < 10; i++) + { + BatchAsyncOperationContext context = new BatchAsyncOperationContext(string.Empty, new ItemBatchOperation(OperationType.Create, i, i.ToString())); + contexts.Add(context); + Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(context)); + } + + await batchAsyncBatcher.DispatchAsync(); + + for (int i = 0; i < 10; i++) + { + BatchAsyncOperationContext context = contexts[i]; + // Some tasks should not be resolved + Assert.IsTrue(context.Task.Status == TaskStatus.RanToCompletion || context.Task.Status == TaskStatus.WaitingForActivation); + if (context.Task.Status == TaskStatus.RanToCompletion) + { + BatchOperationResult result = await context.Task; + Assert.AreEqual(i.ToString(), result.ETag); + } + else + { + // Pass the pending one to another batcher + Assert.IsTrue(await secondAsyncBatcher.TryAddAsync(context)); + } + } + + await secondAsyncBatcher.DispatchAsync(); + // All tasks should be completed + for (int i = 0; i < 10; i++) + { + BatchAsyncOperationContext context = contexts[i]; + Assert.AreEqual(TaskStatus.RanToCompletion, context.Task.Status); + BatchOperationResult result = await context.Task; + Assert.AreEqual(i.ToString(), result.ETag); + } + } + [TestMethod] public void IsEmptyWithNoOperations() { From cdf0ec608f5edc56afdd007a30c177fa68e65df9 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 8 Aug 2019 15:03:34 -0700 Subject: [PATCH 29/41] Refactoring --- .../src/Batch/BatchAsyncBatcher.cs | 26 +++++-------------- .../src/Batch/BatchAsyncContainerExecutor.cs | 8 +++--- .../src/Batch/BatchAsyncOperationContext.cs | 14 +++++++--- .../src/Batch/PartitionKeyBatchResponse.cs | 18 ++++++++----- .../Batch/BatchAsyncBatcherTests.cs | 18 ++++++++----- .../Batch/BatchAsyncOperationContextTests.cs | 4 +-- .../Batch/BatchAsyncStreamerTests.cs | 2 +- .../Batch/PartitionKeyBatchResponseTests.cs | 2 +- ...titionKeyRangeBatchExecutionResultTests.cs | 2 +- 9 files changed, 50 insertions(+), 44 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 66d3b83ccf..1943225fe3 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos using System.Collections.Generic; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Core.Trace; /// /// Maintains a batch of operations and dispatches the batch through an executor. Maps results into the original operation contexts. @@ -85,6 +86,7 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati // Operation index is in the scope of the current batch batchAsyncOperation.Operation.OperationIndex = this.batchOperations.Count; + batchAsyncOperation.CurrentBatcher = this; this.batchOperations.Add(batchAsyncOperation); return true; } @@ -96,38 +98,24 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati { using (PartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken)) { - // If the batch was not successful, we need to set all the responses - if (!batchResponse.IsSuccessStatusCode) - { - BatchOperationResult errorResult = new BatchOperationResult(batchResponse.StatusCode); - foreach (BatchAsyncOperationContext operation in this.batchOperations) - { - operation.Complete(errorResult); - } - - return; - } - for (int index = 0; index < this.batchOperations.Count; index++) { BatchAsyncOperationContext operation = this.batchOperations[index]; - try - { - BatchOperationResult response = batchResponse[operation.Operation.OperationIndex]; - operation.Complete(response); - } - catch (KeyNotFoundException) + BatchOperationResult response = batchResponse[operation.Operation.OperationIndex]; + if (response != null) { + operation.Complete(this, response); } } } } catch (Exception ex) { + DefaultTrace.TraceError("Exception during BatchAsyncBatcher: {0}", ex); // Exceptions happening during execution fail all the Tasks foreach (BatchAsyncOperationContext operation in this.batchOperations) { - operation.Fail(ex); + operation.Fail(this, ex); } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index b149113a29..7d44b83731 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -86,7 +86,7 @@ public async Task AddAsync( string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false); BatchAsyncStreamer streamer = await this.GetOrAddStreamerForPartitionKeyRangeAsync(resolvedPartitionKeyRangeId).ConfigureAwait(false); BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); - await streamer.AddAsync(context).ConfigureAwait(false); + await streamer.AddAsync(context); return await context.Task; } @@ -123,7 +123,7 @@ internal async Task ValidateOperationAsync( Debug.Assert(BatchAsyncContainerExecutor.ValidateOperationEPK(operation, itemRequestOptions)); } - await operation.MaterializeResourceAsync(this.cosmosClientContext.CosmosSerializer, cancellationToken).ConfigureAwait(false); + await operation.MaterializeResourceAsync(this.cosmosClientContext.CosmosSerializer, cancellationToken); int itemByteSize = operation.GetApproximateSerializedLength(); @@ -203,7 +203,7 @@ private async Task ReBatchAsync( { string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(context.Operation, cancellationToken).ConfigureAwait(false); BatchAsyncStreamer streamer = await this.GetOrAddStreamerForPartitionKeyRangeAsync(resolvedPartitionKeyRangeId).ConfigureAwait(false); - await streamer.AddAsync(context).ConfigureAwait(false); + await streamer.AddAsync(context); } private async Task ExecuteAsync( @@ -236,7 +236,7 @@ private async Task ExecuteAsync( } } - return new PartitionKeyBatchResponse(responses, this.cosmosClientContext.CosmosSerializer); + return new PartitionKeyBatchResponse(operations.Count, responses, this.cosmosClientContext.CosmosSerializer); } private async Task ResolvePartitionKeyRangeIdAsync( diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs index db147e5538..2200ab5e63 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs @@ -5,7 +5,7 @@ namespace Microsoft.Azure.Cosmos { using System; - using System.Threading; + using System.Diagnostics; using System.Threading.Tasks; /// @@ -17,6 +17,8 @@ internal class BatchAsyncOperationContext public ItemBatchOperation Operation { get; } + public BatchAsyncBatcher CurrentBatcher { get; set; } + public Task Task => this.taskCompletionSource.Task; private TaskCompletionSource taskCompletionSource = new TaskCompletionSource(); @@ -29,13 +31,19 @@ public BatchAsyncOperationContext( this.PartitionKeyRangeId = partitionKeyRangeId; } - public void Complete(BatchOperationResult result) + public void Complete( + BatchAsyncBatcher completer, + BatchOperationResult result) { + Debug.Assert(this.CurrentBatcher == null || completer == this.CurrentBatcher); this.taskCompletionSource.SetResult(result); } - public void Fail(Exception exception) + public void Fail( + BatchAsyncBatcher completer, + Exception exception) { + Debug.Assert(this.CurrentBatcher == null || completer == this.CurrentBatcher); this.taskCompletionSource.SetException(exception); } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs index b95fba3fbd..2a1b328dad 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs @@ -17,7 +17,7 @@ namespace Microsoft.Azure.Cosmos internal class PartitionKeyBatchResponse : BatchResponse { // Results sorted in the order operations had been added. - private readonly SortedList resultsByOperationIndex; + private readonly BatchOperationResult[] resultsByOperationIndex; private bool isDisposed; /// @@ -40,16 +40,20 @@ internal PartitionKeyBatchResponse( /// /// Initializes a new instance of the class. /// + /// Original operations that generated the server responses. /// Responses from the server. /// Serializer to deserialize response resource body streams. - internal PartitionKeyBatchResponse(IEnumerable serverResponses, CosmosSerializer serializer) + internal PartitionKeyBatchResponse( + int originalOperationsCount, + IEnumerable serverResponses, + CosmosSerializer serializer) { this.StatusCode = serverResponses.Any(r => r.StatusCode != HttpStatusCode.OK) ? (HttpStatusCode)StatusCodes.MultiStatus : HttpStatusCode.OK; this.ServerResponses = serverResponses; - this.resultsByOperationIndex = new SortedList(); + this.resultsByOperationIndex = new BatchOperationResult[originalOperationsCount]; StringBuilder errorMessageBuilder = new StringBuilder(); List activityIds = new List(); @@ -60,7 +64,7 @@ internal PartitionKeyBatchResponse(IEnumerable serverResponses, C for (int index = 0; index < serverResponse.Operations.Count; index++) { int operationIndex = serverResponse.Operations[index].OperationIndex; - if (!this.resultsByOperationIndex.ContainsKey(operationIndex) + if (this.resultsByOperationIndex[operationIndex] == null || this.resultsByOperationIndex[operationIndex].StatusCode == (HttpStatusCode)StatusCodes.TooManyRequests) { this.resultsByOperationIndex[operationIndex] = serverResponse[index]; @@ -102,7 +106,7 @@ internal PartitionKeyBatchResponse(IEnumerable serverResponses, C /// /// Gets the number of operation results. /// - public override int Count => this.resultsByOperationIndex.Count; + public override int Count => this.resultsByOperationIndex.Length; /// public override BatchOperationResult this[int index] => this.resultsByOperationIndex[index]; @@ -137,9 +141,9 @@ public override BatchOperationResult GetOperationResultAtIndex(int index) /// Enumerator over the operation results. public override IEnumerator GetEnumerator() { - foreach (KeyValuePair pair in this.resultsByOperationIndex) + foreach (BatchOperationResult result in this.resultsByOperationIndex) { - yield return pair.Value; + yield return result; } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index db8c142a6f..d87cdf3bdd 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -55,8 +55,7 @@ private Func, CancellationToken, Task< batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); - return response; + return new PartitionKeyBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); }; // The response will include all but 2 operation responses @@ -67,7 +66,7 @@ private Func, CancellationToken, Task< List results = new List(); ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operationCount]; int index = 0; - foreach (BatchAsyncOperationContext operation in operations.Take(operationCount)) + foreach (BatchAsyncOperationContext operation in operations.Skip(1).Take(operationCount)) { results.Add( new BatchOperationResult(HttpStatusCode.OK) @@ -94,8 +93,7 @@ private Func, CancellationToken, Task< batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); - return response; + return new PartitionKeyBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); }; private Func, CancellationToken, Task> ExecutorWithFailure @@ -227,11 +225,19 @@ public async Task DispatchWithLessResponses() await batchAsyncBatcher.DispatchAsync(); + // Responses 1 and 10 should be missing for (int i = 0; i < 10; i++) { BatchAsyncOperationContext context = contexts[i]; // Some tasks should not be resolved - Assert.IsTrue(context.Task.Status == TaskStatus.RanToCompletion || context.Task.Status == TaskStatus.WaitingForActivation); + if(i == 0 || i == 9) + { + Assert.IsTrue(context.Task.Status == TaskStatus.WaitingForActivation); + } + else + { + Assert.IsTrue(context.Task.Status == TaskStatus.RanToCompletion); + } if (context.Task.Status == TaskStatus.RanToCompletion) { BatchOperationResult result = await context.Task; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs index 8017fa3587..867c3d0917 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs @@ -45,7 +45,7 @@ public async Task TaskResultIsSetOnCompleteAsync() BatchOperationResult expected = new BatchOperationResult(HttpStatusCode.OK); - batchAsyncOperationContext.Complete(expected); + batchAsyncOperationContext.Complete(null, expected); Assert.AreEqual(expected, await batchAsyncOperationContext.Task); Assert.AreEqual(TaskStatus.RanToCompletion, batchAsyncOperationContext.Task.Status); @@ -58,7 +58,7 @@ public async Task ExceptionIsSetOnFailAsync() ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(string.Empty, operation); - batchAsyncOperationContext.Fail(failure); + batchAsyncOperationContext.Fail(null, failure); Exception capturedException = await Assert.ThrowsExceptionAsync(() => batchAsyncOperationContext.Task); Assert.AreEqual(failure, capturedException); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index f94cd222ae..0f7002d354 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -56,7 +56,7 @@ private Func, CancellationToken, Task< batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); return response; }; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs index f3b39ac891..a10a70baf4 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs @@ -58,7 +58,7 @@ public async Task StatusCodesAreSetThroughResponseAsync() batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(arrayOperations.Length, new List { batchresponse }, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs index b7c08241f8..38b423589e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs @@ -63,7 +63,7 @@ public async Task StatusCodesAreSetThroughResponseAsync() batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(arrayOperations.Length, new List { batchresponse }, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } From 7335a82d34835f0d8f9dfb156719da0c8ae28ad1 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 8 Aug 2019 20:32:57 -0700 Subject: [PATCH 30/41] Reducing semaphores and fixing bug --- .../src/Batch/BatchAsyncBatcher.cs | 55 +++++++++---------- .../src/Batch/BatchAsyncContainerExecutor.cs | 3 +- .../src/Batch/BatchAsyncStreamer.cs | 37 ++++++------- .../Batch/BatchAsyncBatcherTests.cs | 53 ++++++------------ 4 files changed, 63 insertions(+), 85 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 1943225fe3..43c8a8e178 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -14,16 +14,15 @@ namespace Microsoft.Azure.Cosmos /// Maintains a batch of operations and dispatches the batch through an executor. Maps results into the original operation contexts. /// /// - internal class BatchAsyncBatcher : IDisposable + internal class BatchAsyncBatcher { - private readonly SemaphoreSlim tryAddLimiter; private readonly CosmosSerializer CosmosSerializer; private readonly List batchOperations; private readonly Func, CancellationToken, Task> executor; private readonly int maxBatchByteSize; private readonly int maxBatchOperationCount; - private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); private long currentSize = 0; + private bool dispached = false; public bool IsEmpty => this.batchOperations.Count == 0; @@ -54,42 +53,43 @@ public BatchAsyncBatcher( } this.batchOperations = new List(maxBatchOperationCount); - this.tryAddLimiter = new SemaphoreSlim(1, 1); this.executor = executor; this.maxBatchByteSize = maxBatchByteSize; this.maxBatchOperationCount = maxBatchOperationCount; this.CosmosSerializer = cosmosSerializer; } - public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperation) + public bool TryAdd(BatchAsyncOperationContext batchAsyncOperation) { + if (this.dispached) + { + return false; + } + if (batchAsyncOperation == null) { throw new ArgumentNullException(nameof(batchAsyncOperation)); } - using (await this.tryAddLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) + if (this.batchOperations.Count == this.maxBatchOperationCount) { - if (this.batchOperations.Count == this.maxBatchOperationCount) - { - return false; - } + return false; + } - int itemByteSize = batchAsyncOperation.Operation.GetApproximateSerializedLength(); + int itemByteSize = batchAsyncOperation.Operation.GetApproximateSerializedLength(); - if (itemByteSize + this.currentSize > this.maxBatchByteSize) - { - return false; - } + if (itemByteSize + this.currentSize > this.maxBatchByteSize) + { + return false; + } - this.currentSize += itemByteSize; + this.currentSize += itemByteSize; - // Operation index is in the scope of the current batch - batchAsyncOperation.Operation.OperationIndex = this.batchOperations.Count; - batchAsyncOperation.CurrentBatcher = this; - this.batchOperations.Add(batchAsyncOperation); - return true; - } + // Operation index is in the scope of the current batch + batchAsyncOperation.Operation.OperationIndex = this.batchOperations.Count; + batchAsyncOperation.CurrentBatcher = this; + this.batchOperations.Add(batchAsyncOperation); + return true; } public async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) @@ -118,13 +118,10 @@ public async Task TryAddAsync(BatchAsyncOperationContext batchAsyncOperati operation.Fail(this, ex); } } - } - - public void Dispose() - { - this.cancellationTokenSource.Cancel(); - this.cancellationTokenSource.Dispose(); - this.tryAddLimiter.Dispose(); + finally + { + this.dispached = true; + } } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 7d44b83731..fdc8ff9329 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -222,7 +222,7 @@ private async Task ExecuteAsync( PartitionKeyRangeBatchExecutionResult result = await this.ExecuteServerRequestAsync(serverRequest, cancellationToken); - List responses = new List(serverRequest.Operations.Count); + List responses = new List(result.ServerResponses.Count); if (!result.ContainsSplit()) { responses.AddRange(result.ServerResponses); @@ -272,7 +272,6 @@ private async Task CreateServerRequestAsync( CancellationToken cancellationToken) { ArraySegment operationsArraySegment = new ArraySegment(operations.ToArray()); - PartitionKeyServerBatchRequest request = await PartitionKeyServerBatchRequest.CreateAsync( partitionKeyRangeId, operationsArraySegment, diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 7eaee1b929..002e0c61b6 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -80,17 +80,17 @@ public BatchAsyncStreamer( public async Task AddAsync(BatchAsyncOperationContext context) { - if (!await this.currentBatcher.TryAddAsync(context)) + using (await this.dispatchLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) { - // Batcher is full - BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAndCreateAsync(); - if (toDispatch != null) + while (!this.currentBatcher.TryAdd(context)) { - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + // Batcher is full + BatchAsyncBatcher toDispatch = this.GetBatchToDispatchAndCreate(); + if (toDispatch != null) + { + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + } } - - bool addedContext = await this.currentBatcher.TryAddAsync(context); - Debug.Assert(addedContext, "Could not add context to batcher."); } } @@ -98,7 +98,6 @@ public async Task DisposeAsync() { this.cancellationTokenSource.Cancel(); this.cancellationTokenSource.Dispose(); - this.currentBatcher?.Dispose(); this.currentTimer.CancelTimer(); foreach (Task previousDispatchedTask in this.previousDispatchedTasks) { @@ -124,28 +123,28 @@ private async Task DispatchTimerAsync() return; } - BatchAsyncBatcher toDispatch = await this.GetBatchToDispatchAndCreateAsync(); - if (toDispatch != null) + using (await this.dispatchLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) { - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + BatchAsyncBatcher toDispatch = this.GetBatchToDispatchAndCreate(); + if (toDispatch != null) + { + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + } } this.StartTimer(); } - private async Task GetBatchToDispatchAndCreateAsync() + private BatchAsyncBatcher GetBatchToDispatchAndCreate() { if (this.currentBatcher.IsEmpty) { return null; } - using (await this.dispatchLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) - { - BatchAsyncBatcher previousBatcher = this.currentBatcher; - this.currentBatcher = this.CreateBatchAsyncBatcher(); - return previousBatcher; - } + BatchAsyncBatcher previousBatcher = this.currentBatcher; + this.currentBatcher = this.CreateBatchAsyncBatcher(); + return previousBatcher; } private BatchAsyncBatcher CreateBatchAsyncBatcher() diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index d87cdf3bdd..11c05be1fb 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -135,12 +135,12 @@ public void ValidatesSerializer() } [TestMethod] - public async Task HasFixedSize() + public void HasFixedSize() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsFalse(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); } [TestMethod] @@ -149,26 +149,9 @@ public async Task HasFixedByteSize() await this.ItemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), CancellationToken.None); // Each operation is 2 bytes BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsFalse(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - } - - [TestMethod] - public void TryAddIsThreadSafe() - { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Task firstOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation)); - Task secondOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation)); - Task thirdOperation = batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation)); - - Task.WhenAll(firstOperation, secondOperation, thirdOperation).GetAwaiter().GetResult(); - - int countSucceded = (firstOperation.Result ? 1 : 0) + (secondOperation.Result ? 1 : 0) + (thirdOperation.Result ? 1 : 0); - int countFailed = (!firstOperation.Result ? 1 : 0) + (!secondOperation.Result ? 1 : 0) + (!thirdOperation.Result ? 1 : 0); - - Assert.AreEqual(2, countSucceded); - Assert.AreEqual(1, countFailed); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); } [TestMethod] @@ -177,8 +160,8 @@ public async Task ExceptionsFailOperationsAsync() BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); - await batchAsyncBatcher.TryAddAsync(context1); - await batchAsyncBatcher.TryAddAsync(context2); + batchAsyncBatcher.TryAdd(context1); + batchAsyncBatcher.TryAdd(context2); await batchAsyncBatcher.DispatchAsync(); Assert.AreEqual(TaskStatus.Faulted, context1.Task.Status); @@ -196,7 +179,7 @@ public async Task DispatchProcessInOrderAsync() { BatchAsyncOperationContext context = new BatchAsyncOperationContext(string.Empty, new ItemBatchOperation(OperationType.Create, i, i.ToString())); contexts.Add(context); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(context)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(context)); } await batchAsyncBatcher.DispatchAsync(); @@ -220,7 +203,7 @@ public async Task DispatchWithLessResponses() { BatchAsyncOperationContext context = new BatchAsyncOperationContext(string.Empty, new ItemBatchOperation(OperationType.Create, i, i.ToString())); contexts.Add(context); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(context)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(context)); } await batchAsyncBatcher.DispatchAsync(); @@ -246,7 +229,7 @@ public async Task DispatchWithLessResponses() else { // Pass the pending one to another batcher - Assert.IsTrue(await secondAsyncBatcher.TryAddAsync(context)); + Assert.IsTrue(secondAsyncBatcher.TryAdd(context)); } } @@ -269,20 +252,20 @@ public void IsEmptyWithNoOperations() } [TestMethod] - public async Task IsNotEmptyWithOperations() + public void IsNotEmptyWithOperations() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); Assert.IsFalse(batchAsyncBatcher.IsEmpty); } [TestMethod] - public async Task CannotAddToDisposedBatch() + public async Task CannotAddToDispatchedBatch() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(await batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - batchAsyncBatcher.Dispose(); - await Assert.ThrowsExceptionAsync(() => batchAsyncBatcher.TryAddAsync(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + await batchAsyncBatcher.DispatchAsync(); + Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); } } } From 31466129a0be8f335169e4d9c1c2d65902463919 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Fri, 9 Aug 2019 09:11:13 -0700 Subject: [PATCH 31/41] Rename --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 002e0c61b6..1f49eed1d2 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -75,7 +75,7 @@ public BatchAsyncStreamer( this.dispatchLimiter = new SemaphoreSlim(1, 1); this.currentBatcher = this.CreateBatchAsyncBatcher(); - this.StartTimer(); + this.ResetTimer(); } public async Task AddAsync(BatchAsyncOperationContext context) @@ -107,7 +107,7 @@ public async Task DisposeAsync() this.dispatchLimiter.Dispose(); } - private void StartTimer() + private void ResetTimer() { this.currentTimer = this.timerPool.GetPooledTimer(this.dispatchTimerInSeconds); this.timerTask = this.currentTimer.StartTimerAsync().ContinueWith((task) => @@ -132,7 +132,7 @@ private async Task DispatchTimerAsync() } } - this.StartTimer(); + this.ResetTimer(); } private BatchAsyncBatcher GetBatchToDispatchAndCreate() From c5e17b3f8bbf6370f83e9f5ac052efc82c8de450 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Fri, 9 Aug 2019 12:05:04 -0700 Subject: [PATCH 32/41] Dispatch outside --- .../src/Batch/BatchAsyncStreamer.cs | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 1f49eed1d2..510a36ea3b 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -80,18 +80,20 @@ public BatchAsyncStreamer( public async Task AddAsync(BatchAsyncOperationContext context) { + BatchAsyncBatcher toDispatch = null; using (await this.dispatchLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) { while (!this.currentBatcher.TryAdd(context)) { // Batcher is full - BatchAsyncBatcher toDispatch = this.GetBatchToDispatchAndCreate(); - if (toDispatch != null) - { - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); - } + toDispatch = this.GetBatchToDispatchAndCreate(); } } + + if (toDispatch != null) + { + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + } } public async Task DisposeAsync() @@ -113,7 +115,7 @@ private void ResetTimer() this.timerTask = this.currentTimer.StartTimerAsync().ContinueWith((task) => { return this.DispatchTimerAsync(); - }); + }, this.cancellationTokenSource.Token); } private async Task DispatchTimerAsync() @@ -123,13 +125,15 @@ private async Task DispatchTimerAsync() return; } + BatchAsyncBatcher toDispatch; using (await this.dispatchLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) { - BatchAsyncBatcher toDispatch = this.GetBatchToDispatchAndCreate(); - if (toDispatch != null) - { - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); - } + toDispatch = this.GetBatchToDispatchAndCreate(); + } + + if (toDispatch != null) + { + this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); } this.ResetTimer(); From bd25f19a5e5f43e508e4950437cddf9e0be6ee9d Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 12 Aug 2019 08:15:33 -0700 Subject: [PATCH 33/41] Variable rename --- Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs | 11 ++++++----- .../src/Batch/BatchAsyncOperationContext.cs | 10 +++++++++- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 43c8a8e178..b5084f831b 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -100,11 +100,11 @@ public bool TryAdd(BatchAsyncOperationContext batchAsyncOperation) { for (int index = 0; index < this.batchOperations.Count; index++) { - BatchAsyncOperationContext operation = this.batchOperations[index]; - BatchOperationResult response = batchResponse[operation.Operation.OperationIndex]; + BatchAsyncOperationContext context = this.batchOperations[index]; + BatchOperationResult response = batchResponse[context.Operation.OperationIndex]; if (response != null) { - operation.Complete(this, response); + context.Complete(this, response); } } } @@ -113,13 +113,14 @@ public bool TryAdd(BatchAsyncOperationContext batchAsyncOperation) { DefaultTrace.TraceError("Exception during BatchAsyncBatcher: {0}", ex); // Exceptions happening during execution fail all the Tasks - foreach (BatchAsyncOperationContext operation in this.batchOperations) + foreach (BatchAsyncOperationContext context in this.batchOperations) { - operation.Fail(this, ex); + context.Fail(this, ex); } } finally { + this.batchOperations.Clear(); this.dispached = true; } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs index 2200ab5e63..19dfece54c 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs @@ -11,7 +11,7 @@ namespace Microsoft.Azure.Cosmos /// /// Context for a particular Batch operation. /// - internal class BatchAsyncOperationContext + internal class BatchAsyncOperationContext : IDisposable { public string PartitionKeyRangeId { get; } @@ -37,6 +37,7 @@ public void Complete( { Debug.Assert(this.CurrentBatcher == null || completer == this.CurrentBatcher); this.taskCompletionSource.SetResult(result); + this.Dispose(); } public void Fail( @@ -45,6 +46,13 @@ public void Fail( { Debug.Assert(this.CurrentBatcher == null || completer == this.CurrentBatcher); this.taskCompletionSource.SetException(exception); + this.Dispose(); + } + + public void Dispose() + { + this.Operation.Dispose(); + this.CurrentBatcher = null; } } } From a000f0c71075d84035c157c972fb3e3cf325bfe4 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 12 Aug 2019 09:07:12 -0700 Subject: [PATCH 34/41] Reducing memory footprint --- .../src/Batch/BatchAsyncContainerExecutor.cs | 14 +++++++------- .../src/Batch/BatchAsyncStreamer.cs | 18 ++++++++---------- .../Batch/BatchAsyncStreamerTests.cs | 4 ++-- 3 files changed, 17 insertions(+), 19 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index fdc8ff9329..ba11c1ebdc 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -22,7 +22,7 @@ namespace Microsoft.Azure.Cosmos /// It groups operations by Partition Key Range and sends them to the Batch API and then unifies results as they become available. It uses as batch processor and as batch executing handler. /// /// - internal class BatchAsyncContainerExecutor + internal class BatchAsyncContainerExecutor : IDisposable { private const int DefaultDispatchTimer = 10; private const int MinimumDispatchTimerInSeconds = 1; @@ -84,17 +84,17 @@ public async Task AddAsync( await this.ValidateOperationAsync(operation, itemRequestOptions, cancellationToken); string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false); - BatchAsyncStreamer streamer = await this.GetOrAddStreamerForPartitionKeyRangeAsync(resolvedPartitionKeyRangeId).ConfigureAwait(false); + BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); await streamer.AddAsync(context); return await context.Task; } - public async Task DisposeAsync() + public void Dispose() { foreach (KeyValuePair streamer in this.streamersByPartitionKeyRange) { - await streamer.Value.DisposeAsync(); + streamer.Value.Dispose(); } foreach (KeyValuePair limiter in this.limitersByPartitionkeyRange) @@ -202,7 +202,7 @@ private async Task ReBatchAsync( CancellationToken cancellationToken) { string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(context.Operation, cancellationToken).ConfigureAwait(false); - BatchAsyncStreamer streamer = await this.GetOrAddStreamerForPartitionKeyRangeAsync(resolvedPartitionKeyRangeId).ConfigureAwait(false); + BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); await streamer.AddAsync(context); } @@ -313,7 +313,7 @@ private async Task ExecuteServerRequestAs } } - private async Task GetOrAddStreamerForPartitionKeyRangeAsync(string partitionKeyRangeId) + private BatchAsyncStreamer GetOrAddStreamerForPartitionKeyRange(string partitionKeyRangeId) { if (this.streamersByPartitionKeyRange.TryGetValue(partitionKeyRangeId, out BatchAsyncStreamer streamer)) { @@ -323,7 +323,7 @@ private async Task GetOrAddStreamerForPartitionKeyRangeAsync BatchAsyncStreamer newStreamer = new BatchAsyncStreamer(this.maxServerRequestOperationCount, this.maxServerRequestBodyLength, this.dispatchTimerInSeconds, this.timerPool, this.cosmosClientContext.CosmosSerializer, this.ExecuteAsync); if (!this.streamersByPartitionKeyRange.TryAdd(partitionKeyRangeId, newStreamer)) { - await newStreamer.DisposeAsync(); + newStreamer.Dispose(); } return this.streamersByPartitionKeyRange[partitionKeyRangeId]; diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 510a36ea3b..6b4479d804 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -18,9 +18,8 @@ namespace Microsoft.Azure.Cosmos /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. /// /// - internal class BatchAsyncStreamer + internal class BatchAsyncStreamer : IDisposable { - private readonly List previousDispatchedTasks = new List(); private readonly SemaphoreSlim dispatchLimiter; private readonly int maxBatchOperationCount; private readonly int maxBatchByteSize; @@ -92,20 +91,18 @@ public async Task AddAsync(BatchAsyncOperationContext context) if (toDispatch != null) { - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + // Discarded for Fire & Forget + _ = toDispatch.DispatchAsync(this.cancellationTokenSource.Token); } } - public async Task DisposeAsync() + public void Dispose() { this.cancellationTokenSource.Cancel(); this.cancellationTokenSource.Dispose(); this.currentTimer.CancelTimer(); - foreach (Task previousDispatchedTask in this.previousDispatchedTasks) - { - await previousDispatchedTask; - } - + this.currentTimer = null; + this.timerTask = null; this.dispatchLimiter.Dispose(); } @@ -133,7 +130,8 @@ private async Task DispatchTimerAsync() if (toDispatch != null) { - this.previousDispatchedTasks.Add(toDispatch.DispatchAsync(this.cancellationTokenSource.Token)); + // Discarded for Fire & Forget + _ = toDispatch.DispatchAsync(this.cancellationTokenSource.Token); } this.ResetTimer(); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 0f7002d354..771952a066 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -145,7 +145,7 @@ public async Task DispatchesAsync() } [TestMethod] - public async Task DisposeAsyncShouldDisposeBatcher() + public async Task DisposeShouldDisposeBatcher() { // Expect all operations to complete as their batches get dispached BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); @@ -159,7 +159,7 @@ public async Task DisposeAsyncShouldDisposeBatcher() await Task.WhenAll(contexts); - await batchAsyncStreamer.DisposeAsync(); + batchAsyncStreamer.Dispose(); var newContext = CreateContext(new ItemBatchOperation(OperationType.Create, 0, "0")); // Disposed batcher's internal cancellation was signaled await Assert.ThrowsExceptionAsync(() => batchAsyncStreamer.AddAsync(newContext)); From 784f3b71f8c400ccfb921830af523258947ad7f2 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Mon, 12 Aug 2019 10:37:45 -0700 Subject: [PATCH 35/41] Test --- .../Batch/BatchAsyncContainerExecutorTests.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs index 4d1c882441..393adc41f7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs @@ -65,7 +65,7 @@ public async Task DoOperationsAsync() Assert.IsNotNull(storedDoc.Resource); } - await executor.DisposeAsync(); + executor.Dispose(); } [TestMethod] From fda8fbd34b9b500a84f69e56f33f90442b30af64 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Wed, 14 Aug 2019 10:55:22 -0700 Subject: [PATCH 36/41] Comments --- .../src/Batch/BatchAsyncBatcher.cs | 31 +++++--- .../src/Batch/BatchAsyncContainerExecutor.cs | 76 ++++++++----------- .../src/Batch/BatchAsyncOperationContext.cs | 28 +++++-- .../src/Batch/BatchAsyncStreamer.cs | 26 +++---- ...e.cs => PartitionKeyRangeBatchResponse.cs} | 10 +-- ...=> PartitionKeyRangeServerBatchRequest.cs} | 14 ++-- .../Batch/BatchAsyncBatcherTests.cs | 11 ++- .../Batch/BatchAsyncStreamerTests.cs | 39 +++------- .../Batch/PartitionKeyBatchResponseTests.cs | 4 +- ...titionKeyRangeBatchExecutionResultTests.cs | 2 +- 10 files changed, 116 insertions(+), 125 deletions(-) rename Microsoft.Azure.Cosmos/src/Batch/{PartitionKeyBatchResponse.cs => PartitionKeyRangeBatchResponse.cs} (96%) rename Microsoft.Azure.Cosmos/src/Batch/{PartitionKeyServerBatchRequest.cs => PartitionKeyRangeServerBatchRequest.cs} (83%) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index b5084f831b..60676dd24e 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -18,7 +18,7 @@ internal class BatchAsyncBatcher { private readonly CosmosSerializer CosmosSerializer; private readonly List batchOperations; - private readonly Func, CancellationToken, Task> executor; + private readonly BatchAsyncBatcherExecuteDelegate executor; private readonly int maxBatchByteSize; private readonly int maxBatchOperationCount; private long currentSize = 0; @@ -30,7 +30,7 @@ public BatchAsyncBatcher( int maxBatchOperationCount, int maxBatchByteSize, CosmosSerializer cosmosSerializer, - Func, CancellationToken, Task> executor) + BatchAsyncBatcherExecuteDelegate executor) { if (maxBatchOperationCount < 1) { @@ -59,44 +59,47 @@ public BatchAsyncBatcher( this.CosmosSerializer = cosmosSerializer; } - public bool TryAdd(BatchAsyncOperationContext batchAsyncOperation) + public virtual bool TryAdd(BatchAsyncOperationContext operationContext) { if (this.dispached) { + DefaultTrace.TraceCritical($"Add operation attempted on dispatched batch."); return false; } - if (batchAsyncOperation == null) + if (operationContext == null) { - throw new ArgumentNullException(nameof(batchAsyncOperation)); + throw new ArgumentNullException(nameof(operationContext)); } if (this.batchOperations.Count == this.maxBatchOperationCount) { + DefaultTrace.TraceVerbose($"Batch is full - Max operation count {this.maxBatchOperationCount} reached."); return false; } - int itemByteSize = batchAsyncOperation.Operation.GetApproximateSerializedLength(); + int itemByteSize = operationContext.Operation.GetApproximateSerializedLength(); if (itemByteSize + this.currentSize > this.maxBatchByteSize) { + DefaultTrace.TraceVerbose($"Batch is full - Max byte size {this.maxBatchByteSize} reached."); return false; } this.currentSize += itemByteSize; // Operation index is in the scope of the current batch - batchAsyncOperation.Operation.OperationIndex = this.batchOperations.Count; - batchAsyncOperation.CurrentBatcher = this; - this.batchOperations.Add(batchAsyncOperation); + operationContext.Operation.OperationIndex = this.batchOperations.Count; + operationContext.CurrentBatcher = this; + this.batchOperations.Add(operationContext); return true; } - public async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) + public virtual async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) { try { - using (PartitionKeyBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken)) + using (PartitionKeyRangeBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken)) { for (int index = 0; index < this.batchOperations.Count; index++) { @@ -125,4 +128,10 @@ public bool TryAdd(BatchAsyncOperationContext batchAsyncOperation) } } } + + /// + /// Executor implementation that processes a list of operations. + /// + /// An instance of . + internal delegate Task BatchAsyncBatcherExecuteDelegate(IReadOnlyList operationContexts, CancellationToken cancellationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index ba11c1ebdc..7dda62d95d 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -9,7 +9,6 @@ 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.Routing; @@ -19,7 +18,9 @@ namespace Microsoft.Azure.Cosmos /// Bulk batch executor for operations in the same container. /// /// - /// It groups operations by Partition Key Range and sends them to the Batch API and then unifies results as they become available. It uses as batch processor and as batch executing handler. + /// It maintains one for each Partition Key Range, which allows independent execution of requests. + /// Semaphores are in place to rate limit the operations at the Streamer / Partition Key Range level, this means that we can send parallel and independent requests to different Partition Key Ranges, but for the same Range, requests will be limited. + /// Two delegate implementations define how a particular request should be executed, and how operations should be retried. When the dispatches a batch, the batch will create a request and call the execute delegate, if conditions are met, it might call the retry delegate. /// /// internal class BatchAsyncContainerExecutor : IDisposable @@ -86,7 +87,7 @@ public async Task AddAsync( string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false); BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); - await streamer.AddAsync(context); + streamer.Add(context); return await context.Task; } @@ -170,31 +171,12 @@ private static void AddHeadersToRequestMessage(RequestMessage requestMessage, st /// /// If because of HybridRow serialization overhead, not all operations fit in the request, we send those extra operations in a separate request. /// - private static IEnumerable GetOverflowOperations( - PartitionKeyServerBatchRequest request, - IEnumerable operationsSentToRequest) + private static int GetOverflowOperations( + PartitionKeyRangeServerBatchRequest request, + IReadOnlyList operationsSentToRequest) { - int totalOperations = operationsSentToRequest.Count(); - int operationsThatOverflowed = totalOperations - request.Operations.Count; - if (operationsThatOverflowed == 0) - { - return Enumerable.Empty(); - } - - return operationsSentToRequest.Skip(totalOperations - operationsThatOverflowed); - } - - private static IReadOnlyList GetOperationsToRetry( - IReadOnlyList baseOperations, - IEnumerable indexes) - { - List operations = new List(); - foreach (int index in indexes) - { - operations.Add(baseOperations[index]); - } - - return operations; + int totalOperations = operationsSentToRequest.Count; + return totalOperations - request.Operations.Count; } private async Task ReBatchAsync( @@ -203,21 +185,22 @@ private async Task ReBatchAsync( { string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(context.Operation, cancellationToken).ConfigureAwait(false); BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); - await streamer.AddAsync(context); + streamer.Add(context); } - private async Task ExecuteAsync( - IReadOnlyList operations, + private async Task ExecuteAsync( + IReadOnlyList operationContexts, CancellationToken cancellationToken) { // All operations should be for the same PKRange - string partitionKeyRangeId = operations[0].PartitionKeyRangeId; - PartitionKeyServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operations.Select(o => o.Operation), cancellationToken); + string partitionKeyRangeId = operationContexts[0].PartitionKeyRangeId; + PartitionKeyRangeServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operationContexts, cancellationToken); // Any overflow goes to a new batch - IEnumerable overFlowOperations = BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operations); - foreach (BatchAsyncOperationContext overflowedContext in overFlowOperations) + int overFlowOperations = BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operationContexts); + while (overFlowOperations > 0) { - await this.ReBatchAsync(overflowedContext, cancellationToken); + await this.ReBatchAsync(operationContexts[operationContexts.Count - overFlowOperations], cancellationToken); + overFlowOperations--; } PartitionKeyRangeBatchExecutionResult result = await this.ExecuteServerRequestAsync(serverRequest, cancellationToken); @@ -229,14 +212,13 @@ private async Task ExecuteAsync( } else { - IReadOnlyList retryContexts = BatchAsyncContainerExecutor.GetOperationsToRetry(operations, result.Operations.Select(o => o.OperationIndex)); - foreach (BatchAsyncOperationContext retryContext in retryContexts) + foreach (ItemBatchOperation operationToRetry in result.Operations) { - await this.ReBatchAsync(retryContext, cancellationToken); + await this.ReBatchAsync(operationContexts[operationToRetry.OperationIndex], cancellationToken); } } - return new PartitionKeyBatchResponse(operations.Count, responses, this.cosmosClientContext.CosmosSerializer); + return new PartitionKeyRangeBatchResponse(operationContexts.Count, responses, this.cosmosClientContext.CosmosSerializer); } private async Task ResolvePartitionKeyRangeIdAsync( @@ -266,13 +248,19 @@ private async Task FillOperationPropertiesAsync(ItemBatchOperation operation, Ca } } - private async Task CreateServerRequestAsync( + private async Task CreateServerRequestAsync( string partitionKeyRangeId, - IEnumerable operations, + IReadOnlyList operationContexts, CancellationToken cancellationToken) { - ArraySegment operationsArraySegment = new ArraySegment(operations.ToArray()); - PartitionKeyServerBatchRequest request = await PartitionKeyServerBatchRequest.CreateAsync( + ItemBatchOperation[] operations = new ItemBatchOperation[operationContexts.Count]; + for (int i = 0; i < operationContexts.Count; i++) + { + operations[i] = operationContexts[i].Operation; + } + + ArraySegment operationsArraySegment = new ArraySegment(operations); + PartitionKeyRangeServerBatchRequest request = await PartitionKeyRangeServerBatchRequest.CreateAsync( partitionKeyRangeId, operationsArraySegment, this.maxServerRequestBodyLength, @@ -285,7 +273,7 @@ private async Task CreateServerRequestAsync( } private async Task ExecuteServerRequestAsync( - PartitionKeyServerBatchRequest serverRequest, + PartitionKeyRangeServerBatchRequest serverRequest, CancellationToken cancellationToken) { SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(serverRequest.PartitionKeyRangeId); diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs index 19dfece54c..fe478a4e86 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs @@ -5,8 +5,8 @@ namespace Microsoft.Azure.Cosmos { using System; - using System.Diagnostics; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Core.Trace; /// /// Context for a particular Batch operation. @@ -35,8 +35,11 @@ public void Complete( BatchAsyncBatcher completer, BatchOperationResult result) { - Debug.Assert(this.CurrentBatcher == null || completer == this.CurrentBatcher); - this.taskCompletionSource.SetResult(result); + if (this.AssertBatcher(completer)) + { + this.taskCompletionSource.SetResult(result); + } + this.Dispose(); } @@ -44,8 +47,11 @@ public void Fail( BatchAsyncBatcher completer, Exception exception) { - Debug.Assert(this.CurrentBatcher == null || completer == this.CurrentBatcher); - this.taskCompletionSource.SetException(exception); + if (this.AssertBatcher(completer)) + { + this.taskCompletionSource.SetException(exception); + } + this.Dispose(); } @@ -54,5 +60,17 @@ public void Dispose() this.Operation.Dispose(); this.CurrentBatcher = null; } + + private bool AssertBatcher(BatchAsyncBatcher completer) + { + if (!object.ReferenceEquals(completer, this.CurrentBatcher)) + { + DefaultTrace.TraceCritical($"Operation {this.Operation.Id} was completed by incorrect batcher."); + this.taskCompletionSource.SetException(new Exception($"Operation {this.Operation.Id} was completed by incorrect batcher.")); + return false; + } + + return true; + } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 6b4479d804..037fe496e2 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -5,25 +5,25 @@ namespace Microsoft.Azure.Cosmos { using System; - using System.Collections.Generic; - using System.Diagnostics; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Documents; /// - /// Handles operation queueing and dispatching. + /// Handles operation queueing and dispatching. + /// Fills batches efficiently and maintains a timer for early dispatching in case of partially-filled batches and to optimize for throughput. /// /// - /// will add the operation to the current batcher or if full, dispatch it, create a new one and add the operation to it. + /// There is always one batch at a time being filled. Locking is in place to avoid concurrent threads trying to Add operations while the timer might be Dispatching the current batch. + /// The current batch is dispatched and a new one is readied to be filled by new operations, the dispatched batch runs independently through a fire & forget pattern. /// /// internal class BatchAsyncStreamer : IDisposable { - private readonly SemaphoreSlim dispatchLimiter; + private readonly object dispatchLimiter = new object(); private readonly int maxBatchOperationCount; private readonly int maxBatchByteSize; - private readonly Func, CancellationToken, Task> executor; + private readonly BatchAsyncBatcherExecuteDelegate executor; private readonly int dispatchTimerInSeconds; private readonly CosmosSerializer cosmosSerializer; private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); @@ -38,7 +38,7 @@ public BatchAsyncStreamer( int dispatchTimerInSeconds, TimerPool timerPool, CosmosSerializer cosmosSerializer, - Func, CancellationToken, Task> executor) + BatchAsyncBatcherExecuteDelegate executor) { if (maxBatchOperationCount < 1) { @@ -71,16 +71,15 @@ public BatchAsyncStreamer( this.dispatchTimerInSeconds = dispatchTimerInSeconds; this.timerPool = timerPool; this.cosmosSerializer = cosmosSerializer; - this.dispatchLimiter = new SemaphoreSlim(1, 1); this.currentBatcher = this.CreateBatchAsyncBatcher(); this.ResetTimer(); } - public async Task AddAsync(BatchAsyncOperationContext context) + public void Add(BatchAsyncOperationContext context) { BatchAsyncBatcher toDispatch = null; - using (await this.dispatchLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) + lock (this.dispatchLimiter) { while (!this.currentBatcher.TryAdd(context)) { @@ -103,7 +102,6 @@ public void Dispose() this.currentTimer.CancelTimer(); this.currentTimer = null; this.timerTask = null; - this.dispatchLimiter.Dispose(); } private void ResetTimer() @@ -111,11 +109,11 @@ private void ResetTimer() this.currentTimer = this.timerPool.GetPooledTimer(this.dispatchTimerInSeconds); this.timerTask = this.currentTimer.StartTimerAsync().ContinueWith((task) => { - return this.DispatchTimerAsync(); + this.DispatchTimer(); }, this.cancellationTokenSource.Token); } - private async Task DispatchTimerAsync() + private void DispatchTimer() { if (this.cancellationTokenSource.IsCancellationRequested) { @@ -123,7 +121,7 @@ private async Task DispatchTimerAsync() } BatchAsyncBatcher toDispatch; - using (await this.dispatchLimiter.UsingWaitAsync(this.cancellationTokenSource.Token)) + lock (this.dispatchLimiter) { toDispatch = this.GetBatchToDispatchAndCreate(); } diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchResponse.cs similarity index 96% rename from Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs rename to Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchResponse.cs index 2a1b328dad..af277b05d3 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyBatchResponse.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchResponse.cs @@ -14,21 +14,21 @@ namespace Microsoft.Azure.Cosmos /// /// Response of a cross partition key batch request. /// - internal class PartitionKeyBatchResponse : BatchResponse + internal class PartitionKeyRangeBatchResponse : BatchResponse { // Results sorted in the order operations had been added. private readonly BatchOperationResult[] resultsByOperationIndex; private bool isDisposed; /// - /// Initializes a new instance of the class. + /// Initializes a new instance of the class. /// /// Completion status code of the batch request. /// Provides further details about why the batch was not processed. /// Operations that were supposed to be executed, but weren't. /// The reason for failure if any. // This constructor is expected to be used when the batch is not executed at all (if it is a bad request). - internal PartitionKeyBatchResponse( + internal PartitionKeyRangeBatchResponse( HttpStatusCode statusCode, SubStatusCodes subStatusCode, string errorMessage, @@ -38,12 +38,12 @@ internal PartitionKeyBatchResponse( } /// - /// Initializes a new instance of the class. + /// Initializes a new instance of the class. /// /// Original operations that generated the server responses. /// Responses from the server. /// Serializer to deserialize response resource body streams. - internal PartitionKeyBatchResponse( + internal PartitionKeyRangeBatchResponse( int originalOperationsCount, IEnumerable serverResponses, CosmosSerializer serializer) diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyServerBatchRequest.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeServerBatchRequest.cs similarity index 83% rename from Microsoft.Azure.Cosmos/src/Batch/PartitionKeyServerBatchRequest.cs rename to Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeServerBatchRequest.cs index 98e73c5a0b..24a0adfda9 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyServerBatchRequest.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeServerBatchRequest.cs @@ -8,16 +8,16 @@ namespace Microsoft.Azure.Cosmos using System.Threading; using System.Threading.Tasks; - internal sealed class PartitionKeyServerBatchRequest : ServerBatchRequest + internal sealed class PartitionKeyRangeServerBatchRequest : ServerBatchRequest { /// - /// Initializes a new instance of the class. + /// Initializes a new instance of the class. /// /// The partition key range id associated with all requests. /// Maximum length allowed for the request body. /// Maximum number of operations allowed in the request. /// Serializer to serialize user provided objects to JSON. - public PartitionKeyServerBatchRequest( + public PartitionKeyRangeServerBatchRequest( string partitionKeyRangeId, int maxBodyLength, int maxOperationCount, @@ -33,7 +33,7 @@ public PartitionKeyServerBatchRequest( public string PartitionKeyRangeId { get; } /// - /// Creates an instance of . + /// Creates an instance of . /// In case of direct mode requests, all the operations are expected to belong to the same PartitionKeyRange. /// The body of the request is populated with operations till it reaches the provided maxBodyLength. /// @@ -44,8 +44,8 @@ public PartitionKeyServerBatchRequest( /// Whether to stop adding operations to the request once there is non-continuity in the operation indexes. /// Serializer to serialize user provided objects to JSON. /// representing request cancellation. - /// A newly created instance of . - public static async Task CreateAsync( + /// A newly created instance of . + public static async Task CreateAsync( string partitionKeyRangeId, ArraySegment operations, int maxBodyLength, @@ -54,7 +54,7 @@ public static async Task CreateAsync( CosmosSerializer serializer, CancellationToken cancellationToken) { - PartitionKeyServerBatchRequest request = new PartitionKeyServerBatchRequest(partitionKeyRangeId, maxBodyLength, maxOperationCount, serializer); + PartitionKeyRangeServerBatchRequest request = new PartitionKeyRangeServerBatchRequest(partitionKeyRangeId, maxBodyLength, maxOperationCount, serializer); await request.CreateBodyStreamAsync(operations, cancellationToken, ensureContinuousOperationIndexes); return request; } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 11c05be1fb..1f1f805b7b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -13,7 +13,6 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Threading.Tasks; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; - using Moq; [TestClass] public class BatchAsyncBatcherTests @@ -22,7 +21,7 @@ public class BatchAsyncBatcherTests private ItemBatchOperation ItemBatchOperation = new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); - private Func, CancellationToken, Task> Executor + private BatchAsyncBatcherExecuteDelegate Executor = async (IReadOnlyList operations, CancellationToken cancellation) => { List results = new List(); @@ -55,11 +54,11 @@ private Func, CancellationToken, Task< batchRequest, new CosmosJsonDotNetSerializer()); - return new PartitionKeyBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + return new PartitionKeyRangeBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); }; // The response will include all but 2 operation responses - private Func, CancellationToken, Task> ExecutorWithLessResponses + private BatchAsyncBatcherExecuteDelegate ExecutorWithLessResponses = async (IReadOnlyList operations, CancellationToken cancellation) => { int operationCount = operations.Count - 2; @@ -93,10 +92,10 @@ private Func, CancellationToken, Task< batchRequest, new CosmosJsonDotNetSerializer()); - return new PartitionKeyBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + return new PartitionKeyRangeBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); }; - private Func, CancellationToken, Task> ExecutorWithFailure + private BatchAsyncBatcherExecuteDelegate ExecutorWithFailure = (IReadOnlyList operations, CancellationToken cancellation) => { throw expectedException; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 771952a066..dbb89dbeca 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -23,7 +23,7 @@ public class BatchAsyncStreamerTests private TimerPool TimerPool = new TimerPool(1); // Executor just returns a reponse matching the Id with Etag - private Func, CancellationToken, Task> Executor + private BatchAsyncBatcherExecuteDelegate Executor = async (IReadOnlyList operations, CancellationToken cancellation) => { List results = new List(); @@ -56,11 +56,11 @@ private Func, CancellationToken, Task< batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); return response; }; - private Func, CancellationToken, Task> ExecutorWithFailure + private BatchAsyncBatcherExecuteDelegate ExecutorWithFailure = (IReadOnlyList operations, CancellationToken cancellation) => { throw expectedException; @@ -102,8 +102,8 @@ public void ValidatesSerializer() public async Task ExceptionsOnBatchBubbleUpAsync() { BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); - var context = CreateContext(this.ItemBatchOperation); - await batchAsyncStreamer.AddAsync(context); + BatchAsyncOperationContext context = CreateContext(this.ItemBatchOperation); + batchAsyncStreamer.Add(context); Exception capturedException = await Assert.ThrowsExceptionAsync(() => context.Task); Assert.AreEqual(expectedException, capturedException); } @@ -113,8 +113,8 @@ public async Task TimerDispatchesAsync() { // Bigger batch size than the amount of operations, timer should dispatch BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); - var context = CreateContext(this.ItemBatchOperation); - await batchAsyncStreamer.AddAsync(context); + BatchAsyncOperationContext context = CreateContext(this.ItemBatchOperation); + batchAsyncStreamer.Add(context); BatchOperationResult result = await context.Task; Assert.AreEqual(this.ItemBatchOperation.Id, result.ETag); @@ -128,8 +128,8 @@ public async Task DispatchesAsync() List> contexts = new List>(10); for (int i = 0; i < 10; i++) { - var context = CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())); - await batchAsyncStreamer.AddAsync(context); + BatchAsyncOperationContext context = CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())); + batchAsyncStreamer.Add(context); contexts.Add(context.Task); } @@ -144,27 +144,6 @@ public async Task DispatchesAsync() } } - [TestMethod] - public async Task DisposeShouldDisposeBatcher() - { - // Expect all operations to complete as their batches get dispached - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); - List> contexts = new List>(10); - for (int i = 0; i < 10; i++) - { - var context = CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())); - await batchAsyncStreamer.AddAsync(context); - contexts.Add(context.Task); - } - - await Task.WhenAll(contexts); - - batchAsyncStreamer.Dispose(); - var newContext = CreateContext(new ItemBatchOperation(OperationType.Create, 0, "0")); - // Disposed batcher's internal cancellation was signaled - await Assert.ThrowsExceptionAsync(() => batchAsyncStreamer.AddAsync(newContext)); - } - private static BatchAsyncOperationContext CreateContext(ItemBatchOperation operation) => new BatchAsyncOperationContext(string.Empty, operation); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs index a10a70baf4..5e48bbf41f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs @@ -20,7 +20,7 @@ public class PartitionKeyBatchResponseTests public void StatusCodesAreSet() { const string errorMessage = "some error"; - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(HttpStatusCode.NotFound, SubStatusCodes.ClientTcpChannelFull, errorMessage, null); + PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(HttpStatusCode.NotFound, SubStatusCodes.ClientTcpChannelFull, errorMessage, null); Assert.AreEqual(HttpStatusCode.NotFound, response.StatusCode); Assert.AreEqual(SubStatusCodes.ClientTcpChannelFull, response.SubStatusCode); Assert.AreEqual(errorMessage, response.ErrorMessage); @@ -58,7 +58,7 @@ public async Task StatusCodesAreSetThroughResponseAsync() batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(arrayOperations.Length, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(arrayOperations.Length, new List { batchresponse }, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs index 38b423589e..17f82d9eb2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs @@ -63,7 +63,7 @@ public async Task StatusCodesAreSetThroughResponseAsync() batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyBatchResponse response = new PartitionKeyBatchResponse(arrayOperations.Length, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(arrayOperations.Length, new List { batchresponse }, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } From 15154deb85aa0413545a9414a35b805d26af0404 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Thu, 15 Aug 2019 11:29:27 -0700 Subject: [PATCH 37/41] Batcher refactor --- .../src/Batch/BatchAsyncBatcher.cs | 98 ++++++++- .../src/Batch/BatchAsyncContainerExecutor.cs | 72 +------ .../src/Batch/BatchAsyncStreamer.cs | 14 +- .../Batch/BatchAsyncBatcherTests.cs | 194 ++++++++++++++---- .../Batch/BatchAsyncStreamerTests.cs | 47 +++-- 5 files changed, 283 insertions(+), 142 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 60676dd24e..42250f4bbe 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -11,14 +11,23 @@ namespace Microsoft.Azure.Cosmos using Microsoft.Azure.Cosmos.Core.Trace; /// - /// Maintains a batch of operations and dispatches the batch through an executor. Maps results into the original operation contexts. + /// Maintains a batch of operations and dispatches it as a unit of work. /// + /// + /// The dispatch process consists of: + /// 1. Creating a . + /// 2. Verifying overflow that might happen due to HybridRow serialization. Any operations that did not fit, get sent to the . + /// 3. Execution of the request gets delegated to . + /// 4. If there was a split detected, all operations in the request, are sent to the for re-queueing. + /// 5. The result of the request is used to wire up all responses with the original Tasks for each operation. + /// /// internal class BatchAsyncBatcher { - private readonly CosmosSerializer CosmosSerializer; + private readonly CosmosSerializer cosmosSerializer; private readonly List batchOperations; private readonly BatchAsyncBatcherExecuteDelegate executor; + private readonly BatchAsyncBatcherRetryDelegate retrier; private readonly int maxBatchByteSize; private readonly int maxBatchOperationCount; private long currentSize = 0; @@ -30,7 +39,8 @@ public BatchAsyncBatcher( int maxBatchOperationCount, int maxBatchByteSize, CosmosSerializer cosmosSerializer, - BatchAsyncBatcherExecuteDelegate executor) + BatchAsyncBatcherExecuteDelegate executor, + BatchAsyncBatcherRetryDelegate retrier) { if (maxBatchOperationCount < 1) { @@ -47,6 +57,11 @@ public BatchAsyncBatcher( throw new ArgumentNullException(nameof(executor)); } + if (retrier == null) + { + throw new ArgumentNullException(nameof(retrier)); + } + if (cosmosSerializer == null) { throw new ArgumentNullException(nameof(cosmosSerializer)); @@ -54,9 +69,10 @@ public BatchAsyncBatcher( this.batchOperations = new List(maxBatchOperationCount); this.executor = executor; + this.retrier = retrier; this.maxBatchByteSize = maxBatchByteSize; this.maxBatchOperationCount = maxBatchOperationCount; - this.CosmosSerializer = cosmosSerializer; + this.cosmosSerializer = cosmosSerializer; } public virtual bool TryAdd(BatchAsyncOperationContext operationContext) @@ -97,9 +113,33 @@ public virtual bool TryAdd(BatchAsyncOperationContext operationContext) public virtual async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) { + PartitionKeyRangeServerBatchRequest serverRequest = await this.CreateServerRequestAsync(cancellationToken); + // Any overflow goes to a new batch + int overFlowOperations = this.GetOverflowOperations(serverRequest, this.batchOperations); + while (overFlowOperations > 0) + { + await this.retrier(this.batchOperations[this.batchOperations.Count - overFlowOperations], cancellationToken); + overFlowOperations--; + } + try { - using (PartitionKeyRangeBatchResponse batchResponse = await this.executor(this.batchOperations, cancellationToken)) + PartitionKeyRangeBatchExecutionResult result = await this.executor(serverRequest, cancellationToken); + + List responses = new List(result.ServerResponses.Count); + if (!result.ContainsSplit()) + { + responses.AddRange(result.ServerResponses); + } + else + { + foreach (ItemBatchOperation operationToRetry in result.Operations) + { + await this.retrier(this.batchOperations[operationToRetry.OperationIndex], cancellationToken); + } + } + + using (PartitionKeyRangeBatchResponse batchResponse = new PartitionKeyRangeBatchResponse(serverRequest.Operations.Count, responses, this.cosmosSerializer)) { for (int index = 0; index < this.batchOperations.Count; index++) { @@ -115,9 +155,10 @@ public virtual bool TryAdd(BatchAsyncOperationContext operationContext) catch (Exception ex) { DefaultTrace.TraceError("Exception during BatchAsyncBatcher: {0}", ex); - // Exceptions happening during execution fail all the Tasks - foreach (BatchAsyncOperationContext context in this.batchOperations) + // Exceptions happening during execution fail all the Tasks part of the request (excluding overflow) + foreach (ItemBatchOperation itemBatchOperation in serverRequest.Operations) { + BatchAsyncOperationContext context = this.batchOperations[itemBatchOperation.OperationIndex]; context.Fail(this, ex); } } @@ -127,11 +168,52 @@ public virtual bool TryAdd(BatchAsyncOperationContext operationContext) this.dispached = true; } } + + /// + /// If because of HybridRow serialization overhead, not all operations fit in the request, we send those extra operations in a separate request. + /// + internal virtual int GetOverflowOperations( + PartitionKeyRangeServerBatchRequest request, + IReadOnlyList operationsSentToRequest) + { + int totalOperations = operationsSentToRequest.Count; + return totalOperations - request.Operations.Count; + } + + private async Task CreateServerRequestAsync(CancellationToken cancellationToken) + { + // All operations should be for the same PKRange + string partitionKeyRangeId = this.batchOperations[0].PartitionKeyRangeId; + + ItemBatchOperation[] operations = new ItemBatchOperation[this.batchOperations.Count]; + for (int i = 0; i < this.batchOperations.Count; i++) + { + operations[i] = this.batchOperations[i].Operation; + } + + ArraySegment operationsArraySegment = new ArraySegment(operations); + PartitionKeyRangeServerBatchRequest request = await PartitionKeyRangeServerBatchRequest.CreateAsync( + partitionKeyRangeId, + operationsArraySegment, + this.maxBatchByteSize, + this.maxBatchOperationCount, + ensureContinuousOperationIndexes: false, + serializer: this.cosmosSerializer, + cancellationToken: cancellationToken).ConfigureAwait(false); + + return request; + } } /// /// Executor implementation that processes a list of operations. /// /// An instance of . - internal delegate Task BatchAsyncBatcherExecuteDelegate(IReadOnlyList operationContexts, CancellationToken cancellationToken); + internal delegate Task BatchAsyncBatcherExecuteDelegate(PartitionKeyRangeServerBatchRequest request, CancellationToken cancellationToken); + + /// + /// Delegate to process a request for retry an operation + /// + /// An instance of . + internal delegate Task BatchAsyncBatcherRetryDelegate(BatchAsyncOperationContext operationContext, CancellationToken cancellationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 7dda62d95d..6890be7fe3 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -168,17 +168,6 @@ private static void AddHeadersToRequestMessage(RequestMessage requestMessage, st requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchRequest, bool.TrueString); } - /// - /// If because of HybridRow serialization overhead, not all operations fit in the request, we send those extra operations in a separate request. - /// - private static int GetOverflowOperations( - PartitionKeyRangeServerBatchRequest request, - IReadOnlyList operationsSentToRequest) - { - int totalOperations = operationsSentToRequest.Count; - return totalOperations - request.Operations.Count; - } - private async Task ReBatchAsync( BatchAsyncOperationContext context, CancellationToken cancellationToken) @@ -188,39 +177,6 @@ private async Task ReBatchAsync( streamer.Add(context); } - private async Task ExecuteAsync( - IReadOnlyList operationContexts, - CancellationToken cancellationToken) - { - // All operations should be for the same PKRange - string partitionKeyRangeId = operationContexts[0].PartitionKeyRangeId; - PartitionKeyRangeServerBatchRequest serverRequest = await this.CreateServerRequestAsync(partitionKeyRangeId, operationContexts, cancellationToken); - // Any overflow goes to a new batch - int overFlowOperations = BatchAsyncContainerExecutor.GetOverflowOperations(serverRequest, operationContexts); - while (overFlowOperations > 0) - { - await this.ReBatchAsync(operationContexts[operationContexts.Count - overFlowOperations], cancellationToken); - overFlowOperations--; - } - - PartitionKeyRangeBatchExecutionResult result = await this.ExecuteServerRequestAsync(serverRequest, cancellationToken); - - List responses = new List(result.ServerResponses.Count); - if (!result.ContainsSplit()) - { - responses.AddRange(result.ServerResponses); - } - else - { - foreach (ItemBatchOperation operationToRetry in result.Operations) - { - await this.ReBatchAsync(operationContexts[operationToRetry.OperationIndex], cancellationToken); - } - } - - return new PartitionKeyRangeBatchResponse(operationContexts.Count, responses, this.cosmosClientContext.CosmosSerializer); - } - private async Task ResolvePartitionKeyRangeIdAsync( ItemBatchOperation operation, CancellationToken cancellationToken) @@ -248,31 +204,7 @@ private async Task FillOperationPropertiesAsync(ItemBatchOperation operation, Ca } } - private async Task CreateServerRequestAsync( - string partitionKeyRangeId, - IReadOnlyList operationContexts, - CancellationToken cancellationToken) - { - ItemBatchOperation[] operations = new ItemBatchOperation[operationContexts.Count]; - for (int i = 0; i < operationContexts.Count; i++) - { - operations[i] = operationContexts[i].Operation; - } - - ArraySegment operationsArraySegment = new ArraySegment(operations); - PartitionKeyRangeServerBatchRequest request = await PartitionKeyRangeServerBatchRequest.CreateAsync( - partitionKeyRangeId, - operationsArraySegment, - this.maxServerRequestBodyLength, - this.maxServerRequestOperationCount, - ensureContinuousOperationIndexes: false, - serializer: this.cosmosClientContext.CosmosSerializer, - cancellationToken: cancellationToken).ConfigureAwait(false); - - return request; - } - - private async Task ExecuteServerRequestAsync( + private async Task ExecuteAsync( PartitionKeyRangeServerBatchRequest serverRequest, CancellationToken cancellationToken) { @@ -308,7 +240,7 @@ private BatchAsyncStreamer GetOrAddStreamerForPartitionKeyRange(string partition return streamer; } - BatchAsyncStreamer newStreamer = new BatchAsyncStreamer(this.maxServerRequestOperationCount, this.maxServerRequestBodyLength, this.dispatchTimerInSeconds, this.timerPool, this.cosmosClientContext.CosmosSerializer, this.ExecuteAsync); + BatchAsyncStreamer newStreamer = new BatchAsyncStreamer(this.maxServerRequestOperationCount, this.maxServerRequestBodyLength, this.dispatchTimerInSeconds, this.timerPool, this.cosmosClientContext.CosmosSerializer, this.ExecuteAsync, this.ReBatchAsync); if (!this.streamersByPartitionKeyRange.TryAdd(partitionKeyRangeId, newStreamer)) { newStreamer.Dispose(); diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index 037fe496e2..f568a8d5b4 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -15,7 +15,7 @@ namespace Microsoft.Azure.Cosmos /// /// /// There is always one batch at a time being filled. Locking is in place to avoid concurrent threads trying to Add operations while the timer might be Dispatching the current batch. - /// The current batch is dispatched and a new one is readied to be filled by new operations, the dispatched batch runs independently through a fire & forget pattern. + /// The current batch is dispatched and a new one is readied to be filled by new operations, the dispatched batch runs independently through a fire and forget pattern. /// /// internal class BatchAsyncStreamer : IDisposable @@ -24,6 +24,7 @@ internal class BatchAsyncStreamer : IDisposable private readonly int maxBatchOperationCount; private readonly int maxBatchByteSize; private readonly BatchAsyncBatcherExecuteDelegate executor; + private readonly BatchAsyncBatcherRetryDelegate retrier; private readonly int dispatchTimerInSeconds; private readonly CosmosSerializer cosmosSerializer; private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); @@ -38,7 +39,8 @@ public BatchAsyncStreamer( int dispatchTimerInSeconds, TimerPool timerPool, CosmosSerializer cosmosSerializer, - BatchAsyncBatcherExecuteDelegate executor) + BatchAsyncBatcherExecuteDelegate executor, + BatchAsyncBatcherRetryDelegate retrier) { if (maxBatchOperationCount < 1) { @@ -60,6 +62,11 @@ public BatchAsyncStreamer( throw new ArgumentNullException(nameof(executor)); } + if (retrier == null) + { + throw new ArgumentNullException(nameof(retrier)); + } + if (cosmosSerializer == null) { throw new ArgumentNullException(nameof(cosmosSerializer)); @@ -68,6 +75,7 @@ public BatchAsyncStreamer( this.maxBatchOperationCount = maxBatchOperationCount; this.maxBatchByteSize = maxBatchByteSize; this.executor = executor; + this.retrier = retrier; this.dispatchTimerInSeconds = dispatchTimerInSeconds; this.timerPool = timerPool; this.cosmosSerializer = cosmosSerializer; @@ -149,7 +157,7 @@ private BatchAsyncBatcher GetBatchToDispatchAndCreate() private BatchAsyncBatcher CreateBatchAsyncBatcher() { - return new BatchAsyncBatcher(this.maxBatchOperationCount, this.maxBatchByteSize, this.cosmosSerializer, this.executor); + return new BatchAsyncBatcher(this.maxBatchOperationCount, this.maxBatchByteSize, this.cosmosSerializer, this.executor, this.retrier); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 1f1f805b7b..60be012d1a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -13,30 +13,31 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Threading.Tasks; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; [TestClass] public class BatchAsyncBatcherTests { private static Exception expectedException = new Exception(); - private ItemBatchOperation ItemBatchOperation = new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); + private ItemBatchOperation ItemBatchOperation() => new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); private BatchAsyncBatcherExecuteDelegate Executor - = async (IReadOnlyList operations, CancellationToken cancellation) => + = async (PartitionKeyRangeServerBatchRequest request, CancellationToken cancellationToken) => { List results = new List(); - ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operations.Count]; + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[request.Operations.Count]; int index = 0; - foreach (BatchAsyncOperationContext operation in operations) + foreach (ItemBatchOperation operation in request.Operations) { results.Add( new BatchOperationResult(HttpStatusCode.OK) { ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), - ETag = operation.Operation.Id + ETag = operation.Id }); - arrayOperations[index++] = operation.Operation; + arrayOperations[index++] = operation; } MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); @@ -44,37 +45,76 @@ private BatchAsyncBatcherExecuteDelegate Executor SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( partitionKey: null, operations: new ArraySegment(arrayOperations), - maxBodyLength: (int)responseContent.Length * operations.Count, - maxOperationCount: operations.Count, + maxBodyLength: (int)responseContent.Length * request.Operations.Count, + maxOperationCount: request.Operations.Count, serializer: new CosmosJsonDotNetSerializer(), - cancellationToken: cancellation); + cancellationToken: cancellationToken); BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }, batchRequest, new CosmosJsonDotNetSerializer()); - return new PartitionKeyRangeBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, new List() { batchresponse }); + }; + + private BatchAsyncBatcherExecuteDelegate ExecutorWithSplit + = async (PartitionKeyRangeServerBatchRequest request, CancellationToken cancellationToken) => + { + List results = new List(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[request.Operations.Count]; + int index = 0; + foreach (ItemBatchOperation operation in request.Operations) + { + results.Add( + new BatchOperationResult(HttpStatusCode.Gone) + { + ETag = operation.Id, + SubStatusCode = SubStatusCodes.PartitionKeyRangeGone + }); + + arrayOperations[index++] = operation; + } + + MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); + + SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( + partitionKey: null, + operations: new ArraySegment(arrayOperations), + maxBodyLength: (int)responseContent.Length * request.Operations.Count, + maxOperationCount: request.Operations.Count, + serializer: new CosmosJsonDotNetSerializer(), + cancellationToken: cancellationToken); + + ResponseMessage responseMessage = new ResponseMessage(HttpStatusCode.Gone) { Content = responseContent }; + responseMessage.Headers.SubStatusCode = SubStatusCodes.PartitionKeyRangeGone; + + BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( + responseMessage, + batchRequest, + new CosmosJsonDotNetSerializer()); + + return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, new List() { batchresponse }); }; // The response will include all but 2 operation responses private BatchAsyncBatcherExecuteDelegate ExecutorWithLessResponses - = async (IReadOnlyList operations, CancellationToken cancellation) => + = async (PartitionKeyRangeServerBatchRequest request, CancellationToken cancellationToken) => { - int operationCount = operations.Count - 2; + int operationCount = request.Operations.Count - 2; List results = new List(); ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operationCount]; int index = 0; - foreach (BatchAsyncOperationContext operation in operations.Skip(1).Take(operationCount)) + foreach (ItemBatchOperation operation in request.Operations.Skip(1).Take(operationCount)) { results.Add( new BatchOperationResult(HttpStatusCode.OK) { ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), - ETag = operation.Operation.Id + ETag = operation.Id }); - arrayOperations[index++] = operation.Operation; + arrayOperations[index++] = operation; } MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); @@ -85,29 +125,34 @@ private BatchAsyncBatcherExecuteDelegate ExecutorWithLessResponses maxBodyLength: (int)responseContent.Length * operationCount, maxOperationCount: operationCount, serializer: new CosmosJsonDotNetSerializer(), - cancellationToken: cancellation); + cancellationToken: cancellationToken); BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }, batchRequest, new CosmosJsonDotNetSerializer()); - return new PartitionKeyRangeBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, new List() { batchresponse }); }; private BatchAsyncBatcherExecuteDelegate ExecutorWithFailure - = (IReadOnlyList operations, CancellationToken cancellation) => + = (PartitionKeyRangeServerBatchRequest request, CancellationToken cancellationToken) => { throw expectedException; }; + private BatchAsyncBatcherRetryDelegate Retrier = (BatchAsyncOperationContext operation, CancellationToken cancellation) => + { + return Task.CompletedTask; + }; + [DataTestMethod] [ExpectedException(typeof(ArgumentOutOfRangeException))] [DataRow(0)] [DataRow(-1)] public void ValidatesSize(int size) { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(size, 1, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(size, 1, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); } [DataTestMethod] @@ -116,49 +161,57 @@ public void ValidatesSize(int size) [DataRow(-1)] public void ValidatesByteSize(int size) { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, size, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, size, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); } [TestMethod] [ExpectedException(typeof(ArgumentNullException))] public void ValidatesExecutor() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1, new CosmosJsonDotNetSerializer(), null); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1, new CosmosJsonDotNetSerializer(), null, this.Retrier); + } + + [TestMethod] + [ExpectedException(typeof(ArgumentNullException))] + public void ValidatesRetrier() + { + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1, new CosmosJsonDotNetSerializer(), this.Executor, null); } [TestMethod] [ExpectedException(typeof(ArgumentNullException))] public void ValidatesSerializer() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1, null, this.Executor); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1, null, this.Executor, this.Retrier); } [TestMethod] public void HasFixedSize() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); + Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); } [TestMethod] public async Task HasFixedByteSize() { - await this.ItemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), CancellationToken.None); + ItemBatchOperation itemBatchOperation = this.ItemBatchOperation(); + await itemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), default(CancellationToken)); // Each operation is 2 bytes - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); - Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, itemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, itemBatchOperation))); + Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, itemBatchOperation))); } [TestMethod] public async Task ExceptionsFailOperationsAsync() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); - BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); - BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure, this.Retrier); + BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); + BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); batchAsyncBatcher.TryAdd(context1); batchAsyncBatcher.TryAdd(context2); await batchAsyncBatcher.DispatchAsync(); @@ -172,7 +225,7 @@ public async Task ExceptionsFailOperationsAsync() [TestMethod] public async Task DispatchProcessInOrderAsync() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); List contexts = new List(10); for (int i = 0; i < 10; i++) { @@ -195,8 +248,8 @@ public async Task DispatchProcessInOrderAsync() [TestMethod] public async Task DispatchWithLessResponses() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithLessResponses); - BatchAsyncBatcher secondAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithLessResponses, this.Retrier); + BatchAsyncBatcher secondAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); List contexts = new List(10); for (int i = 0; i < 10; i++) { @@ -246,25 +299,80 @@ public async Task DispatchWithLessResponses() [TestMethod] public void IsEmptyWithNoOperations() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); Assert.IsTrue(batchAsyncBatcher.IsEmpty); } [TestMethod] public void IsNotEmptyWithOperations() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); Assert.IsFalse(batchAsyncBatcher.IsEmpty); } [TestMethod] public async Task CannotAddToDispatchedBatch() { - BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); + Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); + await batchAsyncBatcher.DispatchAsync(); + Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); + } + + [TestMethod] + public async Task RetrierGetsCalledOnSplit() + { + BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); + BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); + + Mock retryDelegate = new Mock(); + + BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithSplit, retryDelegate.Object); + Assert.IsTrue(batchAsyncBatcher.TryAdd(context1)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(context2)); await batchAsyncBatcher.DispatchAsync(); - Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation))); + retryDelegate.Verify(a => a(It.Is(o => o == context1), It.IsAny()), Times.Once); + retryDelegate.Verify(a => a(It.Is(o => o == context2), It.IsAny()), Times.Once); + retryDelegate.Verify(a => a(It.IsAny(), It.IsAny()), Times.Exactly(2)); + } + + [TestMethod] + public async Task RetrierGetsCalledOnOverFlow() + { + BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); + BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); + + Mock retryDelegate = new Mock(); + Mock executeDelegate = new Mock(); + + BatchAsyncBatcherThatOverflows batchAsyncBatcher = new BatchAsyncBatcherThatOverflows(2, 1000, new CosmosJsonDotNetSerializer(), executeDelegate.Object, retryDelegate.Object); + Assert.IsTrue(batchAsyncBatcher.TryAdd(context1)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(context2)); + await batchAsyncBatcher.DispatchAsync(); + retryDelegate.Verify(a => a(It.Is(o => o == context1), It.IsAny()), Times.Never); + retryDelegate.Verify(a => a(It.Is(o => o == context2), It.IsAny()), Times.Once); + retryDelegate.Verify(a => a(It.IsAny(), It.IsAny()), Times.Once); + } + + private class BatchAsyncBatcherThatOverflows : BatchAsyncBatcher + { + public BatchAsyncBatcherThatOverflows( + int maxBatchOperationCount, + int maxBatchByteSize, + CosmosSerializer cosmosSerializer, + BatchAsyncBatcherExecuteDelegate executor, + BatchAsyncBatcherRetryDelegate retrier) : base (maxBatchOperationCount, maxBatchByteSize, cosmosSerializer, executor, retrier) + { + + } + + internal override int GetOverflowOperations( + PartitionKeyRangeServerBatchRequest request, + IReadOnlyList operationsSentToRequest) + { + return 1; + } } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index dbb89dbeca..954e30ca51 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -24,21 +24,21 @@ public class BatchAsyncStreamerTests // Executor just returns a reponse matching the Id with Etag private BatchAsyncBatcherExecuteDelegate Executor - = async (IReadOnlyList operations, CancellationToken cancellation) => + = async (PartitionKeyRangeServerBatchRequest request, CancellationToken cancellationToken) => { List results = new List(); - ItemBatchOperation[] arrayOperations = new ItemBatchOperation[operations.Count]; + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[request.Operations.Count]; int index = 0; - foreach (BatchAsyncOperationContext operation in operations) + foreach (ItemBatchOperation operation in request.Operations) { results.Add( new BatchOperationResult(HttpStatusCode.OK) { ResourceStream = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true), - ETag = operation.Operation.Id + ETag = operation.Id }); - arrayOperations[index++] = operation.Operation; + arrayOperations[index++] = operation; } MemoryStream responseContent = await new BatchResponsePayloadWriter(results).GeneratePayloadAsync(); @@ -46,33 +46,37 @@ private BatchAsyncBatcherExecuteDelegate Executor SinglePartitionKeyServerBatchRequest batchRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync( partitionKey: null, operations: new ArraySegment(arrayOperations), - maxBodyLength: (int)responseContent.Length * operations.Count, - maxOperationCount: operations.Count, + maxBodyLength: (int)responseContent.Length * request.Operations.Count, + maxOperationCount: request.Operations.Count, serializer: new CosmosJsonDotNetSerializer(), - cancellationToken: cancellation); + cancellationToken: cancellationToken); BatchResponse batchresponse = await BatchResponse.PopulateFromContentAsync( new ResponseMessage(HttpStatusCode.OK) { Content = responseContent }, batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(operations.Count, new List { batchresponse }, new CosmosJsonDotNetSerializer()); - return response; + return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, new List() { batchresponse }); }; private BatchAsyncBatcherExecuteDelegate ExecutorWithFailure - = (IReadOnlyList operations, CancellationToken cancellation) => + = async (PartitionKeyRangeServerBatchRequest request, CancellationToken cancellationToken) => { throw expectedException; }; + private BatchAsyncBatcherRetryDelegate Retrier = (BatchAsyncOperationContext operation, CancellationToken cancellation) => + { + return Task.CompletedTask; + }; + [DataTestMethod] [ExpectedException(typeof(ArgumentOutOfRangeException))] [DataRow(0)] [DataRow(-1)] public void ValidatesSize(int size) { - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(size, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(size, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); } [DataTestMethod] @@ -81,27 +85,34 @@ public void ValidatesSize(int size) [DataRow(-1)] public void ValidatesDispatchTimer(int dispatchTimerInSeconds) { - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, dispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, dispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); } [TestMethod] [ExpectedException(typeof(ArgumentNullException))] public void ValidatesExecutor() { - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), null); + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), null, this.Retrier); + } + + [TestMethod] + [ExpectedException(typeof(ArgumentNullException))] + public void ValidatesRetrier() + { + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor, null); } [TestMethod] [ExpectedException(typeof(ArgumentNullException))] public void ValidatesSerializer() { - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, null, this.Executor); + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(1, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, null, this.Executor, this.Retrier); } [TestMethod] public async Task ExceptionsOnBatchBubbleUpAsync() { - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure); + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure, this.Retrier); BatchAsyncOperationContext context = CreateContext(this.ItemBatchOperation); batchAsyncStreamer.Add(context); Exception capturedException = await Assert.ThrowsExceptionAsync(() => context.Task); @@ -112,7 +123,7 @@ public async Task ExceptionsOnBatchBubbleUpAsync() public async Task TimerDispatchesAsync() { // Bigger batch size than the amount of operations, timer should dispatch - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); BatchAsyncOperationContext context = CreateContext(this.ItemBatchOperation); batchAsyncStreamer.Add(context); BatchOperationResult result = await context.Task; @@ -124,7 +135,7 @@ public async Task TimerDispatchesAsync() public async Task DispatchesAsync() { // Expect all operations to complete as their batches get dispached - BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor); + BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); List> contexts = new List>(10); for (int i = 0; i < 10; i++) { From 47184fd2d4096547815749556cf572eaf1cffa4b Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Fri, 16 Aug 2019 11:41:35 -0700 Subject: [PATCH 38/41] Adding interlock checker --- .../src/Util/InterlockIncrementCheck.cs | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 Microsoft.Azure.Cosmos/src/Util/InterlockIncrementCheck.cs diff --git a/Microsoft.Azure.Cosmos/src/Util/InterlockIncrementCheck.cs b/Microsoft.Azure.Cosmos/src/Util/InterlockIncrementCheck.cs new file mode 100644 index 0000000000..fdb2bf50c7 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Util/InterlockIncrementCheck.cs @@ -0,0 +1,56 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos +{ + using System; + using System.Threading; + + /// + /// This class is used to assert that a region of code can only be called concurrently by a limited amount of threads. + /// + internal class InterlockIncrementCheck : IDisposable + { + private readonly int maxConcurrentOperations; + private int counter = 0; + + public InterlockIncrementCheck(int maxConcurrentOperations = 1) + { + if (maxConcurrentOperations < 1) + { + throw new ArgumentOutOfRangeException(nameof(maxConcurrentOperations), "Cannot be lower than 1."); + } + + this.maxConcurrentOperations = maxConcurrentOperations; + } + + /// + /// Increments the internal lock and asserts that only the allowed + /// + /// When more operations than those allowed try to access the context. + /// + /// InterlockIncrementCheck interlockIncrementCheck = new InterlockIncrementCheck(); + /// using (interlockIncrementCheck.EnterLockCheck()) + /// { + /// // protected code + /// } + /// + /// + public InterlockIncrementCheck EnterLockCheck() + { + Interlocked.Increment(ref this.counter); + if (this.counter > this.maxConcurrentOperations) + { + throw new InvalidOperationException($"InterlockIncrementCheck detected {this.counter} with a maximum of {this.maxConcurrentOperations}."); + } + + return this; + } + + public void Dispose() + { + Interlocked.Decrement(ref this.counter); + } + } +} From cc8a483037796cb3899249211ac43e9348262a5e Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Fri, 16 Aug 2019 11:44:42 -0700 Subject: [PATCH 39/41] Refactoring itemoperation and context --- .../src/Batch/BatchAsyncBatcher.cs | 132 ++++++++---------- .../src/Batch/BatchAsyncContainerExecutor.cs | 13 +- .../src/Batch/BatchAsyncStreamer.cs | 4 +- .../src/Batch/ItemBatchOperation.cs | 22 +++ ...ontext.cs => ItemBatchOperationContext.cs} | 20 ++- .../PartitionKeyRangeBatchExecutionResult.cs | 16 +-- .../Batch/PartitionKeyRangeBatchResponse.cs | 65 +++------ .../PartitionKeyRangeServerBatchRequest.cs | 6 +- .../src/Batch/ServerBatchRequest.cs | 5 +- 9 files changed, 135 insertions(+), 148 deletions(-) rename Microsoft.Azure.Cosmos/src/Batch/{BatchAsyncOperationContext.cs => ItemBatchOperationContext.cs} (73%) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 42250f4bbe..08abc40094 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -21,15 +21,16 @@ namespace Microsoft.Azure.Cosmos /// 4. If there was a split detected, all operations in the request, are sent to the for re-queueing. /// 5. The result of the request is used to wire up all responses with the original Tasks for each operation. /// - /// + /// internal class BatchAsyncBatcher { private readonly CosmosSerializer cosmosSerializer; - private readonly List batchOperations; + private readonly List batchOperations; private readonly BatchAsyncBatcherExecuteDelegate executor; private readonly BatchAsyncBatcherRetryDelegate retrier; private readonly int maxBatchByteSize; private readonly int maxBatchOperationCount; + private readonly InterlockIncrementCheck interlockIncrementCheck = new InterlockIncrementCheck(); private long currentSize = 0; private bool dispached = false; @@ -67,7 +68,7 @@ public BatchAsyncBatcher( throw new ArgumentNullException(nameof(cosmosSerializer)); } - this.batchOperations = new List(maxBatchOperationCount); + this.batchOperations = new List(maxBatchOperationCount); this.executor = executor; this.retrier = retrier; this.maxBatchByteSize = maxBatchByteSize; @@ -75,7 +76,7 @@ public BatchAsyncBatcher( this.cosmosSerializer = cosmosSerializer; } - public virtual bool TryAdd(BatchAsyncOperationContext operationContext) + public virtual bool TryAdd(ItemBatchOperation operation) { if (this.dispached) { @@ -83,9 +84,14 @@ public virtual bool TryAdd(BatchAsyncOperationContext operationContext) return false; } - if (operationContext == null) + if (operation == null) { - throw new ArgumentNullException(nameof(operationContext)); + throw new ArgumentNullException(nameof(operation)); + } + + if (operation.Context == null) + { + throw new ArgumentNullException(nameof(operation.Context)); } if (this.batchOperations.Count == this.maxBatchOperationCount) @@ -94,7 +100,7 @@ public virtual bool TryAdd(BatchAsyncOperationContext operationContext) return false; } - int itemByteSize = operationContext.Operation.GetApproximateSerializedLength(); + int itemByteSize = operation.GetApproximateSerializedLength(); if (itemByteSize + this.currentSize > this.maxBatchByteSize) { @@ -105,94 +111,76 @@ public virtual bool TryAdd(BatchAsyncOperationContext operationContext) this.currentSize += itemByteSize; // Operation index is in the scope of the current batch - operationContext.Operation.OperationIndex = this.batchOperations.Count; - operationContext.CurrentBatcher = this; - this.batchOperations.Add(operationContext); + operation.OperationIndex = this.batchOperations.Count; + operation.Context.CurrentBatcher = this; + this.batchOperations.Add(operation); return true; } public virtual async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) { - PartitionKeyRangeServerBatchRequest serverRequest = await this.CreateServerRequestAsync(cancellationToken); - // Any overflow goes to a new batch - int overFlowOperations = this.GetOverflowOperations(serverRequest, this.batchOperations); - while (overFlowOperations > 0) - { - await this.retrier(this.batchOperations[this.batchOperations.Count - overFlowOperations], cancellationToken); - overFlowOperations--; - } - - try + using (this.interlockIncrementCheck.EnterLockCheck()) { - PartitionKeyRangeBatchExecutionResult result = await this.executor(serverRequest, cancellationToken); - - List responses = new List(result.ServerResponses.Count); - if (!result.ContainsSplit()) + // HybridRow serialization might leave some pending operations out of the batch + (PartitionKeyRangeServerBatchRequest serverRequest, ArraySegment pendingOperations) = await this.CreateServerRequestAsync(cancellationToken); + // Any overflow goes to a new batch + foreach (ItemBatchOperation operation in pendingOperations) { - responses.AddRange(result.ServerResponses); + await this.retrier(operation, cancellationToken); } - else + + try { - foreach (ItemBatchOperation operationToRetry in result.Operations) + PartitionKeyRangeBatchExecutionResult result = await this.executor(serverRequest, cancellationToken); + + if (result.IsSplit()) { - await this.retrier(this.batchOperations[operationToRetry.OperationIndex], cancellationToken); + foreach (ItemBatchOperation operationToRetry in result.Operations) + { + await this.retrier(this.batchOperations[operationToRetry.OperationIndex], cancellationToken); + } + + return; } - } - using (PartitionKeyRangeBatchResponse batchResponse = new PartitionKeyRangeBatchResponse(serverRequest.Operations.Count, responses, this.cosmosSerializer)) - { - for (int index = 0; index < this.batchOperations.Count; index++) + using (PartitionKeyRangeBatchResponse batchResponse = new PartitionKeyRangeBatchResponse(serverRequest.Operations.Count, result.ServerResponse, this.cosmosSerializer)) { - BatchAsyncOperationContext context = this.batchOperations[index]; - BatchOperationResult response = batchResponse[context.Operation.OperationIndex]; - if (response != null) + for (int index = 0; index < this.batchOperations.Count; index++) { - context.Complete(this, response); + ItemBatchOperation operation = this.batchOperations[index]; + BatchOperationResult response = batchResponse[operation.OperationIndex]; + if (response != null) + { + operation.Context.Complete(this, response); + } } } } - } - catch (Exception ex) - { - DefaultTrace.TraceError("Exception during BatchAsyncBatcher: {0}", ex); - // Exceptions happening during execution fail all the Tasks part of the request (excluding overflow) - foreach (ItemBatchOperation itemBatchOperation in serverRequest.Operations) + catch (Exception ex) { - BatchAsyncOperationContext context = this.batchOperations[itemBatchOperation.OperationIndex]; - context.Fail(this, ex); + DefaultTrace.TraceError("Exception during BatchAsyncBatcher: {0}", ex); + // Exceptions happening during execution fail all the Tasks part of the request (excluding overflow) + foreach (ItemBatchOperation itemBatchOperation in serverRequest.Operations) + { + ItemBatchOperation operation = this.batchOperations[itemBatchOperation.OperationIndex]; + operation.Context.Fail(this, ex); + } + } + finally + { + this.batchOperations.Clear(); + this.dispached = true; } - } - finally - { - this.batchOperations.Clear(); - this.dispached = true; } } - /// - /// If because of HybridRow serialization overhead, not all operations fit in the request, we send those extra operations in a separate request. - /// - internal virtual int GetOverflowOperations( - PartitionKeyRangeServerBatchRequest request, - IReadOnlyList operationsSentToRequest) - { - int totalOperations = operationsSentToRequest.Count; - return totalOperations - request.Operations.Count; - } - - private async Task CreateServerRequestAsync(CancellationToken cancellationToken) + internal virtual async Task>> CreateServerRequestAsync(CancellationToken cancellationToken) { // All operations should be for the same PKRange - string partitionKeyRangeId = this.batchOperations[0].PartitionKeyRangeId; + string partitionKeyRangeId = this.batchOperations[0].Context.PartitionKeyRangeId; - ItemBatchOperation[] operations = new ItemBatchOperation[this.batchOperations.Count]; - for (int i = 0; i < this.batchOperations.Count; i++) - { - operations[i] = this.batchOperations[i].Operation; - } - - ArraySegment operationsArraySegment = new ArraySegment(operations); - PartitionKeyRangeServerBatchRequest request = await PartitionKeyRangeServerBatchRequest.CreateAsync( + ArraySegment operationsArraySegment = new ArraySegment(this.batchOperations.ToArray()); + return await PartitionKeyRangeServerBatchRequest.CreateAsync( partitionKeyRangeId, operationsArraySegment, this.maxBatchByteSize, @@ -200,8 +188,6 @@ private async Task CreateServerRequestAsync ensureContinuousOperationIndexes: false, serializer: this.cosmosSerializer, cancellationToken: cancellationToken).ConfigureAwait(false); - - return request; } } @@ -215,5 +201,5 @@ private async Task CreateServerRequestAsync /// Delegate to process a request for retry an operation /// /// An instance of . - internal delegate Task BatchAsyncBatcherRetryDelegate(BatchAsyncOperationContext operationContext, CancellationToken cancellationToken); + internal delegate Task BatchAsyncBatcherRetryDelegate(ItemBatchOperation operation, CancellationToken cancellationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs index 6890be7fe3..fe36315a14 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs @@ -86,8 +86,9 @@ public async Task AddAsync( string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false); BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); - BatchAsyncOperationContext context = new BatchAsyncOperationContext(resolvedPartitionKeyRangeId, operation); - streamer.Add(context); + ItemBatchOperationContext context = new ItemBatchOperationContext(resolvedPartitionKeyRangeId); + operation.AttachContext(context); + streamer.Add(operation); return await context.Task; } @@ -169,12 +170,12 @@ private static void AddHeadersToRequestMessage(RequestMessage requestMessage, st } private async Task ReBatchAsync( - BatchAsyncOperationContext context, + ItemBatchOperation operation, CancellationToken cancellationToken) { - string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(context.Operation, cancellationToken).ConfigureAwait(false); + string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false); BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId); - streamer.Add(context); + streamer.Add(operation); } private async Task ResolvePartitionKeyRangeIdAsync( @@ -228,7 +229,7 @@ private async Task ExecuteAsync( BatchResponse serverResponse = await BatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.CosmosSerializer).ConfigureAwait(false); - return new PartitionKeyRangeBatchExecutionResult(serverRequest.PartitionKeyRangeId, serverRequest.Operations, new List() { serverResponse }); + return new PartitionKeyRangeBatchExecutionResult(serverRequest.PartitionKeyRangeId, serverRequest.Operations, serverResponse); } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs index f568a8d5b4..7ccb45c313 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncStreamer.cs @@ -84,12 +84,12 @@ public BatchAsyncStreamer( this.ResetTimer(); } - public void Add(BatchAsyncOperationContext context) + public void Add(ItemBatchOperation operation) { BatchAsyncBatcher toDispatch = null; lock (this.dispatchLimiter) { - while (!this.currentBatcher.TryAdd(context)) + while (!this.currentBatcher.TryAdd(operation)) { // Batcher is full toDispatch = this.GetBatchToDispatchAndCreate(); diff --git a/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs index 05402387fd..fa0967a490 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperation.cs @@ -88,6 +88,14 @@ internal Memory ResourceBody } } + /// + /// Operational context used in stream operations. + /// + /// + /// + /// + internal ItemBatchOperationContext Context { get; private set; } + /// /// Disposes the current . /// @@ -289,6 +297,20 @@ internal virtual async Task MaterializeResourceAsync(CosmosSerializer serializer } } + /// + /// Attached a context to the current operation to track resolution. + /// + /// If the operation already had an attached context. + internal void AttachContext(ItemBatchOperationContext context) + { + if (this.Context != null) + { + throw new InvalidOperationException("Cannot modify the current context of an operation."); + } + + this.Context = context; + } + /// /// Disposes the disposable members held by this class. /// diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperationContext.cs similarity index 73% rename from Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs rename to Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperationContext.cs index fe478a4e86..e1d264fdbc 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncOperationContext.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/ItemBatchOperationContext.cs @@ -11,23 +11,18 @@ namespace Microsoft.Azure.Cosmos /// /// Context for a particular Batch operation. /// - internal class BatchAsyncOperationContext : IDisposable + internal class ItemBatchOperationContext : IDisposable { public string PartitionKeyRangeId { get; } - public ItemBatchOperation Operation { get; } - public BatchAsyncBatcher CurrentBatcher { get; set; } public Task Task => this.taskCompletionSource.Task; private TaskCompletionSource taskCompletionSource = new TaskCompletionSource(); - public BatchAsyncOperationContext( - string partitionKeyRangeId, - ItemBatchOperation operation) + public ItemBatchOperationContext(string partitionKeyRangeId) { - this.Operation = operation; this.PartitionKeyRangeId = partitionKeyRangeId; } @@ -47,7 +42,7 @@ public void Fail( BatchAsyncBatcher completer, Exception exception) { - if (this.AssertBatcher(completer)) + if (this.AssertBatcher(completer, exception)) { this.taskCompletionSource.SetException(exception); } @@ -57,16 +52,17 @@ public void Fail( public void Dispose() { - this.Operation.Dispose(); this.CurrentBatcher = null; } - private bool AssertBatcher(BatchAsyncBatcher completer) + private bool AssertBatcher( + BatchAsyncBatcher completer, + Exception innerException = null) { if (!object.ReferenceEquals(completer, this.CurrentBatcher)) { - DefaultTrace.TraceCritical($"Operation {this.Operation.Id} was completed by incorrect batcher."); - this.taskCompletionSource.SetException(new Exception($"Operation {this.Operation.Id} was completed by incorrect batcher.")); + DefaultTrace.TraceCritical($"Operation was completed by incorrect batcher."); + this.taskCompletionSource.SetException(new Exception($"Operation was completed by incorrect batcher.", innerException)); return false; } diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs index cfe12e1a78..99b971b09f 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchExecutionResult.cs @@ -12,24 +12,24 @@ internal class PartitionKeyRangeBatchExecutionResult { public string PartitionKeyRangeId { get; } - public IReadOnlyList ServerResponses { get; } + public BatchResponse ServerResponse { get; } public IEnumerable Operations { get; } public PartitionKeyRangeBatchExecutionResult( string pkRangeId, IEnumerable operations, - List serverResponses) + BatchResponse serverResponse) { this.PartitionKeyRangeId = pkRangeId; - this.ServerResponses = serverResponses; + this.ServerResponse = serverResponse; this.Operations = operations; } - internal bool ContainsSplit() => this.ServerResponses != null && this.ServerResponses.Any(serverResponse => - serverResponse.StatusCode == HttpStatusCode.Gone - && (serverResponse.SubStatusCode == Documents.SubStatusCodes.CompletingSplit - || serverResponse.SubStatusCode == Documents.SubStatusCodes.CompletingPartitionMigration - || serverResponse.SubStatusCode == Documents.SubStatusCodes.PartitionKeyRangeGone)); + internal bool IsSplit() => this.ServerResponse != null && + this.ServerResponse.StatusCode == HttpStatusCode.Gone + && (this.ServerResponse.SubStatusCode == Documents.SubStatusCodes.CompletingSplit + || this.ServerResponse.SubStatusCode == Documents.SubStatusCodes.CompletingPartitionMigration + || this.ServerResponse.SubStatusCode == Documents.SubStatusCodes.PartitionKeyRangeGone); } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchResponse.cs index af277b05d3..c2ff0fc1ab 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchResponse.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchResponse.cs @@ -6,7 +6,6 @@ namespace Microsoft.Azure.Cosmos { using System; using System.Collections.Generic; - using System.Linq; using System.Net; using System.Text; using Microsoft.Azure.Documents; @@ -41,67 +40,55 @@ internal PartitionKeyRangeBatchResponse( /// Initializes a new instance of the class. /// /// Original operations that generated the server responses. - /// Responses from the server. + /// Response from the server. /// Serializer to deserialize response resource body streams. internal PartitionKeyRangeBatchResponse( int originalOperationsCount, - IEnumerable serverResponses, + BatchResponse serverResponse, CosmosSerializer serializer) { - this.StatusCode = serverResponses.Any(r => r.StatusCode != HttpStatusCode.OK) - ? (HttpStatusCode)StatusCodes.MultiStatus - : HttpStatusCode.OK; + this.StatusCode = serverResponse.StatusCode; - this.ServerResponses = serverResponses; + this.ServerResponse = serverResponse; this.resultsByOperationIndex = new BatchOperationResult[originalOperationsCount]; StringBuilder errorMessageBuilder = new StringBuilder(); List activityIds = new List(); List itemBatchOperations = new List(); - foreach (BatchResponse serverResponse in serverResponses) + // We expect number of results == number of operations here + for (int index = 0; index < serverResponse.Operations.Count; index++) { - // We expect number of results == number of operations here - for (int index = 0; index < serverResponse.Operations.Count; index++) + int operationIndex = serverResponse.Operations[index].OperationIndex; + if (this.resultsByOperationIndex[operationIndex] == null + || this.resultsByOperationIndex[operationIndex].StatusCode == (HttpStatusCode)StatusCodes.TooManyRequests) { - int operationIndex = serverResponse.Operations[index].OperationIndex; - if (this.resultsByOperationIndex[operationIndex] == null - || this.resultsByOperationIndex[operationIndex].StatusCode == (HttpStatusCode)StatusCodes.TooManyRequests) - { - this.resultsByOperationIndex[operationIndex] = serverResponse[index]; - } + this.resultsByOperationIndex[operationIndex] = serverResponse[index]; } + } - itemBatchOperations.AddRange(serverResponse.Operations); - this.RequestCharge += serverResponse.RequestCharge; + itemBatchOperations.AddRange(serverResponse.Operations); + this.RequestCharge += serverResponse.RequestCharge; - if (!string.IsNullOrEmpty(serverResponse.ErrorMessage)) - { - errorMessageBuilder.AppendFormat("{0}; ", serverResponse.ErrorMessage); - } - - activityIds.Add(serverResponse.ActivityId); + if (!string.IsNullOrEmpty(serverResponse.ErrorMessage)) + { + errorMessageBuilder.AppendFormat("{0}; ", serverResponse.ErrorMessage); } - this.ActivityIds = activityIds; + this.ActivityId = serverResponse.ActivityId; this.ErrorMessage = errorMessageBuilder.Length > 2 ? errorMessageBuilder.ToString(0, errorMessageBuilder.Length - 2) : null; this.Operations = itemBatchOperations; this.Serializer = serializer; } /// - /// Gets the ActivityIds that identify the server requests made to execute the batch request. + /// Gets the ActivityId that identifies the server request made to execute the batch request. /// -#pragma warning disable CA1721 // Property names should not match get methods - public virtual IEnumerable ActivityIds { get; } -#pragma warning restore CA1721 // Property names should not match get methods - - /// - public override string ActivityId => this.ActivityIds.First(); + public override string ActivityId { get; } internal override CosmosSerializer Serializer { get; } // for unit testing only - internal IEnumerable ServerResponses { get; private set; } + internal BatchResponse ServerResponse { get; private set; } /// /// Gets the number of operation results. @@ -149,7 +136,7 @@ public override IEnumerator GetEnumerator() internal override IEnumerable GetActivityIds() { - return this.ActivityIds; + return new string[1] { this.ActivityId }; } /// @@ -161,15 +148,7 @@ protected override void Dispose(bool disposing) if (disposing && !this.isDisposed) { this.isDisposed = true; - if (this.ServerResponses != null) - { - foreach (BatchResponse response in this.ServerResponses) - { - response.Dispose(); - } - - this.ServerResponses = null; - } + this.ServerResponse?.Dispose(); } base.Dispose(disposing); diff --git a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeServerBatchRequest.cs b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeServerBatchRequest.cs index 24a0adfda9..c4719d09d6 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeServerBatchRequest.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeServerBatchRequest.cs @@ -45,7 +45,7 @@ public PartitionKeyRangeServerBatchRequest( /// Serializer to serialize user provided objects to JSON. /// representing request cancellation. /// A newly created instance of . - public static async Task CreateAsync( + public static async Task>> CreateAsync( string partitionKeyRangeId, ArraySegment operations, int maxBodyLength, @@ -55,8 +55,8 @@ public static async Task CreateAsync( CancellationToken cancellationToken) { PartitionKeyRangeServerBatchRequest request = new PartitionKeyRangeServerBatchRequest(partitionKeyRangeId, maxBodyLength, maxOperationCount, serializer); - await request.CreateBodyStreamAsync(operations, cancellationToken, ensureContinuousOperationIndexes); - return request; + ArraySegment pendingOperations = await request.CreateBodyStreamAsync(operations, cancellationToken, ensureContinuousOperationIndexes); + return new Tuple>(request, pendingOperations); } } } diff --git a/Microsoft.Azure.Cosmos/src/Batch/ServerBatchRequest.cs b/Microsoft.Azure.Cosmos/src/Batch/ServerBatchRequest.cs index 29ebc57e06..dd4e8818d1 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/ServerBatchRequest.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/ServerBatchRequest.cs @@ -70,7 +70,8 @@ public MemoryStream TransferBodyStream() /// Operations to be added; read-only. /// representing request cancellation. /// Whether to stop adding operations to the request once there is non-continuity in the operation indexes. - protected async Task CreateBodyStreamAsync( + /// Any pending operations that were not included in the request. + protected async Task> CreateBodyStreamAsync( ArraySegment operations, CancellationToken cancellationToken, bool ensureContinuousOperationIndexes = false) @@ -132,6 +133,8 @@ protected async Task CreateBodyStreamAsync( { throw new RequestEntityTooLargeException(RMResources.RequestTooLarge); } + + return new ArraySegment(operations.Array, materializedCount, operations.Count - materializedCount); } private Result WriteOperation(long index, out ReadOnlyMemory buffer) From 974eaef2dcc8d7b7ea5c02daab7416012c5a2873 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Fri, 16 Aug 2019 11:44:54 -0700 Subject: [PATCH 40/41] New tests --- .../Batch/BatchAsyncContainerExecutorTests.cs | 1 - .../Batch/BatchAsyncBatcherTests.cs | 131 ++++++++++-------- .../Batch/BatchAsyncOperationContextTests.cs | 24 +++- .../Batch/BatchAsyncStreamerTests.cs | 24 ++-- .../Batch/PartitionKeyBatchResponseTests.cs | 2 +- ...titionKeyRangeBatchExecutionResultTests.cs | 6 +- ...artitionKeyRangeServerBatchRequestTests.cs | 63 +++++++++ .../InterlockIncrementCheckTests.cs | 65 +++++++++ 8 files changed, 242 insertions(+), 74 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeServerBatchRequestTests.cs create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InterlockIncrementCheckTests.cs diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs index 393adc41f7..3fbee77ab8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchAsyncContainerExecutorTests.cs @@ -11,7 +11,6 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using System.Threading.Tasks; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; - using Newtonsoft.Json; [TestClass] #pragma warning disable CA1001 // Types that own disposable fields should be disposable diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 60be012d1a..4244c84f61 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -20,7 +20,15 @@ public class BatchAsyncBatcherTests { private static Exception expectedException = new Exception(); - private ItemBatchOperation ItemBatchOperation() => new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); + private ItemBatchOperation CreateItemBatchOperation(bool withContext = false) { + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0, string.Empty, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); + if (withContext) + { + operation.AttachContext(new ItemBatchOperationContext(string.Empty)); + } + + return operation; + } private BatchAsyncBatcherExecuteDelegate Executor = async (PartitionKeyRangeServerBatchRequest request, CancellationToken cancellationToken) => @@ -55,7 +63,7 @@ private BatchAsyncBatcherExecuteDelegate Executor batchRequest, new CosmosJsonDotNetSerializer()); - return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, new List() { batchresponse }); + return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, batchresponse); }; private BatchAsyncBatcherExecuteDelegate ExecutorWithSplit @@ -94,7 +102,7 @@ private BatchAsyncBatcherExecuteDelegate ExecutorWithSplit batchRequest, new CosmosJsonDotNetSerializer()); - return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, new List() { batchresponse }); + return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, batchresponse); }; // The response will include all but 2 operation responses @@ -132,7 +140,7 @@ private BatchAsyncBatcherExecuteDelegate ExecutorWithLessResponses batchRequest, new CosmosJsonDotNetSerializer()); - return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, new List() { batchresponse }); + return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, batchresponse); }; private BatchAsyncBatcherExecuteDelegate ExecutorWithFailure @@ -141,7 +149,7 @@ private BatchAsyncBatcherExecuteDelegate ExecutorWithFailure throw expectedException; }; - private BatchAsyncBatcherRetryDelegate Retrier = (BatchAsyncOperationContext operation, CancellationToken cancellation) => + private BatchAsyncBatcherRetryDelegate Retrier = (ItemBatchOperation operation, CancellationToken cancellation) => { return Task.CompletedTask; }; @@ -189,31 +197,35 @@ public void ValidatesSerializer() public void HasFixedSize() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); - Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(this.CreateItemBatchOperation(true))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(this.CreateItemBatchOperation(true))); + Assert.IsFalse(batchAsyncBatcher.TryAdd(this.CreateItemBatchOperation(true))); } [TestMethod] public async Task HasFixedByteSize() { - ItemBatchOperation itemBatchOperation = this.ItemBatchOperation(); + ItemBatchOperation itemBatchOperation = this.CreateItemBatchOperation(true); await itemBatchOperation.MaterializeResourceAsync(new CosmosJsonDotNetSerializer(), default(CancellationToken)); // Each operation is 2 bytes BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(3, 4, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, itemBatchOperation))); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, itemBatchOperation))); - Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, itemBatchOperation))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(itemBatchOperation)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(itemBatchOperation)); + Assert.IsFalse(batchAsyncBatcher.TryAdd(itemBatchOperation)); } [TestMethod] public async Task ExceptionsFailOperationsAsync() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure, this.Retrier); - BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); - BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); - batchAsyncBatcher.TryAdd(context1); - batchAsyncBatcher.TryAdd(context2); + ItemBatchOperation operation1 = this.CreateItemBatchOperation(); + ItemBatchOperation operation2 = this.CreateItemBatchOperation(); + ItemBatchOperationContext context1 = new ItemBatchOperationContext(string.Empty); + operation1.AttachContext(context1); + ItemBatchOperationContext context2 = new ItemBatchOperationContext(string.Empty); + operation2.AttachContext(context2); + batchAsyncBatcher.TryAdd(operation1); + batchAsyncBatcher.TryAdd(operation2); await batchAsyncBatcher.DispatchAsync(); Assert.AreEqual(TaskStatus.Faulted, context1.Task.Status); @@ -226,19 +238,21 @@ public async Task ExceptionsFailOperationsAsync() public async Task DispatchProcessInOrderAsync() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); - List contexts = new List(10); + List contexts = new List(10); for (int i = 0; i < 10; i++) { - BatchAsyncOperationContext context = new BatchAsyncOperationContext(string.Empty, new ItemBatchOperation(OperationType.Create, i, i.ToString())); + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, i, i.ToString()); + ItemBatchOperationContext context = new ItemBatchOperationContext(string.Empty); + operation.AttachContext(context); contexts.Add(context); - Assert.IsTrue(batchAsyncBatcher.TryAdd(context)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(operation)); } await batchAsyncBatcher.DispatchAsync(); for (int i = 0; i < 10; i++) { - BatchAsyncOperationContext context = contexts[i]; + ItemBatchOperationContext context = contexts[i]; Assert.AreEqual(TaskStatus.RanToCompletion, context.Task.Status); BatchOperationResult result = await context.Task; Assert.AreEqual(i.ToString(), result.ETag); @@ -250,12 +264,14 @@ public async Task DispatchWithLessResponses() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithLessResponses, this.Retrier); BatchAsyncBatcher secondAsyncBatcher = new BatchAsyncBatcher(10, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); - List contexts = new List(10); + List operations = new List(10); for (int i = 0; i < 10; i++) { - BatchAsyncOperationContext context = new BatchAsyncOperationContext(string.Empty, new ItemBatchOperation(OperationType.Create, i, i.ToString())); - contexts.Add(context); - Assert.IsTrue(batchAsyncBatcher.TryAdd(context)); + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, i, i.ToString()); + ItemBatchOperationContext context = new ItemBatchOperationContext(string.Empty); + operation.AttachContext(context); + operations.Add(operation); + Assert.IsTrue(batchAsyncBatcher.TryAdd(operation)); } await batchAsyncBatcher.DispatchAsync(); @@ -263,25 +279,25 @@ public async Task DispatchWithLessResponses() // Responses 1 and 10 should be missing for (int i = 0; i < 10; i++) { - BatchAsyncOperationContext context = contexts[i]; + ItemBatchOperation operation = operations[i]; // Some tasks should not be resolved if(i == 0 || i == 9) { - Assert.IsTrue(context.Task.Status == TaskStatus.WaitingForActivation); + Assert.IsTrue(operation.Context.Task.Status == TaskStatus.WaitingForActivation); } else { - Assert.IsTrue(context.Task.Status == TaskStatus.RanToCompletion); + Assert.IsTrue(operation.Context.Task.Status == TaskStatus.RanToCompletion); } - if (context.Task.Status == TaskStatus.RanToCompletion) + if (operation.Context.Task.Status == TaskStatus.RanToCompletion) { - BatchOperationResult result = await context.Task; + BatchOperationResult result = await operation.Context.Task; Assert.AreEqual(i.ToString(), result.ETag); } else { // Pass the pending one to another batcher - Assert.IsTrue(secondAsyncBatcher.TryAdd(context)); + Assert.IsTrue(secondAsyncBatcher.TryAdd(operation)); } } @@ -289,9 +305,9 @@ public async Task DispatchWithLessResponses() // All tasks should be completed for (int i = 0; i < 10; i++) { - BatchAsyncOperationContext context = contexts[i]; - Assert.AreEqual(TaskStatus.RanToCompletion, context.Task.Status); - BatchOperationResult result = await context.Task; + ItemBatchOperation operation = operations[i]; + Assert.AreEqual(TaskStatus.RanToCompletion, operation.Context.Task.Status); + BatchOperationResult result = await operation.Context.Task; Assert.AreEqual(i.ToString(), result.ETag); } } @@ -307,7 +323,7 @@ public void IsEmptyWithNoOperations() public void IsNotEmptyWithOperations() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); + Assert.IsTrue(batchAsyncBatcher.TryAdd(this.CreateItemBatchOperation(true))); Assert.IsFalse(batchAsyncBatcher.IsEmpty); } @@ -315,44 +331,50 @@ public void IsNotEmptyWithOperations() public async Task CannotAddToDispatchedBatch() { BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(1, 1000, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); - Assert.IsTrue(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); + ItemBatchOperation operation = this.CreateItemBatchOperation(); + operation.AttachContext(new ItemBatchOperationContext(string.Empty)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(operation)); await batchAsyncBatcher.DispatchAsync(); - Assert.IsFalse(batchAsyncBatcher.TryAdd(new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()))); + Assert.IsFalse(batchAsyncBatcher.TryAdd(this.CreateItemBatchOperation())); } [TestMethod] public async Task RetrierGetsCalledOnSplit() { - BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); - BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); + ItemBatchOperation operation1 = this.CreateItemBatchOperation(); + ItemBatchOperation operation2 = this.CreateItemBatchOperation(); + operation1.AttachContext(new ItemBatchOperationContext(string.Empty)); + operation2.AttachContext(new ItemBatchOperationContext(string.Empty)); Mock retryDelegate = new Mock(); BatchAsyncBatcher batchAsyncBatcher = new BatchAsyncBatcher(2, 1000, new CosmosJsonDotNetSerializer(), this.ExecutorWithSplit, retryDelegate.Object); - Assert.IsTrue(batchAsyncBatcher.TryAdd(context1)); - Assert.IsTrue(batchAsyncBatcher.TryAdd(context2)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(operation1)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(operation2)); await batchAsyncBatcher.DispatchAsync(); - retryDelegate.Verify(a => a(It.Is(o => o == context1), It.IsAny()), Times.Once); - retryDelegate.Verify(a => a(It.Is(o => o == context2), It.IsAny()), Times.Once); - retryDelegate.Verify(a => a(It.IsAny(), It.IsAny()), Times.Exactly(2)); + retryDelegate.Verify(a => a(It.Is(o => o == operation1), It.IsAny()), Times.Once); + retryDelegate.Verify(a => a(It.Is(o => o == operation2), It.IsAny()), Times.Once); + retryDelegate.Verify(a => a(It.IsAny(), It.IsAny()), Times.Exactly(2)); } [TestMethod] public async Task RetrierGetsCalledOnOverFlow() { - BatchAsyncOperationContext context1 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); - BatchAsyncOperationContext context2 = new BatchAsyncOperationContext(string.Empty, this.ItemBatchOperation()); + ItemBatchOperation operation1 = this.CreateItemBatchOperation(); + ItemBatchOperation operation2 = this.CreateItemBatchOperation(); + operation1.AttachContext(new ItemBatchOperationContext(string.Empty)); + operation2.AttachContext(new ItemBatchOperationContext(string.Empty)); Mock retryDelegate = new Mock(); Mock executeDelegate = new Mock(); BatchAsyncBatcherThatOverflows batchAsyncBatcher = new BatchAsyncBatcherThatOverflows(2, 1000, new CosmosJsonDotNetSerializer(), executeDelegate.Object, retryDelegate.Object); - Assert.IsTrue(batchAsyncBatcher.TryAdd(context1)); - Assert.IsTrue(batchAsyncBatcher.TryAdd(context2)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(operation1)); + Assert.IsTrue(batchAsyncBatcher.TryAdd(operation2)); await batchAsyncBatcher.DispatchAsync(); - retryDelegate.Verify(a => a(It.Is(o => o == context1), It.IsAny()), Times.Never); - retryDelegate.Verify(a => a(It.Is(o => o == context2), It.IsAny()), Times.Once); - retryDelegate.Verify(a => a(It.IsAny(), It.IsAny()), Times.Once); + retryDelegate.Verify(a => a(It.Is(o => o == operation1), It.IsAny()), Times.Never); + retryDelegate.Verify(a => a(It.Is(o => o == operation2), It.IsAny()), Times.Once); + retryDelegate.Verify(a => a(It.IsAny(), It.IsAny()), Times.Once); } private class BatchAsyncBatcherThatOverflows : BatchAsyncBatcher @@ -367,11 +389,12 @@ public BatchAsyncBatcherThatOverflows( } - internal override int GetOverflowOperations( - PartitionKeyRangeServerBatchRequest request, - IReadOnlyList operationsSentToRequest) + internal override async Task>> CreateServerRequestAsync(CancellationToken cancellationToken) { - return 1; + (PartitionKeyRangeServerBatchRequest serverRequest, ArraySegment pendingOperations) = await base.CreateServerRequestAsync(cancellationToken); + + // Returning a pending operation to retry + return new Tuple>(serverRequest, new ArraySegment(serverRequest.Operations.ToArray(), 1, 1)); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs index 867c3d0917..d060ee7fe8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncOperationContextTests.cs @@ -18,10 +18,11 @@ public void PartitionKeyRangeIdIsSetOnInitialization() { string expectedPkRangeId = Guid.NewGuid().ToString(); ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); - BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(expectedPkRangeId, operation); + ItemBatchOperationContext batchAsyncOperationContext = new ItemBatchOperationContext(expectedPkRangeId); + operation.AttachContext(batchAsyncOperationContext); Assert.IsNotNull(batchAsyncOperationContext.Task); - Assert.AreEqual(operation, batchAsyncOperationContext.Operation); + Assert.AreEqual(batchAsyncOperationContext, operation.Context); Assert.AreEqual(expectedPkRangeId, batchAsyncOperationContext.PartitionKeyRangeId); Assert.AreEqual(TaskStatus.WaitingForActivation, batchAsyncOperationContext.Task.Status); } @@ -30,10 +31,11 @@ public void PartitionKeyRangeIdIsSetOnInitialization() public void TaskIsCreatedOnInitialization() { ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); - BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(string.Empty, operation); + ItemBatchOperationContext batchAsyncOperationContext = new ItemBatchOperationContext(string.Empty); + operation.AttachContext(batchAsyncOperationContext); Assert.IsNotNull(batchAsyncOperationContext.Task); - Assert.AreEqual(operation, batchAsyncOperationContext.Operation); + Assert.AreEqual(batchAsyncOperationContext, operation.Context); Assert.AreEqual(TaskStatus.WaitingForActivation, batchAsyncOperationContext.Task.Status); } @@ -41,7 +43,8 @@ public void TaskIsCreatedOnInitialization() public async Task TaskResultIsSetOnCompleteAsync() { ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); - BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(string.Empty, operation); + ItemBatchOperationContext batchAsyncOperationContext = new ItemBatchOperationContext(string.Empty); + operation.AttachContext(batchAsyncOperationContext); BatchOperationResult expected = new BatchOperationResult(HttpStatusCode.OK); @@ -56,7 +59,8 @@ public async Task ExceptionIsSetOnFailAsync() { Exception failure = new Exception("It failed"); ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); - BatchAsyncOperationContext batchAsyncOperationContext = new BatchAsyncOperationContext(string.Empty, operation); + ItemBatchOperationContext batchAsyncOperationContext = new ItemBatchOperationContext(string.Empty); + operation.AttachContext(batchAsyncOperationContext); batchAsyncOperationContext.Fail(null, failure); @@ -64,5 +68,13 @@ public async Task ExceptionIsSetOnFailAsync() Assert.AreEqual(failure, capturedException); Assert.AreEqual(TaskStatus.Faulted, batchAsyncOperationContext.Task.Status); } + + [TestMethod] + public void CannotAttachMoreThanOnce() + { + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, 0); + operation.AttachContext(new ItemBatchOperationContext(string.Empty)); + Assert.ThrowsException(() => operation.AttachContext(new ItemBatchOperationContext(string.Empty))); + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 954e30ca51..7680b7f6e5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -56,7 +56,7 @@ private BatchAsyncBatcherExecuteDelegate Executor batchRequest, new CosmosJsonDotNetSerializer()); - return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, new List() { batchresponse }); + return new PartitionKeyRangeBatchExecutionResult(request.PartitionKeyRangeId, request.Operations, batchresponse); }; private BatchAsyncBatcherExecuteDelegate ExecutorWithFailure @@ -65,7 +65,7 @@ private BatchAsyncBatcherExecuteDelegate ExecutorWithFailure throw expectedException; }; - private BatchAsyncBatcherRetryDelegate Retrier = (BatchAsyncOperationContext operation, CancellationToken cancellation) => + private BatchAsyncBatcherRetryDelegate Retrier = (ItemBatchOperation operation, CancellationToken cancellation) => { return Task.CompletedTask; }; @@ -113,8 +113,8 @@ public void ValidatesSerializer() public async Task ExceptionsOnBatchBubbleUpAsync() { BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.ExecutorWithFailure, this.Retrier); - BatchAsyncOperationContext context = CreateContext(this.ItemBatchOperation); - batchAsyncStreamer.Add(context); + ItemBatchOperationContext context = AttachContext(this.ItemBatchOperation); + batchAsyncStreamer.Add(this.ItemBatchOperation); Exception capturedException = await Assert.ThrowsExceptionAsync(() => context.Task); Assert.AreEqual(expectedException, capturedException); } @@ -124,8 +124,8 @@ public async Task TimerDispatchesAsync() { // Bigger batch size than the amount of operations, timer should dispatch BatchAsyncStreamer batchAsyncStreamer = new BatchAsyncStreamer(2, MaxBatchByteSize, DispatchTimerInSeconds, this.TimerPool, new CosmosJsonDotNetSerializer(), this.Executor, this.Retrier); - BatchAsyncOperationContext context = CreateContext(this.ItemBatchOperation); - batchAsyncStreamer.Add(context); + ItemBatchOperationContext context = AttachContext(this.ItemBatchOperation); + batchAsyncStreamer.Add(this.ItemBatchOperation); BatchOperationResult result = await context.Task; Assert.AreEqual(this.ItemBatchOperation.Id, result.ETag); @@ -139,8 +139,9 @@ public async Task DispatchesAsync() List> contexts = new List>(10); for (int i = 0; i < 10; i++) { - BatchAsyncOperationContext context = CreateContext(new ItemBatchOperation(OperationType.Create, i, i.ToString())); - batchAsyncStreamer.Add(context); + ItemBatchOperation operation = new ItemBatchOperation(OperationType.Create, i, i.ToString()); + ItemBatchOperationContext context = AttachContext(operation); + batchAsyncStreamer.Add(operation); contexts.Add(context.Task); } @@ -155,6 +156,11 @@ public async Task DispatchesAsync() } } - private static BatchAsyncOperationContext CreateContext(ItemBatchOperation operation) => new BatchAsyncOperationContext(string.Empty, operation); + private static ItemBatchOperationContext AttachContext(ItemBatchOperation operation) + { + ItemBatchOperationContext context = new ItemBatchOperationContext(string.Empty); + operation.AttachContext(context); + return context; + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs index 5e48bbf41f..3deb8b54e4 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyBatchResponseTests.cs @@ -58,7 +58,7 @@ public async Task StatusCodesAreSetThroughResponseAsync() batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(arrayOperations.Length, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(arrayOperations.Length, batchresponse, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs index 17f82d9eb2..ca67af38d2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeBatchExecutionResultTests.cs @@ -63,7 +63,7 @@ public async Task StatusCodesAreSetThroughResponseAsync() batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(arrayOperations.Length, new List { batchresponse }, new CosmosJsonDotNetSerializer()); + PartitionKeyRangeBatchResponse response = new PartitionKeyRangeBatchResponse(arrayOperations.Length, batchresponse, new CosmosJsonDotNetSerializer()); Assert.AreEqual(HttpStatusCode.OK, response.StatusCode); } @@ -101,9 +101,9 @@ private async Task ConstainsSplitIsTrueInternal(HttpStatusCode statusCode, batchRequest, new CosmosJsonDotNetSerializer()); - PartitionKeyRangeBatchExecutionResult result = new PartitionKeyRangeBatchExecutionResult("0", arrayOperations, new List { batchresponse }); + PartitionKeyRangeBatchExecutionResult result = new PartitionKeyRangeBatchExecutionResult("0", arrayOperations, batchresponse); - return result.ContainsSplit(); + return result.IsSplit(); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeServerBatchRequestTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeServerBatchRequestTests.cs new file mode 100644 index 0000000000..84ced2be2e --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/PartitionKeyRangeServerBatchRequestTests.cs @@ -0,0 +1,63 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Collections.Generic; + using System.IO; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class PartitionKeyRangeServerBatchRequestTests + { + private static ItemBatchOperation CreateItemBatchOperation(string id = "") + { + return new ItemBatchOperation(OperationType.Create, 0, id, new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true)); + } + + [TestMethod] + public async Task FitsAllOperations() + { + List operations = new List() + { + CreateItemBatchOperation(), + CreateItemBatchOperation() + }; + + (PartitionKeyRangeServerBatchRequest request , ArraySegment pendingOperations) = await PartitionKeyRangeServerBatchRequest.CreateAsync("0", new ArraySegment(operations.ToArray()), 200000, 2, false, new CosmosJsonDotNetSerializer(), default(CancellationToken)); + + Assert.AreEqual(operations.Count, request.Operations.Count); + CollectionAssert.AreEqual(operations, request.Operations.ToArray()); + Assert.AreEqual(0, pendingOperations.Count); + } + + /// + /// Verifies that the pending operations contain items that did not fit on the request + /// + [TestMethod] + public async Task OverflowsBasedOnCount() + { + List operations = new List() + { + CreateItemBatchOperation("1"), + CreateItemBatchOperation("2"), + CreateItemBatchOperation("3") + }; + + // Setting max count to 1 + (PartitionKeyRangeServerBatchRequest request, ArraySegment pendingOperations) = await PartitionKeyRangeServerBatchRequest.CreateAsync("0", new ArraySegment(operations.ToArray()), 200000, 1, false, new CosmosJsonDotNetSerializer(), default(CancellationToken)); + + Assert.AreEqual(1, request.Operations.Count); + Assert.AreEqual(operations[0].Id, request.Operations[0].Id); + Assert.AreEqual(2, pendingOperations.Count); + Assert.AreEqual(operations[1].Id, pendingOperations[0].Id); + Assert.AreEqual(operations[2].Id, pendingOperations[1].Id); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InterlockIncrementCheckTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InterlockIncrementCheckTests.cs new file mode 100644 index 0000000000..2c9535c2e4 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InterlockIncrementCheckTests.cs @@ -0,0 +1,65 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ +namespace Microsoft.Azure.Cosmos.Tests +{ + using System; + using System.Collections.Generic; + using System.Threading.Tasks; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class InterlockIncrementCheckTests + { + [DataRow(0)] + [DataRow(-1)] + [DataTestMethod] + [ExpectedException(typeof(ArgumentOutOfRangeException))] + public void ValidatesMaxConcurrentOperations(int maxConcurrentOperations) + { + new InterlockIncrementCheck(maxConcurrentOperations); + } + + [TestMethod] + public void AllowsMultipleChainedOperations() + { + InterlockIncrementCheck check = new InterlockIncrementCheck(); + using (check.EnterLockCheck()) + { + } + + using (check.EnterLockCheck()) + { + } + } + + [TestMethod] + public async Task AllowsMultipleConcurrentOperations() + { + InterlockIncrementCheck check = new InterlockIncrementCheck(2); + List tasks = new List(2); + tasks.Add(this.RunLock(check)); + tasks.Add(this.RunLock(check)); + await Task.WhenAll(tasks); + } + + [TestMethod] + [ExpectedException(typeof(InvalidOperationException))] + public async Task ThrowsOnMultipleConcurrentOperations() + { + InterlockIncrementCheck check = new InterlockIncrementCheck(1); + List tasks = new List(2); + tasks.Add(this.RunLock(check)); + tasks.Add(this.RunLock(check)); + await Task.WhenAll(tasks); + } + + private async Task RunLock(InterlockIncrementCheck check) + { + using (check.EnterLockCheck()) + { + await Task.Delay(500); + } + } + } +} From 3a6ce2738484fa3f62aa4445cd746c0a3631a3c9 Mon Sep 17 00:00:00 2001 From: Matias Quaranta Date: Tue, 20 Aug 2019 13:03:26 -0700 Subject: [PATCH 41/41] Addressing comments --- .../src/Batch/BatchAsyncBatcher.cs | 69 ++++++++++++------- .../src/Util/InterlockIncrementCheck.cs | 11 +-- .../InterlockIncrementCheckTests.cs | 19 +---- 3 files changed, 49 insertions(+), 50 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs index 08abc40094..e595c16404 100644 --- a/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs +++ b/Microsoft.Azure.Cosmos/src/Batch/BatchAsyncBatcher.cs @@ -96,7 +96,7 @@ public virtual bool TryAdd(ItemBatchOperation operation) if (this.batchOperations.Count == this.maxBatchOperationCount) { - DefaultTrace.TraceVerbose($"Batch is full - Max operation count {this.maxBatchOperationCount} reached."); + DefaultTrace.TraceInformation($"Batch is full - Max operation count {this.maxBatchOperationCount} reached."); return false; } @@ -104,7 +104,7 @@ public virtual bool TryAdd(ItemBatchOperation operation) if (itemByteSize + this.currentSize > this.maxBatchByteSize) { - DefaultTrace.TraceVerbose($"Batch is full - Max byte size {this.maxBatchByteSize} reached."); + DefaultTrace.TraceInformation($"Batch is full - Max byte size {this.maxBatchByteSize} reached."); return false; } @@ -119,14 +119,34 @@ public virtual bool TryAdd(ItemBatchOperation operation) public virtual async Task DispatchAsync(CancellationToken cancellationToken = default(CancellationToken)) { - using (this.interlockIncrementCheck.EnterLockCheck()) + this.interlockIncrementCheck.EnterLockCheck(); + + PartitionKeyRangeServerBatchRequest serverRequest = null; + ArraySegment pendingOperations; + + try { - // HybridRow serialization might leave some pending operations out of the batch - (PartitionKeyRangeServerBatchRequest serverRequest, ArraySegment pendingOperations) = await this.CreateServerRequestAsync(cancellationToken); - // Any overflow goes to a new batch - foreach (ItemBatchOperation operation in pendingOperations) + try { - await this.retrier(operation, cancellationToken); + // HybridRow serialization might leave some pending operations out of the batch + Tuple> createRequestResponse = await this.CreateServerRequestAsync(cancellationToken); + serverRequest = createRequestResponse.Item1; + pendingOperations = createRequestResponse.Item2; + // Any overflow goes to a new batch + foreach (ItemBatchOperation operation in pendingOperations) + { + await this.retrier(operation, cancellationToken); + } + } + catch (Exception ex) + { + // Exceptions happening during request creation, fail the entire list + foreach (ItemBatchOperation itemBatchOperation in this.batchOperations) + { + itemBatchOperation.Context.Fail(this, ex); + } + + throw; } try @@ -137,7 +157,7 @@ public virtual bool TryAdd(ItemBatchOperation operation) { foreach (ItemBatchOperation operationToRetry in result.Operations) { - await this.retrier(this.batchOperations[operationToRetry.OperationIndex], cancellationToken); + await this.retrier(operationToRetry, cancellationToken); } return; @@ -145,32 +165,33 @@ public virtual bool TryAdd(ItemBatchOperation operation) using (PartitionKeyRangeBatchResponse batchResponse = new PartitionKeyRangeBatchResponse(serverRequest.Operations.Count, result.ServerResponse, this.cosmosSerializer)) { - for (int index = 0; index < this.batchOperations.Count; index++) + foreach (ItemBatchOperation itemBatchOperation in batchResponse.Operations) { - ItemBatchOperation operation = this.batchOperations[index]; - BatchOperationResult response = batchResponse[operation.OperationIndex]; - if (response != null) - { - operation.Context.Complete(this, response); - } + BatchOperationResult response = batchResponse[itemBatchOperation.OperationIndex]; + itemBatchOperation.Context.Complete(this, response); } } } catch (Exception ex) { - DefaultTrace.TraceError("Exception during BatchAsyncBatcher: {0}", ex); // Exceptions happening during execution fail all the Tasks part of the request (excluding overflow) foreach (ItemBatchOperation itemBatchOperation in serverRequest.Operations) { - ItemBatchOperation operation = this.batchOperations[itemBatchOperation.OperationIndex]; - operation.Context.Fail(this, ex); + itemBatchOperation.Context.Fail(this, ex); } + + throw; } - finally - { - this.batchOperations.Clear(); - this.dispached = true; - } + + } + catch (Exception ex) + { + DefaultTrace.TraceError("Exception during BatchAsyncBatcher: {0}", ex); + } + finally + { + this.batchOperations.Clear(); + this.dispached = true; } } diff --git a/Microsoft.Azure.Cosmos/src/Util/InterlockIncrementCheck.cs b/Microsoft.Azure.Cosmos/src/Util/InterlockIncrementCheck.cs index fdb2bf50c7..40839298d6 100644 --- a/Microsoft.Azure.Cosmos/src/Util/InterlockIncrementCheck.cs +++ b/Microsoft.Azure.Cosmos/src/Util/InterlockIncrementCheck.cs @@ -10,7 +10,7 @@ namespace Microsoft.Azure.Cosmos /// /// This class is used to assert that a region of code can only be called concurrently by a limited amount of threads. /// - internal class InterlockIncrementCheck : IDisposable + internal class InterlockIncrementCheck { private readonly int maxConcurrentOperations; private int counter = 0; @@ -37,20 +37,13 @@ public InterlockIncrementCheck(int maxConcurrentOperations = 1) /// } /// /// - public InterlockIncrementCheck EnterLockCheck() + public void EnterLockCheck() { Interlocked.Increment(ref this.counter); if (this.counter > this.maxConcurrentOperations) { throw new InvalidOperationException($"InterlockIncrementCheck detected {this.counter} with a maximum of {this.maxConcurrentOperations}."); } - - return this; - } - - public void Dispose() - { - Interlocked.Decrement(ref this.counter); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InterlockIncrementCheckTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InterlockIncrementCheckTests.cs index 2c9535c2e4..25a2682c06 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InterlockIncrementCheckTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/InterlockIncrementCheckTests.cs @@ -20,19 +20,6 @@ public void ValidatesMaxConcurrentOperations(int maxConcurrentOperations) new InterlockIncrementCheck(maxConcurrentOperations); } - [TestMethod] - public void AllowsMultipleChainedOperations() - { - InterlockIncrementCheck check = new InterlockIncrementCheck(); - using (check.EnterLockCheck()) - { - } - - using (check.EnterLockCheck()) - { - } - } - [TestMethod] public async Task AllowsMultipleConcurrentOperations() { @@ -56,10 +43,8 @@ public async Task ThrowsOnMultipleConcurrentOperations() private async Task RunLock(InterlockIncrementCheck check) { - using (check.EnterLockCheck()) - { - await Task.Delay(500); - } + check.EnterLockCheck(); + await Task.Delay(500); } } }