[Preview] AI Integration: Adds latency based filters to add request diagnostics in Activity and some minor changes

Microsoft.Azure.Cosmos.Samples/Tools/Benchmark/BenchmarkConfig.cs

@@ -98,9 +98,6 @@ public class BenchmarkConfig
         [Option(Required = false, HelpText = "Disable core SDK logging")]
         public bool DisableCoreSdkLogging { get; set; }
 
-        [Option(Required = false, HelpText = "Enable Open Telemetry")]
-        public bool EnableOpenTelemetry { get; set; }
-
         [Option(Required = false, HelpText = "Enable Client Telemetry")]
         public bool EnableTelemetry { get; set; }
 
@@ -196,11 +193,6 @@ internal Microsoft.Azure.Cosmos.CosmosClient CreateCosmosClient(string accountKe
                 MaxTcpConnectionsPerEndpoint = this.MaxTcpConnectionsPerEndpoint
             };
 
-            if (this.EnableOpenTelemetry)
-            {
-                clientOptions.EnableOpenTelemetry = true;
-            }
-
             if (this.EnableTelemetry)
             {
                 Environment.SetEnvironmentVariable(

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Client;
     using Newtonsoft.Json;
+    using Telemetry;
 
     /// <summary>
     /// Defines all the configurable options that the CosmosClient requires.
@@ -925,13 +926,15 @@ public override bool CanConvert(Type objectType)
         }
 
         /// <summary>
-        /// Enable OpenTelemetry and start emiting activities for each operations
+        /// Distributed Tracing Options. <see cref="Microsoft.Azure.Cosmos.DistributedTracingOptions"/>
         /// </summary>
-#if PREVIEW
-        public
-#else
-        internal
-#endif
-            bool EnableOpenTelemetry { get; set; }
+        internal DistributedTracingOptions DistributedTracingOptions { get; set; }
+
+        /// <summary>
+        /// Gets or sets value indicating whether distributed tracing activities (<see cref="System.Diagnostics.Activity"/>) are going to be created for the SDK methods calls and HTTP calls.
+        /// By default true for Preview package
+        /// </summary>
+        internal bool EnableDistributedTracing { get; set; }
+
     }
 }

Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs

@@ -10,12 +10,12 @@ namespace Microsoft.Azure.Cosmos.Fluent
     using System.Net.Http;
     using System.Threading;
     using System.Threading.Tasks;
-    using Azure;
     using global::Azure;
     using global::Azure.Core;
     using Microsoft.Azure.Cosmos.Core.Trace;
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Client;
+    using Telemetry;
 
     /// <summary>
     /// This is a Builder class that creates a cosmos client
@@ -409,7 +409,7 @@ public CosmosClientBuilder WithConnectionModeDirect(TimeSpan? idleTcpConnectionT
 
             return this;
         }
-
+        
         /// <summary>
         /// This can be used to weaken the database account consistency level for read operations.
         /// If this is not set the database account consistency level will be used for all requests.
@@ -420,20 +420,18 @@ public CosmosClientBuilder WithConsistencyLevel(Cosmos.ConsistencyLevel consiste
         {
             this.clientOptions.ConsistencyLevel = consistencyLevel;
             return this;
+
         }
 
         /// <summary>
-        /// Enable OpenTelemetry and start emiting activities for each operations
+        /// If Open Telemetry listener is subscribed for Azure.Cosmos namespace, There are <see cref="Microsoft.Azure.Cosmos.DistributedTracingOptions"/> you can leverage to control it.<br></br>
         /// </summary>
+        /// <param name="options">Tracing Options <see cref="Microsoft.Azure.Cosmos.DistributedTracingOptions"/></param>
         /// <returns>The current <see cref="CosmosClientBuilder"/>.</returns>
-#if PREVIEW
-        public
-#else
-        internal
-#endif
-            CosmosClientBuilder EnableOpenTelemetry()
-        {
-            this.clientOptions.EnableOpenTelemetry = true;
+        internal CosmosClientBuilder WithDistributingTracing(DistributedTracingOptions options)
+        {
+            this.clientOptions.DistributedTracingOptions = options;
+
             return this;
         }
 

Microsoft.Azure.Cosmos/src/RequestOptions/RequestOptions.cs

@@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos
     using System;
     using System.Collections.Generic;
     using Microsoft.Azure.Documents;
+    using Telemetry;
 
     /// <summary>
     /// The default cosmos request options
@@ -41,6 +42,11 @@ public class RequestOptions
         /// </summary>
         public Action<Headers> AddRequestHeaders { get; set; }
 
+        /// <summary>
+        /// Set Request Level Distributed Tracing Options.
+        /// </summary>
+        internal DistributedTracingOptions DistributedTracingOptions { get; set; }
+
         /// <summary>
         /// Gets or sets the boolean to use effective partition key routing in the cosmos db request.
         /// </summary>

Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs

@@ -17,8 +17,6 @@ namespace Microsoft.Azure.Cosmos
     using Microsoft.Azure.Cosmos.Resource.CosmosExceptions;
     using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Cosmos.Telemetry;
-    using Microsoft.Azure.Cosmos.Telemetry.Diagnostics;
-    using Microsoft.Azure.Cosmos.Telemetry.OpenTelemetry;
     using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Documents;
 
@@ -266,7 +264,8 @@ private async Task<TResult> OperationHelperWithRootTraceAsync<TResult>(
                     trace,
                     task,
                     openTelemetry,
-                    operationName);
+                    operationName,
+                    requestOptions);
             }
         }
 
@@ -295,7 +294,8 @@ private Task<TResult> OperationHelperWithRootTraceWithSynchronizationContextAsyn
                         trace,
                         task,
                         openTelemetry,
-                        operationName);
+                        operationName,
+                        requestOptions);
                 }
             });
         }
@@ -471,22 +471,24 @@ private async Task<TResult> RunWithDiagnosticsHelperAsync<TResult>(
             ITrace trace,
             Func<ITrace, Task<TResult>> task,
             Func<TResult, OpenTelemetryAttributes> openTelemetry,
-            string operationName)
+            string operationName,
+            RequestOptions requestOptions)
         {
             using (OpenTelemetryCoreRecorder recorder = 
                                 OpenTelemetryRecorderFactory.CreateRecorder(
                                     operationName: operationName,
-                                    isFeatureEnabled: this.clientOptions.EnableOpenTelemetry))
+                                    requestOptions: requestOptions,
+                                    clientContext: this.isDisposed ? null : this))
             using (new ActivityScope(Guid.NewGuid()))
             {
                 try
                 {
+                    // Record Operation Name
+                    recorder.Record(OpenTelemetryAttributeKeys.DbOperation, operationName);
+
                     TResult result = await task(trace).ConfigureAwait(false);
                     if (openTelemetry != null && recorder.IsEnabled)
                     {
-                        // Record client and other information
-                        recorder.Record(operationName, this);
-
                         // Record request response information
                         OpenTelemetryAttributes response = openTelemetry(result);
                         recorder.Record(response);

Microsoft.Azure.Cosmos/src/Resource/CosmosClientContext.cs

@@ -9,9 +9,9 @@ namespace Microsoft.Azure.Cosmos
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Handlers;
-    using Microsoft.Azure.Cosmos.Telemetry.OpenTelemetry;
     using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Documents;
+    using Telemetry;
 
     /// <summary>
     /// This class is used to get access to different client level operations without directly referencing the client object.

Microsoft.Azure.Cosmos/src/Resource/Database/DatabaseInlineCore.cs

@@ -30,7 +30,7 @@ public override Task<ContainerResponse> CreateContainerAsync(
                 nameof(CreateContainerAsync),
                 requestOptions,
                 (trace) => base.CreateContainerAsync(containerProperties, throughput, requestOptions, trace, cancellationToken),
-                (response) => new OpenTelemetryResponse<ContainerProperties>(response));
+                (response) => new OpenTelemetryResponse<ContainerProperties>(response, this.Id));
         }
 
         public override Task<ContainerResponse> CreateContainerAsync(string id,
@@ -43,7 +43,7 @@ public override Task<ContainerResponse> CreateContainerAsync(string id,
                 nameof(CreateContainerAsync),
                 requestOptions,
                 (trace) => base.CreateContainerAsync(id, partitionKeyPath, throughput, requestOptions, trace, cancellationToken),
-                (response) => new OpenTelemetryResponse<ContainerProperties>(response));
+                (response) => new OpenTelemetryResponse<ContainerProperties>(response, this.Id));
         }
 
         public override Task<ContainerResponse> CreateContainerIfNotExistsAsync(
@@ -56,7 +56,7 @@ public override Task<ContainerResponse> CreateContainerIfNotExistsAsync(
                 nameof(CreateContainerIfNotExistsAsync),
                 requestOptions,
                 (trace) => base.CreateContainerIfNotExistsAsync(containerProperties, throughput, requestOptions, trace, cancellationToken),
-                (response) => new OpenTelemetryResponse<ContainerProperties>(response));
+                (response) => new OpenTelemetryResponse<ContainerProperties>(response, this.Id));
         }
 
         public override Task<ContainerResponse> CreateContainerIfNotExistsAsync(
@@ -70,7 +70,7 @@ public override Task<ContainerResponse> CreateContainerIfNotExistsAsync(
                 nameof(CreateContainerIfNotExistsAsync),
                 requestOptions,
                 (trace) => base.CreateContainerIfNotExistsAsync(id, partitionKeyPath, throughput, requestOptions, trace, cancellationToken),
-                (response) => new OpenTelemetryResponse<ContainerProperties>(response));
+                (response) => new OpenTelemetryResponse<ContainerProperties>(response, this.Id));
         }
 
         public override Task<ResponseMessage> CreateContainerStreamAsync(

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/DistributedTracingOptions.cs

@@ -0,0 +1,33 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+    using System.Runtime.CompilerServices;
+
+    /// <summary>
+    /// Open Telemetry Configuration
+    /// It needs to be public once AppInsight is ready
+    /// </summary>
+    internal sealed class DistributedTracingOptions
+    {
+        /// <summary>
+        /// Default Latency threshold for other than query Operation
+        /// </summary>
+        internal static readonly TimeSpan DefaultCrudLatencyThreshold = TimeSpan.FromMilliseconds(100);
+
+        /// <summary>
+        /// Default Latency threshold for QUERY operation
+        /// </summary>
+        internal static readonly TimeSpan DefaultQueryTimeoutThreshold = TimeSpan.FromMilliseconds(500);
+
+        /// <summary>
+        /// Latency Threshold to generate (<see cref="System.Diagnostics.Tracing.EventSource"/>) with Request diagnostics in distributing Tracing.<br></br>
+        /// If it is not set then by default it will generate (<see cref="System.Diagnostics.Tracing.EventSource"/>) for query operation which are taking more than 500 ms and non-query operations taking more than 100 ms.
+        /// </summary>
+        public TimeSpan? DiagnosticsLatencyThreshold { get; set; }
+
+    }
+}

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/Filters/DiagnosticsFilterHelper.cs

@@ -5,23 +5,32 @@
 namespace Microsoft.Azure.Cosmos.Telemetry.Diagnostics
 {
     using System;
-    using System.Net;
+    using Documents;
 
     internal static class DiagnosticsFilterHelper
     {
-        private static readonly TimeSpan latencyThresholdInMs = TimeSpan.FromMilliseconds(250);
-
         /// <summary>
         /// Allow only when either of below is <b>True</b><br></br>
-        /// 1) Latency is not more than 100 ms<br></br>
+        /// 1) Latency is not more than 100/250 (query) ms<br></br>
         /// 3) HTTP status code is not Success<br></br>
         /// </summary>
         /// <returns>true or false</returns>
-        public static bool IsAllowed(
-            TimeSpan latency, 
-            HttpStatusCode statuscode)
+        public static bool IsTracingNeeded(
+            DistributedTracingOptions config,
+            OpenTelemetryAttributes response)
         {
-            return latency > DiagnosticsFilterHelper.latencyThresholdInMs || !statuscode.IsSuccess();
+            TimeSpan latencyThreshold;
+
+            if (config?.DiagnosticsLatencyThreshold != null)
+            {
+                latencyThreshold = config.DiagnosticsLatencyThreshold.Value;
+            }
+            else
+            {
+                latencyThreshold = response.OperationType == OperationType.Query ? DistributedTracingOptions.DefaultQueryTimeoutThreshold : DistributedTracingOptions.DefaultCrudLatencyThreshold;
+            }
+
+            return response.Diagnostics.GetClientElapsedTime() > latencyThreshold || !response.StatusCode.IsSuccess();
         }
     }
 }

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributeKeys.cs

@@ -2,7 +2,7 @@
 // Copyright (c) Microsoft Corporation.  All rights reserved.
 //------------------------------------------------------------
 
-namespace Microsoft.Azure.Cosmos.Telemetry.Diagnostics
+namespace Microsoft.Azure.Cosmos.Telemetry
 {
     internal sealed class OpenTelemetryAttributeKeys
     {

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributes.cs

@@ -2,19 +2,29 @@
 // Copyright (c) Microsoft Corporation.  All rights reserved.
 // ------------------------------------------------------------
 
-namespace Microsoft.Azure.Cosmos.Telemetry.OpenTelemetry
+namespace Microsoft.Azure.Cosmos.Telemetry
 {
     using System.Net;
+    using Microsoft.Azure.Documents;
 
     internal class OpenTelemetryAttributes
     {
-        internal const string NotAvailable = "NA";
-
+        internal const string NotAvailable = "information not available";
+
+        /// <summary>
+        /// For testing purpose only, to make initialization of this class easy 
+        /// </summary>
+        internal OpenTelemetryAttributes()
+        {
+        }
+
         internal OpenTelemetryAttributes(RequestMessage requestMessage)
         {
             this.RequestContentLength = requestMessage?.Headers?.ContentLength ?? OpenTelemetryAttributes.NotAvailable;
             this.ContainerName = requestMessage?.ContainerId ?? OpenTelemetryAttributes.NotAvailable;
             this.DatabaseName = requestMessage?.DatabaseId ?? OpenTelemetryAttributes.NotAvailable;
+
+            this.OperationType = requestMessage?.OperationType ?? OperationType.Invalid;
         }
 
         /// <summary>
@@ -56,5 +66,10 @@ internal OpenTelemetryAttributes(RequestMessage requestMessage)
         /// ItemCount
         /// </summary>
         internal CosmosDiagnostics Diagnostics { get; set; }
+
+        /// <summary>
+        /// OperationType
+        /// </summary>
+        internal OperationType OperationType { get; set; }
     }
 }