[Internal] AI Integration: Adds CorrelationId and Activity Id Attributes for query operation

Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs

@@ -73,7 +73,8 @@ internal CosmosQueryResponseMessageHeaders CloneKnownProperties(
                 SubStatusCodeLiteral = this.SubStatusCodeLiteral,
                 ContentType = this.ContentType,
                 QueryMetricsText = QueryMetricsText,
-                IndexUtilizationText = IndexUtilizationText
+                IndexUtilizationText = IndexUtilizationText,
+                CorrelatedActivityId = this.CorrelatedActivityId
             };
         }
 
@@ -108,7 +109,8 @@ internal static CosmosQueryResponseMessageHeaders ConvertToQueryHeaders(
                 SubStatusCodeLiteral = sourceHeaders.SubStatusCodeLiteral ?? (substatusCode.HasValue ? substatusCode.Value.ToString() : null),
                 ContentType = sourceHeaders.ContentType,
                 QueryMetricsText = sourceHeaders.QueryMetricsText,
-                IndexUtilizationText = sourceHeaders.IndexUtilizationText
+                IndexUtilizationText = sourceHeaders.IndexUtilizationText,
+                CorrelatedActivityId = sourceHeaders.CorrelatedActivityId
             };
         }
     }

Microsoft.Azure.Cosmos/src/Headers/Headers.cs

@@ -36,6 +36,15 @@ public virtual string ContinuationToken
             internal set => this.CosmosMessageHeaders.Continuation = value;
         }
 
+        /// <summary>
+        /// Gets or Set the CoorelatedActivityId in the current <see cref="ResponseMessage"/>.
+        /// </summary>
+        internal virtual string CorrelatedActivityId
+        {
+            get => this.CosmosMessageHeaders.Get(HttpConstants.HttpHeaders.CorrelatedActivityId);
+            set => this.CosmosMessageHeaders.Set(HttpConstants.HttpHeaders.CorrelatedActivityId, value);
+        }
+
         /// <summary>
         /// Gets the request charge for this request from the Azure Cosmos DB service.
         /// </summary>

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

@@ -235,6 +235,7 @@ public override async Task<ResponseMessage> ReadNextAsync(ITrace trace, Cancella
                 {
                     RequestCharge = tryGetQueryPage.Result.RequestCharge,
                     ActivityId = tryGetQueryPage.Result.ActivityId,
+                    CorrelatedActivityId = this.correlatedActivityId.ToString(),
                     SubStatusCode = Documents.SubStatusCodes.Unknown
                 };
 

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

@@ -24,7 +24,7 @@ internal sealed class OpenTelemetryAttributeKeys
         public const string ConnectionMode = "db.cosmosdb.connection_mode";
         public const string OperationType = "db.cosmosdb.operation_type";
 
-        // Request Specifics
+        // Request/Response Specifics
         public const string ContainerName = "db.cosmosdb.container";
         public const string RequestContentLength = "db.cosmosdb.request_content_length_bytes";
         public const string ResponseContentLength = "db.cosmosdb.response_content_length_bytes";
@@ -35,6 +35,8 @@ internal sealed class OpenTelemetryAttributeKeys
         public const string RetryCount = "db.cosmosdb.retry_count";
         public const string ItemCount = "db.cosmosdb.item_count";
         public const string RequestDiagnostics = "db.cosmosdb.request_diagnostics";
+        public const string ActivityId = "db.cosmosdb.activity_id";
+        public const string CorrelatedActivityId = "db.cosmosdb.correlated_activity_id";
 
         // Exceptions
         public const string ExceptionType = "exception.type";

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

@@ -54,5 +54,20 @@ internal OpenTelemetryAttributes(RequestMessage requestMessage)
         /// SubStatusCode
         /// </summary>
         internal int SubStatusCode { get; set; }
+
+        /// <summary>
+        /// ActivityId
+        /// </summary>
+        internal string ActivityId { get; set; }
+
+        /// <summary>
+        /// CorrelatedActivityId
+        /// </summary>
+        internal string CorrelatedActivityId { get; set; }
+
+        /// <summary>
+        /// OperationType
+        /// </summary>
+        internal Documents.OperationType OperationType { get; set; }
     }
 }

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

@@ -14,12 +14,13 @@ namespace Microsoft.Azure.Cosmos.Telemetry
     internal struct OpenTelemetryCoreRecorder : IDisposable
     {
         private const string CosmosDb = "cosmosdb";
-
+        
         private readonly DiagnosticScope scope;
         private readonly DistributedTracingOptions config;
 
         private readonly Documents.OperationType operationType;
-
+        private OpenTelemetryAttributes response = null;
+
         internal static IDictionary<Type, Action<Exception, DiagnosticScope>> OTelCompatibleExceptions = new Dictionary<Type, Action<Exception, DiagnosticScope>>()
         {
             { typeof(CosmosNullReferenceException), (exception, scope) => CosmosNullReferenceException.RecordOtelAttributes((CosmosNullReferenceException)exception, scope)},
@@ -41,15 +42,14 @@ public OpenTelemetryCoreRecorder(
             this.config = config;
             this.operationType = operationType;
 
-            if (this.IsEnabled)
+            if (scope.IsEnabled)
             {
                 this.scope.Start();
 
                 this.Record(
                         operationName: operationName,
                         containerName: containerName,
                         databaseName: databaseName,
-                        operationType: operationType,
                         clientContext: clientContext);
             }
         }
@@ -70,21 +70,18 @@ public void Record(string key, string value)
         /// <param name="operationName"></param>
         /// <param name="containerName"></param>
         /// <param name="databaseName"></param>
-        /// <param name="operationType"></param>
         /// <param name="clientContext"></param>
         public void Record(
             string operationName,
             string containerName,
             string databaseName,
-            Documents.OperationType operationType,
             CosmosClientContext clientContext)
         {
             if (this.IsEnabled)
             {
                 this.scope.AddAttribute(OpenTelemetryAttributeKeys.DbOperation, operationName);
                 this.scope.AddAttribute(OpenTelemetryAttributeKeys.DbName, databaseName);
                 this.scope.AddAttribute(OpenTelemetryAttributeKeys.ContainerName, containerName);
-                this.scope.AddAttribute(OpenTelemetryAttributeKeys.OperationType, operationType);
 
                 // Other information
                 this.scope.AddAttribute(OpenTelemetryAttributeKeys.DbSystemName, OpenTelemetryCoreRecorder.CosmosDb);
@@ -106,18 +103,7 @@ public void Record(OpenTelemetryAttributes response)
         {
             if (this.IsEnabled)
             {
-                this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestContentLength, response.RequestContentLength);
-                this.scope.AddAttribute(OpenTelemetryAttributeKeys.ResponseContentLength, response.ResponseContentLength);
-                this.scope.AddAttribute(OpenTelemetryAttributeKeys.StatusCode, (int)response.StatusCode);
-                this.scope.AddAttribute(OpenTelemetryAttributeKeys.SubStatusCode, (int)response.SubStatusCode);
-                this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestCharge, response.RequestCharge);
-                this.scope.AddAttribute(OpenTelemetryAttributeKeys.ItemCount, response.ItemCount);
-
-                if (response.Diagnostics != null)
-                {
-                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.Region, ClientTelemetryHelper.GetContactedRegions(response.Diagnostics.GetContactedRegions()));
-                    CosmosDbEventSource.RecordDiagnosticsForRequests(this.config, this.operationType, response);
-                }
+                this.response = response;
             }
         }
 
@@ -168,6 +154,29 @@ public void Dispose()
         {
             if (this.scope.IsEnabled)
             {
+                Documents.OperationType operationType 
+                    = (this.response == null || this.response?.OperationType == Documents.OperationType.Invalid) ? this.operationType : this.response.OperationType;
+
+                this.scope.AddAttribute(OpenTelemetryAttributeKeys.OperationType, operationType);
+
+                if (this.response != null)
+                {
+                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestContentLength, this.response.RequestContentLength);
+                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.ResponseContentLength, this.response.ResponseContentLength);
+                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.StatusCode, (int)this.response.StatusCode);
+                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.SubStatusCode, (int)this.response.SubStatusCode);
+                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestCharge, this.response.RequestCharge);
+                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.ItemCount, this.response.ItemCount);
+                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.ActivityId, this.response.ActivityId);
+                    this.scope.AddAttribute(OpenTelemetryAttributeKeys.CorrelatedActivityId, this.response.CorrelatedActivityId);
+
+                    if (this.response.Diagnostics != null)
+                    {
+                        this.scope.AddAttribute(OpenTelemetryAttributeKeys.Region, ClientTelemetryHelper.GetContactedRegions(this.response.Diagnostics.GetContactedRegions()));
+                        CosmosDbEventSource.RecordDiagnosticsForRequests(this.config, operationType, this.response);
+                    }
+                }
+
                 this.scope.Dispose();
             }
         }

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

@@ -18,7 +18,9 @@ internal OpenTelemetryResponse(TransactionalBatchResponse responseMessage)
                   diagnostics: responseMessage.Diagnostics,
                   itemCount: responseMessage.Headers?.ItemCount,
                   requestMessage: null,
-                  subStatusCode: (int)responseMessage.Headers?.SubStatusCode)
+                  subStatusCode: (int)responseMessage.Headers?.SubStatusCode,
+                  activityId: responseMessage.Headers?.ActivityId,
+                  correlationId: responseMessage.Headers?.CorrelatedActivityId)
         {
         }
 
@@ -30,7 +32,11 @@ internal OpenTelemetryResponse(ResponseMessage responseMessage)
                   diagnostics: responseMessage.Diagnostics,
                   itemCount: responseMessage.Headers?.ItemCount,
                   requestMessage: responseMessage.RequestMessage,
-                  subStatusCode: (int)responseMessage.Headers?.SubStatusCode)
+                  subStatusCode: (int)responseMessage.Headers?.SubStatusCode,
+                  activityId: responseMessage.Headers?.ActivityId,
+                  correlationId: responseMessage.Headers?.CorrelatedActivityId,
+                  operationType: responseMessage is QueryResponse ? Documents.OperationType.Query : Documents.OperationType.Invalid
+                 )
         {
         }
 
@@ -41,7 +47,10 @@ private OpenTelemetryResponse(
             CosmosDiagnostics diagnostics,
             string itemCount,
             RequestMessage requestMessage,
-            int subStatusCode)
+            int subStatusCode,
+            string activityId,
+            string correlationId,
+            Documents.OperationType operationType = Documents.OperationType.Invalid)
             : base(requestMessage)
         {
             this.StatusCode = statusCode;
@@ -50,6 +59,9 @@ private OpenTelemetryResponse(
             this.Diagnostics = diagnostics;
             this.ItemCount = itemCount; 
             this.SubStatusCode = subStatusCode;
+            this.ActivityId = activityId;
+            this.CorrelatedActivityId = correlationId;
+            this.OperationType = operationType;
         }
 
         private static string GetPayloadSize(ResponseMessage response)

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse{T}.cs

@@ -17,7 +17,10 @@ internal OpenTelemetryResponse(FeedResponse<T> responseMessage)
                diagnostics: responseMessage.Diagnostics,
                itemCount: responseMessage.Headers?.ItemCount,
                requestMessage: responseMessage.RequestMessage,
-               subStatusCode: (int)responseMessage.Headers?.SubStatusCode)
+               subStatusCode: (int)responseMessage.Headers?.SubStatusCode,
+               activityId: responseMessage.Headers?.ActivityId,
+               correlatedActivityId: responseMessage.Headers?.CorrelatedActivityId,
+               operationType: responseMessage is QueryResponse<T> ? Documents.OperationType.Query : Documents.OperationType.Invalid)
         {
         }
 
@@ -29,7 +32,10 @@ internal OpenTelemetryResponse(Response<T> responseMessage)
                   diagnostics: responseMessage.Diagnostics,
                   itemCount: responseMessage.Headers?.ItemCount,
                   requestMessage: responseMessage.RequestMessage,
-                  subStatusCode: (int)responseMessage.Headers?.SubStatusCode)
+                  subStatusCode: (int)responseMessage.Headers?.SubStatusCode,
+                  activityId: responseMessage.Headers?.ActivityId,
+                  correlatedActivityId: responseMessage.Headers?.CorrelatedActivityId,
+                  operationType: responseMessage is QueryResponse ? Documents.OperationType.Query : Documents.OperationType.Invalid)
         {
         }
 
@@ -40,7 +46,10 @@ private OpenTelemetryResponse(
            CosmosDiagnostics diagnostics,
            string itemCount,
            RequestMessage requestMessage,
-           int subStatusCode)
+           int subStatusCode,
+           string activityId,
+           string correlatedActivityId,
+           Documents.OperationType operationType)
            : base(requestMessage)
         {
             this.StatusCode = statusCode;
@@ -49,6 +58,9 @@ private OpenTelemetryResponse(
             this.Diagnostics = diagnostics;
             this.ItemCount = itemCount;
             this.SubStatusCode = subStatusCode;
+            this.ActivityId = activityId;
+            this.CorrelatedActivityId = correlatedActivityId;
+            this.OperationType = operationType;
         }
     }
 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/BaselineTests.cs

@@ -116,17 +116,13 @@ public void ExecuteTestSuite(IEnumerable<TInput> inputs, [CallerMemberName] stri
             string outputText =
                 Regex.Replace(
                     Regex.Replace(
-                        Regex.Replace(
                             File.ReadAllText(outputPath), @"\s+", string.Empty), 
-                    @"<ATTRIBUTE-VALUE>[\w\W]*?</ATTRIBUTE-VALUE>", string.Empty),
                 @"<OPERATION>[\w\W]*?</OPERATION>", string.Empty); // in changefeed test in was changing
 
             string baselineText =
                 Regex.Replace(
                     Regex.Replace(
-                        Regex.Replace(
                             File.ReadAllText(baselinePath), @"\s+", string.Empty), 
-                    @"<ATTRIBUTE-VALUE>[\w\W]*?</ATTRIBUTE-VALUE>", string.Empty),
                 @"<OPERATION>[\w\W]*?</OPERATION>", string.Empty);
 
             int commonPrefixLength = 0;

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

@@ -129,7 +129,7 @@
     }
   ]
 }]]></Json>
-      <OTelActivities><ACTIVITY><OPERATION>Operation.ExecuteAsync</OPERATION><ATTRIBUTE-KEY>kind</ATTRIBUTE-KEY><ATTRIBUTE-KEY>az.namespace</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.operation</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.name</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.container</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.operation_type</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.system</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.machine_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>net.peer.name</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.client_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.user_agent</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.connection_mode</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.request_content_length_bytes</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.response_content_length_bytes</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.status_code</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.sub_status_code</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.request_charge</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.item_count</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.regions_contacted</ATTRIBUTE-KEY></ACTIVITY>
+      <OTelActivities><ACTIVITY><OPERATION>Operation.ExecuteAsync</OPERATION><ATTRIBUTE-KEY>kind</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>client</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>az.namespace</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>Microsoft.DocumentDB</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.operation</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>ExecuteAsync</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.name</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.container</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.system</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>cosmosdb</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.cosmosdb.machine_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>net.peer.name</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>127.0.0.1</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.cosmosdb.client_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.user_agent</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.connection_mode</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>Direct</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.cosmosdb.operation_type</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>Replace</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.cosmosdb.request_content_length_bytes</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.response_content_length_bytes</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.status_code</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.sub_status_code</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.request_charge</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.item_count</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.activity_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.correlated_activity_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.regions_contacted</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>South Central US</ATTRIBUTE-VALUE></ACTIVITY>
 <EVENT>Ideally, this should contain request diagnostics but request diagnostics is subject to change with each request as it contains few unique id. So just putting this tag with this static text to make sure event is getting generated for each test.</EVENT>
 </OTelActivities>
     </Output>