Adding diagnostic context

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

@@ -157,7 +157,19 @@ public virtual bool TryAdd(ItemBatchOperation operation)
                         foreach (ItemBatchOperation itemBatchOperation in batchResponse.Operations)
                         {
                             TransactionalBatchOperationResult response = batchResponse[itemBatchOperation.OperationIndex];
-                            itemBatchOperation.Context.Diagnostics.AppendDiagnostics(batchResponse.Diagnostics);
+
+                            // Bulk has diagnostics per a item operation.
+                            // Batch has a single diagnostics for the execute operation
+                            if (itemBatchOperation.DiagnosticsContext != null)
+                            {
+                                response.DiagnosticsContext = itemBatchOperation.DiagnosticsContext;
+                                response.DiagnosticsContext.Append(batchResponse.DiagnosticsContext);
+                            }
+                            else
+                            {
+                                response.DiagnosticsContext = batchResponse.DiagnosticsContext;
+                            }
+
                             if (!response.IsSuccessStatusCode)
                             {
                                 Documents.ShouldRetryResult shouldRetry = await itemBatchOperation.Context.ShouldRetryAsync(response, cancellationToken);

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

@@ -220,13 +220,15 @@ private async Task<PartitionKeyRangeBatchExecutionResult> ExecuteAsync(
             PartitionKeyRangeServerBatchRequest serverRequest,
             CancellationToken cancellationToken)
         {
+            CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContext();
+            CosmosDiagnosticScope limiterScope = diagnosticsContext.CreateScope("BatchAsyncContainerExecutor.Limiter");
             SemaphoreSlim limiter = this.GetOrAddLimiterForPartitionKeyRange(serverRequest.PartitionKeyRangeId);
             using (await limiter.UsingWaitAsync(cancellationToken))
             {
+                limiterScope.Dispose();
                 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,
@@ -236,11 +238,15 @@ private async Task<PartitionKeyRangeBatchExecutionResult> ExecuteAsync(
                         partitionKey: null,
                         streamPayload: serverRequestPayload,
                         requestEnricher: requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest.PartitionKeyRangeId),
+                        diagnosticsScope: diagnosticsContext,
                         cancellationToken: cancellationToken).ConfigureAwait(false);
 
-                    TransactionalBatchResponse serverResponse = await TransactionalBatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.SerializerCore).ConfigureAwait(false);
+                    using (diagnosticsContext.CreateScope("BatchAsyncContainerExecutor.ToResponse"))
+                    {
+                        TransactionalBatchResponse serverResponse = await TransactionalBatchResponse.FromResponseMessageAsync(responseMessage, serverRequest, this.cosmosClientContext.SerializerCore).ConfigureAwait(false);
 
-                    return new PartitionKeyRangeBatchExecutionResult(serverRequest.PartitionKeyRangeId, serverRequest.Operations, serverResponse);
+                        return new PartitionKeyRangeBatchExecutionResult(serverRequest.PartitionKeyRangeId, serverRequest.Operations, serverResponse);
+                    }
                 }
             }
         }
@@ -252,7 +258,15 @@ private BatchAsyncStreamer GetOrAddStreamerForPartitionKeyRange(string partition
                 return streamer;
             }
 
-            BatchAsyncStreamer newStreamer = new BatchAsyncStreamer(this.maxServerRequestOperationCount, this.maxServerRequestBodyLength, this.dispatchTimerInSeconds, this.timerPool, this.cosmosClientContext.SerializerCore, this.ExecuteAsync, this.ReBatchAsync);
+            BatchAsyncStreamer newStreamer = new BatchAsyncStreamer(
+                this.maxServerRequestOperationCount,
+                this.maxServerRequestBodyLength,
+                this.dispatchTimerInSeconds,
+                this.timerPool,
+                this.cosmosClientContext.SerializerCore,
+                this.ExecuteAsync,
+                this.ReBatchAsync);
+
             if (!this.streamersByPartitionKeyRange.TryAdd(partitionKeyRangeId, newStreamer))
             {
                 newStreamer.Dispose();

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

@@ -209,7 +209,14 @@ public virtual Task<TransactionalBatchResponse> ExecuteAsync(
             RequestOptions requestOptions,
             CancellationToken cancellationToken = default(CancellationToken))
         {
-            BatchExecutor executor = new BatchExecutor(this.container, this.partitionKey, this.operations, requestOptions);
+            CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContext();
+            BatchExecutor executor = new BatchExecutor(
+                container: this.container,
+                partitionKey: this.partitionKey,
+                operations: this.operations,
+                batchOptions: requestOptions,
+                diagnosticsContext: diagnosticsContext);
+
             this.operations = new List<ItemBatchOperation>();
             return executor.ExecuteAsync(cancellationToken);
         }

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

@@ -25,36 +25,47 @@ internal sealed class BatchExecutor
 
         private readonly RequestOptions batchOptions;
 
+        private readonly CosmosDiagnosticsContext diagnosticsContext;
+
         public BatchExecutor(
             ContainerCore container,
             PartitionKey partitionKey,
             IReadOnlyList<ItemBatchOperation> operations,
-            RequestOptions batchOptions)
+            RequestOptions batchOptions,
+            CosmosDiagnosticsContext diagnosticsContext)
         {
             this.container = container;
             this.clientContext = this.container.ClientContext;
             this.inputOperations = operations;
             this.partitionKey = partitionKey;
             this.batchOptions = batchOptions;
+            this.diagnosticsContext = diagnosticsContext;
         }
 
         public async Task<TransactionalBatchResponse> ExecuteAsync(CancellationToken cancellationToken)
         {
-            BatchExecUtils.EnsureValid(this.inputOperations, this.batchOptions);
-
-            PartitionKey? serverRequestPartitionKey = this.partitionKey;
-            if (this.batchOptions != null && this.batchOptions.IsEffectivePartitionKeyRouting)
+            using (this.diagnosticsContext.CreateOverallScope("BatchExecuteAsync"))
             {
-                serverRequestPartitionKey = null;
-            }
+                BatchExecUtils.EnsureValid(this.inputOperations, this.batchOptions);
 
-            SinglePartitionKeyServerBatchRequest serverRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync(
-                      serverRequestPartitionKey,
-                      new ArraySegment<ItemBatchOperation>(this.inputOperations.ToArray()),
-                      this.clientContext.SerializerCore,
-                      cancellationToken);
+                PartitionKey? serverRequestPartitionKey = this.partitionKey;
+                if (this.batchOptions != null && this.batchOptions.IsEffectivePartitionKeyRouting)
+                {
+                    serverRequestPartitionKey = null;
+                }
 
-            return await this.ExecuteServerRequestAsync(serverRequest, cancellationToken);
+                SinglePartitionKeyServerBatchRequest serverRequest;
+                using (this.diagnosticsContext.CreateScope("CreateBatchRequest"))
+                {
+                    serverRequest = await SinglePartitionKeyServerBatchRequest.CreateAsync(
+                          serverRequestPartitionKey,
+                          new ArraySegment<ItemBatchOperation>(this.inputOperations.ToArray()),
+                          this.clientContext.SerializerCore,
+                          cancellationToken);
+                }
+
+                return await this.ExecuteServerRequestAsync(serverRequest, cancellationToken);
+            }
         }
 
         /// <summary>
@@ -84,12 +95,16 @@ private async Task<TransactionalBatchResponse> ExecuteServerRequestAsync(
                         requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchAtomic, bool.TrueString);
                         requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchOrdered, bool.TrueString);
                     },
+                    diagnosticsScope: this.diagnosticsContext,
                     cancellationToken);
 
-                return await TransactionalBatchResponse.FromResponseMessageAsync(
-                    responseMessage,
-                    serverRequest,
-                    this.clientContext.SerializerCore);
+                using (this.diagnosticsContext.CreateScope("TransactionalBatchResponse"))
+                {
+                    return await TransactionalBatchResponse.FromResponseMessageAsync(
+                        responseMessage,
+                        serverRequest,
+                        this.clientContext.SerializerCore);
+                }
             }
         }
     }

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

@@ -31,14 +31,16 @@ public ItemBatchOperation(
             PartitionKey partitionKey,
             string id = null,
             Stream resourceStream = null,
-            TransactionalBatchItemRequestOptions requestOptions = null)
+            TransactionalBatchItemRequestOptions requestOptions = null,
+            CosmosDiagnosticsContext diagnosticsContext = null)
         {
             this.OperationType = operationType;
             this.OperationIndex = operationIndex;
             this.PartitionKey = partitionKey;
             this.Id = id;
             this.ResourceStream = resourceStream;
             this.RequestOptions = requestOptions;
+            this.DiagnosticsContext = diagnosticsContext;
         }
 
         public ItemBatchOperation(
@@ -53,6 +55,7 @@ public ItemBatchOperation(
             this.Id = id;
             this.ResourceStream = resourceStream;
             this.RequestOptions = requestOptions;
+            this.DiagnosticsContext = null;
         }
 
         public PartitionKey? PartitionKey { get; internal set; }
@@ -67,6 +70,8 @@ public ItemBatchOperation(
 
         public int OperationIndex { get; internal set; }
 
+        internal CosmosDiagnosticsContext DiagnosticsContext { get; }
+
         internal string PartitionKeyJson { get; set; }
 
         internal Documents.PartitionKey ParsedPartitionKey { get; set; }

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

@@ -21,8 +21,6 @@ internal class ItemBatchOperationContext : IDisposable
 
         public Task<TransactionalBatchOperationResult> OperationTask => this.taskCompletionSource.Task;
 
-        public ItemBatchOperationStatistics Diagnostics { get; } = new ItemBatchOperationStatistics();
-
         private readonly IDocumentClientRetryPolicy retryPolicy;
 
         private TaskCompletionSource<TransactionalBatchOperationResult> taskCompletionSource = new TaskCompletionSource<TransactionalBatchOperationResult>();
@@ -58,8 +56,6 @@ public void Complete(
         {
             if (this.AssertBatcher(completer))
             {
-                this.Diagnostics.Complete();
-                result.Diagnostics = this.Diagnostics;
                 this.taskCompletionSource.SetResult(result);
             }
 

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

@@ -70,6 +70,8 @@ internal PartitionKeyRangeBatchResponse(
         /// <inheritdoc />
         public override CosmosDiagnostics Diagnostics => this.serverResponse.Diagnostics;
 
+        internal override CosmosDiagnosticsContext DiagnosticsContext => this.serverResponse.DiagnosticsContext;
+
         internal override CosmosSerializerCore SerializerCore { get; }
 
         /// <summary>

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

@@ -46,7 +46,7 @@ public PartitionKeyRangeServerBatchRequest(
         /// <param name="ensureContinuousOperationIndexes">Whether to stop adding operations to the request once there is non-continuity in the operation indexes.</param>
         /// <param name="serializerCore">Serializer to serialize user provided objects to JSON.</param>
         /// <param name="cancellationToken"><see cref="CancellationToken"/> representing request cancellation.</param>
-        /// <returns>A newly created instance of <see cref="PartitionKeyRangeServerBatchRequest"/>.</returns>
+        /// <returns>A newly created instance of <see cref="PartitionKeyRangeServerBatchRequest"/> and the overflow ItemBatchOperation not being processed.</returns>
         public static async Task<Tuple<PartitionKeyRangeServerBatchRequest, ArraySegment<ItemBatchOperation>>> CreateAsync(
             string partitionKeyRangeId,
             ArraySegment<ItemBatchOperation> operations,

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

@@ -16,7 +16,8 @@ namespace Microsoft.Azure.Cosmos
     /// </summary>
     public class TransactionalBatchOperationResult
     {
-        internal TransactionalBatchOperationResult(HttpStatusCode statusCode)
+        internal TransactionalBatchOperationResult(
+            HttpStatusCode statusCode)
         {
             this.StatusCode = statusCode;
         }
@@ -92,7 +93,7 @@ public virtual bool IsSuccessStatusCode
         /// <summary>
         /// Gets the cosmos diagnostic information for the current request to Azure Cosmos DB service
         /// </summary>
-        internal virtual CosmosDiagnostics Diagnostics { get; set; }
+        internal virtual CosmosDiagnosticsContext DiagnosticsContext { get; set; }
 
         internal static Result ReadOperationResult(Memory<byte> input, out TransactionalBatchOperationResult batchOperationResult)
         {
@@ -200,13 +201,25 @@ private static Result ReadOperationResult(ref RowReader reader, out Transactiona
 
         internal ResponseMessage ToResponseMessage()
         {
-            ResponseMessage responseMessage = new ResponseMessage(this.StatusCode);
-            responseMessage.Headers.SubStatusCode = this.SubStatusCode;
-            responseMessage.Headers.ETag = this.ETag;
-            responseMessage.Headers.RetryAfter = this.RetryAfter;
-            responseMessage.Headers.RequestCharge = this.RequestCharge;
-            responseMessage.Content = this.ResourceStream;
-            responseMessage.Diagnostics = this.Diagnostics;
+            Headers headers = new Headers()
+            {
+                SubStatusCode = this.SubStatusCode,
+                ETag = this.ETag,
+                RetryAfter = this.RetryAfter,
+                RequestCharge = this.RequestCharge,
+            };
+
+            ResponseMessage responseMessage = new ResponseMessage(
+                statusCode: this.StatusCode,
+                requestMessage: null,
+                errorMessage: null,
+                error: null,
+                headers: headers,
+                diagnostics: this.DiagnosticsContext ?? new CosmosDiagnosticsContext())
+            {
+                Content = this.ResourceStream
+            };
+
             return responseMessage;
         }
     }

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

@@ -33,18 +33,20 @@ public class TransactionalBatchResponse : IReadOnlyList<TransactionalBatchOperat
         /// <param name="subStatusCode">Provides further details about why the batch was not processed.</param>
         /// <param name="errorMessage">The reason for failure.</param>
         /// <param name="operations">Operations that were to be executed.</param>
+        /// <param name="diagnosticsContext">Diagnostics for the operation</param>
         internal TransactionalBatchResponse(
             HttpStatusCode statusCode,
             SubStatusCodes subStatusCode,
             string errorMessage,
-            IReadOnlyList<ItemBatchOperation> operations)
+            IReadOnlyList<ItemBatchOperation> operations,
+            CosmosDiagnosticsContext diagnosticsContext)
             : this(statusCode,
                   subStatusCode,
                   errorMessage,
                   requestCharge: 0,
                   retryAfter: null,
                   activityId: Guid.Empty.ToString(),
-                  cosmosDiagnostics: null,
+                  diagnosticsContext: diagnosticsContext,
                   operations: operations,
                   serializer: null)
         {
@@ -65,7 +67,7 @@ private TransactionalBatchResponse(
             double requestCharge,
             TimeSpan? retryAfter,
             string activityId,
-            CosmosDiagnostics cosmosDiagnostics,
+            CosmosDiagnosticsContext diagnosticsContext,
             IReadOnlyList<ItemBatchOperation> operations,
             CosmosSerializerCore serializer)
         {
@@ -77,7 +79,8 @@ private TransactionalBatchResponse(
             this.RequestCharge = requestCharge;
             this.RetryAfter = retryAfter;
             this.ActivityId = activityId;
-            this.Diagnostics = cosmosDiagnostics;
+            this.Diagnostics = diagnosticsContext;
+            this.DiagnosticsContext = diagnosticsContext ?? throw new ArgumentNullException(nameof(diagnosticsContext));
         }
 
         /// <summary>
@@ -132,6 +135,8 @@ public virtual bool IsSuccessStatusCode
         /// </summary>
         public virtual CosmosDiagnostics Diagnostics { get; }
 
+        internal virtual CosmosDiagnosticsContext DiagnosticsContext { get; }
+
         internal virtual SubStatusCodes SubStatusCode { get; }
 
         internal virtual CosmosSerializerCore SerializerCore { get; }
@@ -248,7 +253,7 @@ internal static async Task<TransactionalBatchResponse> FromResponseMessageAsync(
                                 responseMessage.Headers.RequestCharge,
                                 responseMessage.Headers.RetryAfter,
                                 responseMessage.Headers.ActivityId,
-                                responseMessage.Diagnostics,
+                                responseMessage.DiagnosticsContext,
                                 serverRequest.Operations,
                                 serializer);
                         }
@@ -264,7 +269,7 @@ internal static async Task<TransactionalBatchResponse> FromResponseMessageAsync(
                         responseMessage.Headers.RequestCharge,
                         responseMessage.Headers.RetryAfter,
                         responseMessage.Headers.ActivityId,
-                        responseMessage.Diagnostics,
+                        responseMessage.DiagnosticsContext,
                         serverRequest.Operations,
                         serializer);
                 }
@@ -282,7 +287,7 @@ internal static async Task<TransactionalBatchResponse> FromResponseMessageAsync(
                             responseMessage.Headers.RequestCharge,
                             responseMessage.Headers.RetryAfter,
                             responseMessage.Headers.ActivityId,
-                            responseMessage.Diagnostics,
+                            responseMessage.DiagnosticsContext,
                             serverRequest.Operations,
                             serializer);
                     }
@@ -378,7 +383,7 @@ record =>
                 responseMessage.Headers.RequestCharge,
                 responseMessage.Headers.RetryAfter,
                 responseMessage.Headers.ActivityId,
-                responseMessage.Diagnostics,
+                responseMessage.DiagnosticsContext,
                 serverRequest.Operations,
                 serializer);