Add GetClientElapsedTime to CosmosDiagnostics

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

@@ -134,7 +134,11 @@ internal virtual async Task ValidateOperationAsync(
                     throw new InvalidOperationException(ClientResources.UnsupportedBulkRequestOptions);
                 }
 
-                Debug.Assert(itemRequestOptions.DiagnosticContext == null, "Disable Diagnostics is not supported on Bulk operations");
+                if (itemRequestOptions.DiagnosticContext != null)
+                {
+                    throw new ArgumentException("DiagnosticContext is not allowed when AllowBulkExecution is set to true");
+                }
+
                 Debug.Assert(BatchAsyncContainerExecutor.ValidateOperationEPK(operation, itemRequestOptions));
             }
 
@@ -242,7 +246,7 @@ private async Task<PartitionKeyRangeBatchExecutionResult> ExecuteAsync(
                         partitionKey: null,
                         streamPayload: serverRequestPayload,
                         requestEnricher: requestMessage => BatchAsyncContainerExecutor.AddHeadersToRequestMessage(requestMessage, serverRequest.PartitionKeyRangeId),
-                        diagnosticsScope: diagnosticsContext,
+                        diagnosticsContext: diagnosticsContext,
                         cancellationToken: cancellationToken).ConfigureAwait(false);
 
                     using (diagnosticsContext.CreateScope("BatchAsyncContainerExecutor.ToResponse"))

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

@@ -95,7 +95,7 @@ private async Task<TransactionalBatchResponse> ExecuteServerRequestAsync(
                         requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchAtomic, bool.TrueString);
                         requestMessage.Headers.Add(HttpConstants.HttpHeaders.IsBatchOrdered, bool.TrueString);
                     },
-                    diagnosticsScope: this.diagnosticsContext,
+                    diagnosticsContext: this.diagnosticsContext,
                     cancellationToken);
 
                 using (this.diagnosticsContext.CreateScope("TransactionalBatchResponse"))

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

@@ -79,7 +79,7 @@ private TransactionalBatchResponse(
             this.RequestCharge = requestCharge;
             this.RetryAfter = retryAfter;
             this.ActivityId = activityId;
-            this.Diagnostics = diagnosticsContext;
+            this.Diagnostics = diagnosticsContext.Diagnostics;
             this.DiagnosticsContext = diagnosticsContext ?? throw new ArgumentNullException(nameof(diagnosticsContext));
         }
 

Microsoft.Azure.Cosmos/src/CosmosClient.cs

@@ -819,7 +819,7 @@ private Task<ResponseMessage> CreateDatabaseStreamInternalAsync(
                 partitionKey: null,
                 streamPayload: streamPayload,
                 requestEnricher: (httpRequestMessage) => httpRequestMessage.AddThroughputHeader(throughput),
-                diagnosticsScope: null,
+                diagnosticsContext: null,
                 cancellationToken: cancellationToken);
         }
 

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosDiagnosticScope.cs

@@ -15,16 +15,14 @@ namespace Microsoft.Azure.Cosmos.Diagnostics
     internal sealed class CosmosDiagnosticScope : CosmosDiagnosticsInternal, IDisposable
     {
         private readonly Stopwatch ElapsedTimeStopWatch;
-        private readonly Action<TimeSpan> ElapsedTimeCallback;
         private bool isDisposed = false;
 
         public CosmosDiagnosticScope(
             string name,
-            Action<TimeSpan> elapsedTimeCallback = null)
+            Stopwatch stopwatch = null)
         {
             this.Id = name;
-            this.ElapsedTimeStopWatch = Stopwatch.StartNew();
-            this.ElapsedTimeCallback = elapsedTimeCallback;
+            this.ElapsedTimeStopWatch = stopwatch ?? Stopwatch.StartNew();
         }
 
         public string Id { get; }
@@ -48,7 +46,6 @@ public void Dispose()
             }
 
             this.ElapsedTimeStopWatch.Stop();
-            this.ElapsedTimeCallback?.Invoke(this.ElapsedTimeStopWatch.Elapsed);
             this.isDisposed = true;
         }
 

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosDiagnostics.cs

@@ -1,14 +1,21 @@
 //------------------------------------------------------------
 // Copyright (c) Microsoft Corporation.  All rights reserved.
 //------------------------------------------------------------
-
 namespace Microsoft.Azure.Cosmos
 {
+    using System;
+
     /// <summary>
     ///  Contains the cosmos diagnostic information for the current request to Azure Cosmos DB service.
     /// </summary>
     public abstract class CosmosDiagnostics
     {
+        /// <summary>
+        /// This represent the current elapsed latency of the request. 
+        /// </summary>
+        /// <returns>The elapsed time since the start of the request.</returns>
+        public abstract TimeSpan GetElapsedClientLatency();
+
         /// <summary>
         /// Gets the string field <see cref="CosmosDiagnostics"/> instance in the Azure CosmosDB database service.
         /// </summary>

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosDiagnosticsContext.cs

@@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos
     using System;
     using System.Collections;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using Microsoft.Azure.Cosmos.Diagnostics;
     using Microsoft.Azure.Documents;
     using static Microsoft.Azure.Cosmos.CosmosClientSideRequestStatistics;
@@ -21,10 +22,12 @@ internal abstract class CosmosDiagnosticsContext : CosmosDiagnosticsInternal, IE
 
         public abstract int FailedRequestCount { get; protected set; }
 
-        public abstract TimeSpan? TotalElapsedTime { get; protected set; }
-
         public abstract string UserAgent { get; protected set; }
 
+        internal abstract CosmosDiagnostics Diagnostics { get; }
+
+        public abstract Stopwatch OverallClientRequestTime { get; }
+
         internal abstract CosmosDiagnosticScope CreateOverallScope(string name);
 
         internal abstract CosmosDiagnosticScope CreateScope(string name);
@@ -42,7 +45,7 @@ internal abstract class CosmosDiagnosticsContext : CosmosDiagnosticsInternal, IE
         internal abstract void AddDiagnosticsInternal(CosmosDiagnosticsContext newContext);
 
         internal abstract void SetSdkUserAgent(string userAgent);
-
+       
         public abstract IEnumerator<CosmosDiagnosticsInternal> GetEnumerator();
 
         IEnumerator IEnumerable.GetEnumerator()

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosDiagnosticsContextCore.cs

@@ -5,6 +5,9 @@ namespace Microsoft.Azure.Cosmos
 {
     using System;
     using System.Collections.Generic;
+    using System.Collections.ObjectModel;
+    using System.Diagnostics;
+    using System.Linq;
     using Microsoft.Azure.Cosmos.Diagnostics;
 
     /// <summary>
@@ -37,6 +40,8 @@ public CosmosDiagnosticsContextCore()
         {
             this.StartUtc = DateTime.UtcNow;
             this.ContextList = new List<CosmosDiagnosticsInternal>();
+            this.Diagnostics = new CosmosDiagnosticsCore(this);
+            this.OverallClientRequestTime = Stopwatch.StartNew();
         }
 
         public override DateTime StartUtc { get; }
@@ -45,10 +50,12 @@ public CosmosDiagnosticsContextCore()
 
         public override int FailedRequestCount { get; protected set; }
 
-        public override TimeSpan? TotalElapsedTime { get; protected set; }
-
         public override string UserAgent { get; protected set; } = CosmosDiagnosticsContextCore.DefaultUserAgentString;
 
+        internal override CosmosDiagnostics Diagnostics { get; }
+
+        public override Stopwatch OverallClientRequestTime { get; }
+
         internal override CosmosDiagnosticScope CreateOverallScope(string name)
         {
             CosmosDiagnosticScope scope;
@@ -59,8 +66,8 @@ internal override CosmosDiagnosticScope CreateOverallScope(string name)
             }
             else
             {
-                scope = new CosmosDiagnosticScope(name, this.SetElapsedTime);
                 this.isOverallScopeSet = true;
+                scope = new CosmosDiagnosticScope(name, this.OverallClientRequestTime);
             }
 
             this.ContextList.Add(scope);
@@ -158,11 +165,6 @@ private void AddRequestCount(int statusCode)
             }
         }
 
-        private void SetElapsedTime(TimeSpan totalElapsedTime)
-        {
-            this.TotalElapsedTime = totalElapsedTime;
-        }
-
         private void AddSummaryInfo(CosmosDiagnosticsContext newContext)
         {
             if (Object.ReferenceEquals(this, newContext))
@@ -175,12 +177,6 @@ private void AddSummaryInfo(CosmosDiagnosticsContext newContext)
                 this.SetSdkUserAgent(newContext.UserAgent);
             }
 
-            // Use the larger of the total elapsed times
-            if (this.TotalElapsedTime < newContext.TotalElapsedTime)
-            {
-                this.TotalElapsedTime = newContext.TotalElapsedTime;
-            }
-
             this.TotalRequestCount += newContext.TotalRequestCount;
             this.FailedRequestCount += newContext.FailedRequestCount;
         }

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosDiagnosticsCore.cs

@@ -0,0 +1,33 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+    using System.Collections.Generic;
+
+    /// <summary>
+    ///  Contains the cosmos diagnostic information for the current request to Azure Cosmos DB service.
+    /// </summary>
+    internal class CosmosDiagnosticsCore : CosmosDiagnostics
+    {
+        internal CosmosDiagnosticsCore(CosmosDiagnosticsContext diagnosticsContext)
+        {
+            this.Context = diagnosticsContext ?? throw new ArgumentNullException(nameof(diagnosticsContext));
+        }
+
+        internal CosmosDiagnosticsContext Context { get; }
+
+        /// <inheritdoc/>
+        public override TimeSpan GetElapsedClientLatency()
+        {
+            return this.Context.OverallClientRequestTime.Elapsed;
+        }
+
+        /// <inheritdoc/>
+        public override string ToString()
+        {
+            return this.Context.ToString();
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosDiagnosticsInternal.cs

@@ -10,7 +10,7 @@ namespace Microsoft.Azure.Cosmos.Diagnostics
     /// <summary>
     /// Extends <see cref="CosmosDiagnostics"/> to expose internal APIs.
     /// </summary>
-    internal abstract class CosmosDiagnosticsInternal : CosmosDiagnostics
+    internal abstract class CosmosDiagnosticsInternal
     {
         public abstract void Accept(CosmosDiagnosticsInternalVisitor visitor);
 

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosDiagnosticsSerializerVisitor.cs

@@ -67,16 +67,16 @@ public override void Visit(CosmosDiagnosticsContext cosmosDiagnosticsContext)
             this.jsonWriter.WritePropertyName("StartUtc");
             this.jsonWriter.WriteValue(cosmosDiagnosticsContext.StartUtc.ToString("o", CultureInfo.InvariantCulture));
 
-            this.jsonWriter.WritePropertyName("ElapsedTime");
-
-            if (cosmosDiagnosticsContext.TotalElapsedTime.HasValue)
+            if (cosmosDiagnosticsContext.OverallClientRequestTime.IsRunning)
             {
-                this.jsonWriter.WriteValue(cosmosDiagnosticsContext.TotalElapsedTime.Value);
+                this.jsonWriter.WritePropertyName("CurrentElapsedTime");
             }
             else
             {
-                this.jsonWriter.WriteValue("Timer Never Stopped.");
+                this.jsonWriter.WritePropertyName("TotalElapsedTime");
             }
+
+            this.jsonWriter.WriteValue(cosmosDiagnosticsContext.OverallClientRequestTime.Elapsed);
 
             this.jsonWriter.WritePropertyName("UserAgent");
             this.jsonWriter.WriteValue(cosmosDiagnosticsContext.UserAgent);

Microsoft.Azure.Cosmos/src/Diagnostics/EmptyCosmosDiagnosticsContext.cs

@@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using Microsoft.Azure.Cosmos.Diagnostics;
 
     /// <summary>
@@ -18,20 +19,27 @@ internal sealed class EmptyCosmosDiagnosticsContext : CosmosDiagnosticsContext
 
         public static readonly CosmosDiagnosticsContext Singleton = new EmptyCosmosDiagnosticsContext();
 
+        private static readonly Stopwatch Stopwatch = new Stopwatch();
+
+        private static readonly DateTime DefaultStartUtc = new DateTime(0);
+
         private EmptyCosmosDiagnosticsContext()
         {
+            this.Diagnostics = new CosmosDiagnosticsCore(this);
         }
 
-        public override DateTime StartUtc { get; } = new DateTime(0);
+        public override DateTime StartUtc { get; } = EmptyCosmosDiagnosticsContext.DefaultStartUtc;
 
         public override int TotalRequestCount { get; protected set; }
 
         public override int FailedRequestCount { get; protected set; }
 
-        public override TimeSpan? TotalElapsedTime { get; protected set; }
-
         public override string UserAgent { get; protected set; } = "Empty Context";
 
+        internal override CosmosDiagnostics Diagnostics { get; }
+
+        public override Stopwatch OverallClientRequestTime => EmptyCosmosDiagnosticsContext.Stopwatch;
+
         internal override CosmosDiagnosticScope CreateOverallScope(string name)
         {
             return EmptyCosmosDiagnosticsContext.DefaultScope;

Microsoft.Azure.Cosmos/src/FeedIteratorCore.cs

@@ -96,7 +96,7 @@ public override async Task<ResponseMessage> ReadNextAsync(CancellationToken canc
                        request.Headers.Add(HttpConstants.HttpHeaders.IsQuery, bool.TrueString);
                    }
                },
-               diagnosticsScope: null,
+               diagnosticsContext: null,
                cancellationToken: cancellationToken);
 
             this.ContinuationToken = response.Headers.ContinuationToken;

Microsoft.Azure.Cosmos/src/Handler/ResponseMessage.cs

@@ -122,7 +122,7 @@ public virtual Stream Content
         /// <summary>
         /// Gets the cosmos diagnostic information for the current request to Azure Cosmos DB service
         /// </summary>
-        public virtual CosmosDiagnostics Diagnostics => this.DiagnosticsContext;
+        public virtual CosmosDiagnostics Diagnostics => this.DiagnosticsContext.Diagnostics;
 
         internal CosmosDiagnosticsContext DiagnosticsContext { get; }
 

Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs

@@ -743,7 +743,7 @@ private static Task<Response<T>> ResponseHelperAsync<T>(T value)
                            System.Net.HttpStatusCode.OK,
                            new Headers(),
                            value,
-                           new CosmosDiagnosticsContextCore()));
+                           new CosmosDiagnosticsContextCore().Diagnostics));
         }
 
         private static MethodInfo GetMethodInfoOf<T1, T2>(Func<T1, T2> func)

Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs

@@ -184,15 +184,15 @@ internal async Task<Response<T>> AggregateResultAsync(CancellationToken cancella
                 headers.RequestCharge += response.RequestCharge;
 
                 // If the first page has a diagnostic context use that. Else create a new one and add the diagnostic to it.
-                if (response.Diagnostics is CosmosDiagnosticsContext responseDiagnosticContext)
+                if (response.Diagnostics is CosmosDiagnosticsCore diagnosticsCore)
                 {
                     if (diagnosticsContext == null)
                     {
-                        diagnosticsContext = responseDiagnosticContext;
+                        diagnosticsContext = diagnosticsCore.Context;
                     }
                     else
                     {
-                        diagnosticsContext.AddDiagnosticsInternal(responseDiagnosticContext);
+                        diagnosticsContext.AddDiagnosticsInternal(diagnosticsCore.Context);
                     }
 
                 }
@@ -208,7 +208,7 @@ internal async Task<Response<T>> AggregateResultAsync(CancellationToken cancella
                 System.Net.HttpStatusCode.OK,
                 headers,
                 result.FirstOrDefault(),
-                diagnosticsContext);
+                diagnosticsContext.Diagnostics);
         }
 
         private FeedIteratorInternal CreateStreamIterator(bool isContinuationExcpected)

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -153,7 +153,7 @@ internal override async Task<QueryResponseCore> ExecuteItemQueryAsync<RequestOpt
                     cosmosRequestMessage.Headers.Add(HttpConstants.HttpHeaders.ContentType, MediaTypes.QueryJson);
                     cosmosRequestMessage.Headers.Add(HttpConstants.HttpHeaders.IsQuery, bool.TrueString);
                 },
-                diagnosticsScope: queryRequestOptions?.DiagnosticContext,
+                diagnosticsContext: null,
                 cancellationToken: cancellationToken);
 
             schedulingStopwatch.Stop();
@@ -192,7 +192,7 @@ internal override async Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequ
                     requestMessage.Headers.Add(HttpConstants.HttpHeaders.QueryVersion, new Version(major: 1, minor: 0).ToString());
                     requestMessage.UseGatewayMode = true;
                 },
-                diagnosticsScope: null,
+                diagnosticsContext: null,
                 cancellationToken: cancellationToken))
             {
                 // Syntax exception are argument exceptions and thrown to the user.

Microsoft.Azure.Cosmos/src/Resource/Conflict/ConflictsCore.cs

@@ -58,7 +58,7 @@ public override Task<ResponseMessage> DeleteAsync(
                 partitionKey: partitionKey,
                 streamPayload: null,
                 requestEnricher: null,
-                diagnosticsScope: null,
+                diagnosticsContext: null,
                 cancellationToken: cancellationToken);
         }
 
@@ -169,7 +169,7 @@ public override async Task<ItemResponse<T>> ReadCurrentAsync<T>(
                 partitionKey: partitionKey,
                 streamPayload: null,
                 requestEnricher: null,
-                diagnosticsScope: null,
+                diagnosticsContext: null,
                 cancellationToken: cancellationToken);
 
             return await this.clientContext.ResponseFactory.CreateItemResponseAsync<T>(response);