Added batchsize and batchioperation name info

Azure · Aug 2, 2024 · 3ed23f5 · 3ed23f5
1 parent aba2b2c
commit 3ed23f5
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 3 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Batch/TransactionalBatchResponse.cs b/Microsoft.Azure.Cosmos/src/Batch/TransactionalBatchResponse.cs
@@ -392,6 +392,27 @@ private static async Task<TransactionalBatchResponse> PopulateFromContentAsync(
             return response;
         }
 
+        internal int GetBatchSize()
+        {
+            return this.Operations.Count;
+        }
+
+        internal OperationType? GetBatchOperationName()
+        {
+            HashSet<OperationType> operationNames = new ();
+            foreach (ItemBatchOperation operation in this.Operations)
+            {
+                operationNames.Add(operation.OperationType);
+            }
+
+            if (operationNames.Count == 1)
+            {
+                return this.Operations[0].OperationType;
+            }
+
+            return null;
+        }
+
         /// <summary>
         /// Disposes the disposable members held by this class.
         /// </summary>

diff --git a/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributeKeys.cs b/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributeKeys.cs
@@ -36,6 +36,7 @@ internal sealed class OpenTelemetryAttributeKeys
         public const string ItemCount = "db.cosmosdb.item_count";
         public const string ActivityId = "db.cosmosdb.activity_id";
         public const string CorrelatedActivityId = "db.cosmosdb.correlated_activity_id";
+        public const string BatchSize = "db.cosmosdb.batch_size";
 
         // Exceptions
         public const string ExceptionType = "exception.type";

diff --git a/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributes.cs b/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributes.cs
@@ -69,5 +69,15 @@ internal OpenTelemetryAttributes(RequestMessage requestMessage)
         /// OperationType
         /// </summary>
         internal Documents.OperationType OperationType { get; set; }
+
+        /// <summary>
+        /// Batch Size
+        /// </summary>
+        internal int? BatchSize { get; set; }
+
+        /// <summary>
+        /// Will have value for homogeneous batch operation and will be null for heterogeneous batch operation
+        /// </summary>
+        internal Documents.OperationType? BatchOperationName { get; set; }
     }
 }
diff --git a/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryCoreRecorder.cs b/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryCoreRecorder.cs
@@ -222,10 +222,20 @@ public void Dispose()
                 OperationType operationType
                     = (this.response == null || this.response?.OperationType == OperationType.Invalid) ? this.operationType : this.response.OperationType;
 
-                this.scope.AddAttribute(OpenTelemetryAttributeKeys.OperationType, Enum.GetName(typeof(OperationType), operationType));
+                string operationName = Enum.GetName(typeof(OperationType), operationType);
+                if (this.response.BatchOperationName != null)
+                {
+                    string batchOpsName = Enum.GetName(typeof(OperationType), this.response.BatchOperationName);
+                    operationName = $"{operationName}.{batchOpsName}";
+                }
+                this.scope.AddAttribute(OpenTelemetryAttributeKeys.OperationType, operationName);
 
                 if (this.response != null)
                 {
+                    if (this.response.BatchSize is not null)
+                    {
+                        this.scope.AddIntegerAttribute(OpenTelemetryAttributeKeys.BatchSize, (int)this.response.BatchSize);
+                    }
                     this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestContentLength, this.response.RequestContentLength);
                     this.scope.AddAttribute(OpenTelemetryAttributeKeys.ResponseContentLength, this.response.ResponseContentLength);
                     this.scope.AddIntegerAttribute(OpenTelemetryAttributeKeys.StatusCode, (int)this.response.StatusCode);

diff --git a/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs b/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs
@@ -22,7 +22,9 @@ internal OpenTelemetryResponse(TransactionalBatchResponse responseMessage)
                   requestMessage: null,
                   subStatusCode: OpenTelemetryResponse.GetHeader(responseMessage)?.SubStatusCode,
                   activityId: OpenTelemetryResponse.GetHeader(responseMessage)?.ActivityId,
-                  correlationId: OpenTelemetryResponse.GetHeader(responseMessage)?.CorrelatedActivityId)
+                  correlationId: OpenTelemetryResponse.GetHeader(responseMessage)?.CorrelatedActivityId,
+                  batchSize: responseMessage.GetBatchSize(),
+                  batchOperationName: responseMessage.GetBatchOperationName())
         {
         }
 
@@ -52,7 +54,9 @@ private OpenTelemetryResponse(
             Documents.SubStatusCodes? subStatusCode,
             string activityId,
             string correlationId,
-            Documents.OperationType operationType = Documents.OperationType.Invalid)
+            Documents.OperationType operationType = Documents.OperationType.Invalid,
+            int? batchSize = null,
+            Documents.OperationType? batchOperationName = null)
             : base(requestMessage)
         {
             this.StatusCode = statusCode;
@@ -64,6 +68,8 @@ private OpenTelemetryResponse(
             this.ActivityId = activityId;
             this.CorrelatedActivityId = correlationId;
             this.OperationType = operationType;
+            this.BatchSize = batchSize;
+            this.BatchOperationName = batchOperationName;
         }
 
         private static string GetPayloadSize(ResponseMessage response)