[Preview] AI integration: Adds IsDistributedTracingEnabled flag as public API to enable/disable this feature

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -988,17 +988,29 @@ public override bool CanConvert(Type objectType)
                 return objectType == typeof(DateTime);
             }
         }
-
+        
         /// <summary>
         /// Distributed Tracing Options. <see cref="Microsoft.Azure.Cosmos.DistributedTracingOptions"/>
         /// </summary>
+        /// <remarks> Applicable only when Operation level distributed tracing is enabled through <see cref="Microsoft.Azure.Cosmos.CosmosClientOptions.IsDistributedTracingEnabled"/></remarks>
         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
+        /// Gets or sets the flag to generate operation level <see cref="System.Diagnostics.Activity"/> for methods calls using the Source Name "Azure.Cosmos.Operation".
         /// </summary>
-        internal bool EnableDistributedTracing { get; set; }
+        /// <value>
+        /// The default value is true (for preview package).
+        /// </value>
+        /// <remarks>This flag is there to disable it from source. Please Refer https://opentelemetry.io/docs/instrumentation/net/exporters/ to know more about open telemetry exporters</remarks>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+            bool IsDistributedTracingEnabled { get; set; }
+#if PREVIEW
+        = true;
+#endif
 
     }
 }

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

@@ -424,12 +424,30 @@ public CosmosClientBuilder WithConsistencyLevel(Cosmos.ConsistencyLevel consiste
         }
 
         /// <summary>
-        /// 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>
+        /// Sets whether Distributed Tracing for "Azure.Cosmos.Operation" source is enabled.
         /// </summary>
-        /// <param name="options">Tracing Options <see cref="Microsoft.Azure.Cosmos.DistributedTracingOptions"/></param>
+        /// <param name="isEnabled">Whether <see cref="CosmosClientOptions.IsDistributedTracingEnabled"/> is enabled.</param>
         /// <returns>The current <see cref="CosmosClientBuilder"/>.</returns>
-        internal CosmosClientBuilder WithDistributingTracing(DistributedTracingOptions options)
+#if PREVIEW
+        public
+#else
+        internal
+#endif 
+            CosmosClientBuilder WithDistributedTracing(bool isEnabled = true)
+        {
+            this.clientOptions.IsDistributedTracingEnabled = isEnabled;
+            return this;
+        }
+
+        /// <summary>
+        /// Enables Distributed Tracing with a Configuration ref. <see cref="DistributedTracingOptions"/>
+        /// </summary>
+        /// <param name="options"><see cref="DistributedTracingOptions"/>.</param>
+        /// <returns>The current <see cref="CosmosClientBuilder"/>.</returns>]
+        /// <remarks>Refer https://opentelemetry.io/docs/instrumentation/net/exporters/ to know more about open telemetry exporters</remarks>
+        internal CosmosClientBuilder WithDistributedTracingOptions(DistributedTracingOptions options)
         {
+            this.clientOptions.IsDistributedTracingEnabled = true;
             this.clientOptions.DistributedTracingOptions = options;
 
             return this;

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

@@ -7,8 +7,7 @@ namespace Microsoft.Azure.Cosmos
     using System;
 
     /// <summary>
-    /// Open Telemetry Configuration
-    /// It needs to be public once AppInsight is ready
+    /// Options for configuring the distributed tracing and event tracing
     /// </summary>
     internal sealed class DistributedTracingOptions
     {
@@ -23,9 +22,9 @@ internal sealed class DistributedTracingOptions
         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.
+        /// SDK generates <see cref="System.Diagnostics.Tracing.EventSource"/> (Event Source Name is "Azure-Cosmos-Operation-Request-Diagnostics") with Request Diagnostics String, If Operation level distributed tracing is not disabled i.e. <see cref="Microsoft.Azure.Cosmos.CosmosClientOptions.IsDistributedTracingEnabled"/>
         /// </summary>
-        public TimeSpan? DiagnosticsLatencyThreshold { get; set; }
+        /// <remarks>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.</remarks>
+        public TimeSpan? LatencyThresholdForDiagnosticEvent { get; set; }
     }
 }

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

@@ -22,9 +22,9 @@ public static bool IsTracingNeeded(
         {
             TimeSpan latencyThreshold;
 
-            if (config?.DiagnosticsLatencyThreshold != null)
+            if (config?.LatencyThresholdForDiagnosticEvent != null)
             {
-                latencyThreshold = config.DiagnosticsLatencyThreshold.Value;
+                latencyThreshold = config.LatencyThresholdForDiagnosticEvent.Value;
             }
             else
             {

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

@@ -4,6 +4,7 @@
 
 namespace Microsoft.Azure.Cosmos.Telemetry
 {
+    using System;
     using global::Azure.Core.Pipeline;
 
     /// <summary>
@@ -15,19 +16,19 @@ internal static class OpenTelemetryRecorderFactory
         /// Singleton to make sure we only have one instance of the DiagnosticScopeFactory and pattern matching of listener happens only once
         /// </summary>
         private static DiagnosticScopeFactory ScopeFactory { get; set; } 
-
+        
         public static OpenTelemetryCoreRecorder CreateRecorder(string operationName,
             string containerName,
             string databaseName,
             Documents.OperationType operationType,
             RequestOptions requestOptions, 
             CosmosClientContext clientContext)
         {
-            if (clientContext is { ClientOptions.EnableDistributedTracing: true })
+            if (clientContext is { ClientOptions.IsDistributedTracingEnabled: true })
             {
-                ScopeFactory = new DiagnosticScopeFactory(clientNamespace: OpenTelemetryAttributeKeys.DiagnosticNamespace,
-                    resourceProviderNamespace: OpenTelemetryAttributeKeys.ResourceProviderNamespace,
-                    isActivityEnabled: true);
+                OpenTelemetryRecorderFactory.ScopeFactory ??= new DiagnosticScopeFactory(clientNamespace: OpenTelemetryAttributeKeys.DiagnosticNamespace,
+                        resourceProviderNamespace: OpenTelemetryAttributeKeys.ResourceProviderNamespace,
+                        isActivityEnabled: true);
 
                 // If there is no source then it will return default otherwise a valid diagnostic scope
                 DiagnosticScope scope = OpenTelemetryRecorderFactory

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/TestBaseline/EndToEndTraceWriterBaselineTests.BulkOperationsAsync.xml

@@ -1590,6 +1590,10 @@
             maxRetryWaitTimeOnThrottledRequests: TimeSpan.FromSeconds(1),
             maxRetryAttemptsOnThrottledRequests: 3)
             .WithBulkExecution(true)
+            .WithDistributedTracingOptions(new DistributedTracingOptions()
+            {
+                LatencyThresholdForDiagnosticEvent = TimeSpan.FromMilliseconds(0)
+            })
             .WithTransportClientHandlerFactory(transportClient => new TransportClientWrapper(
                 transportClient,
                 (uri, resourceOperation, request) => TransportClientHelper.ReturnThrottledStoreResponseOnItemOperation(
@@ -1599,12 +1603,6 @@
                     exceptionActivityId,
                     errorMessage))));
 
-    throttleClient.ClientOptions.EnableDistributedTracing = true;
-    throttleClient.ClientOptions.DistributedTracingOptions = new DistributedTracingOptions()
-    {
-        DiagnosticsLatencyThreshold = TimeSpan.FromMilliseconds(0)
-    };
-
     ItemRequestOptions requestOptions = new ItemRequestOptions();
     Container containerWithThrottleException = throttleClient.GetContainer(
         database.Id,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/TestBaseline/EndToEndTraceWriterBaselineTests.PointOperationsExceptionsAsync.xml

@@ -160,7 +160,12 @@
     string errorMessage = "Mock throttle exception" + Guid.NewGuid().ToString();
     Guid exceptionActivityId = Guid.NewGuid();
     using CosmosClient throttleClient = TestCommon.CreateCosmosClient(builder =>
-        builder.WithThrottlingRetryOptions(
+        builder
+        .WithDistributedTracingOptions(new DistributedTracingOptions()
+        {
+            LatencyThresholdForDiagnosticEvent = TimeSpan.FromMilliseconds(0)
+        })
+        .WithThrottlingRetryOptions(
             maxRetryWaitTimeOnThrottledRequests: TimeSpan.FromSeconds(1),
             maxRetryAttemptsOnThrottledRequests: 3)
             .WithTransportClientHandlerFactory(transportClient => new TransportClientWrapper(
@@ -171,13 +176,7 @@
                     request,
                     exceptionActivityId,
                     errorMessage))));
-
-    throttleClient.ClientOptions.EnableDistributedTracing = true;
-    throttleClient.ClientOptions.DistributedTracingOptions = new DistributedTracingOptions()
-    {
-        DiagnosticsLatencyThreshold = TimeSpan.FromMilliseconds(0)
-    };
-
+
     ItemRequestOptions requestOptions = new ItemRequestOptions();
     Container containerWithThrottleException = throttleClient.GetContainer(
         database.Id,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Tracing/AssertActivity.cs

@@ -1,4 +1,8 @@
-namespace Microsoft.Azure.Cosmos.Tracing
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Tracing
 {
     using System.Collections.Generic;
     using System.Diagnostics;
@@ -13,14 +17,16 @@ internal static class AssertActivity
         public static void IsValid(Activity activity)
         {
             Assert.IsTrue(activity.OperationName == activity.DisplayName);
-            Assert.IsNotNull(activity.GetTagItem("db.cosmosdb.connection_mode"));
+
+            Assert.IsFalse(string.IsNullOrEmpty(activity.GetTagItem("db.cosmosdb.connection_mode").ToString()), $"connection mode is emtpy for {activity.OperationName}");
+
             if (activity.GetTagItem("db.cosmosdb.connection_mode").ToString() == ConnectionMode.Gateway.ToString())
             {
-                Assert.AreEqual(ActivityKind.Internal, activity.Kind);
+                Assert.AreEqual(ActivityKind.Internal, activity.Kind, $" Actual Kind is {activity.Kind} but expected is {ActivityKind.Internal} for {activity.OperationName}");
             }
-            else
+            else if (activity.GetTagItem("db.cosmosdb.connection_mode").ToString() == ConnectionMode.Direct.ToString())
             {
-                Assert.AreEqual(ActivityKind.Client, activity.Kind);
+                Assert.AreEqual(ActivityKind.Client, activity.Kind, $" Actual Kind is {activity.Kind} but expected is {ActivityKind.Client} for {activity.OperationName}");
             }
 
             IList<string> expectedTags = new List<string>
@@ -71,8 +77,10 @@ private static void AssertDatabaseAndContainerName(string name, KeyValuePair<str
             IList<string> exceptionsForContainerAttribute = new List<string>
             {
                 "Operation.CreateDatabaseAsync",
+                "Operation.CreateDatabaseIfNotExistsAsync",
                 "Operation.ReadAsync",
-                "Operation.DeleteAsync"
+                "Operation.DeleteAsync",
+                "Operation.DeleteStreamAsync"
             };
 
             if ((tag.Key == OpenTelemetryAttributeKeys.ContainerName && !exceptionsForContainerAttribute.Contains(name)) ||

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Tracing/CustomListener.cs

@@ -18,7 +18,7 @@ namespace Microsoft.Azure.Cosmos.Tests
     /// <summary>
     /// It is a custom listener for Activities and Event. It is used to validate the Activities generated by cosmosDb SDK.
     /// </summary>
-    public class CustomListener :
+    internal class CustomListener :
         EventListener, // Override Event Listener to capture Event source events
         IObserver<KeyValuePair<string, object>>, // Override IObserver to capture Activity events
         IObserver<DiagnosticListener>,
@@ -27,8 +27,8 @@ public class CustomListener :
         private readonly Func<string, bool> sourceNameFilter;
         private readonly string eventName;
 
-        private List<IDisposable> subscriptions = new();
-        private List<ProducedDiagnosticScope> Scopes { get; } = new();
+        private ConcurrentBag<IDisposable> subscriptions = new();
+        private ConcurrentBag<ProducedDiagnosticScope> Scopes { get; } = new();
 
         public static ConcurrentBag<Activity> CollectedActivities { private set; get; } = new();
         private static ConcurrentBag<string> CollectedEvents { set; get; } = new();
@@ -75,7 +75,7 @@ public void OnNext(KeyValuePair<string, object> value)
                 string startSuffix = ".Start";
                 string stopSuffix = ".Stop";
                 string exceptionSuffix = ".Exception";
-
+                
                 if (value.Key.EndsWith(startSuffix))
                 {
                     string name = value.Key[..^startSuffix.Length];
@@ -159,7 +159,6 @@ protected override void OnEventSourceCreated(EventSource eventSource)
         protected override void OnEventWritten(EventWrittenEventArgs eventData)
         {
             StringBuilder builder = new StringBuilder();
-            Console.WriteLine(eventData.Payload[0].ToString());
             builder.Append("<EVENT>")
                    .Append("Ideally, this should contain request diagnostics but request diagnostics is " +
                    "subject to change with each request as it contains few unique id. " +
@@ -182,7 +181,7 @@ public override void Dispose()
                 return;
             }
 
-            List<IDisposable> subscriptions;
+            ConcurrentBag<IDisposable> subscriptions;
             lock (this.Scopes)
             {
                 subscriptions = this.subscriptions;