[PREVIEW] AI Integration: Adds open telemetry support for most of the…

… Operations (#3277) * container related operation * database operation * item operations * fix test * code refactor * regain preview tag * code refactor and fix * add unit tests * put preview tag again * fix assertion * fix test * fix code * userinlinecore * permission core * modify assertion message * remove opentelemetry test decorator * code refactor * throw same exception * fix test * fix tests * enable open telemetry for all the operations * added itemcount * code refactor * add baseline tests for open telemetry * add a flag to enable open telemetry featur * remove unit tests * fix test * fix bug * basdeline tests * fix baseline expectations * refactor code for exception * review comments * code refactor * rename enavle open telemetry option and remove cosmos exception * add doc * doc fix * updated preview contracts Co-authored-by: Sourabh Jain <[email protected]>
Azure · Jul 6, 2022 · 66be1a3 · 66be1a3
1 parent 795978a
commit 66be1a3
Show file tree

Hide file tree

Showing 52 changed files with 1,257 additions and 218 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Batch/BatchCore.cs b/Microsoft.Azure.Cosmos/src/Batch/BatchCore.cs
@@ -231,7 +231,8 @@ public override Task<TransactionalBatchResponse> ExecuteAsync(
 
                     this.operations = new List<ItemBatchOperation>();
                     return executor.ExecuteAsync(trace,  cancellationToken);
-                });
+                },
+                (response) => new OpenTelemetryResponse(response));
         }
 
         /// <summary>

diff --git a/Microsoft.Azure.Cosmos/src/ChangeFeed/ChangeFeedIteratorCore.cs b/Microsoft.Azure.Cosmos/src/ChangeFeed/ChangeFeedIteratorCore.cs
@@ -214,11 +214,12 @@ public ChangeFeedIteratorCore(
 
         public override async Task<ResponseMessage> ReadNextAsync(CancellationToken cancellationToken = default)
         {
-            return await this.clientContext.OperationHelperAsync("Change Feed Iterator Read Next Async",
-                                requestOptions: this.changeFeedRequestOptions,
-                                task: (trace) => this.ReadNextInternalAsync(trace, cancellationToken),
-                                traceComponent: TraceComponent.ChangeFeed,
-                                traceLevel: TraceLevel.Info);
+                return await this.clientContext.OperationHelperAsync("Change Feed Iterator Read Next Async",
+                                                requestOptions: this.changeFeedRequestOptions,
+                                                task: (trace) => this.ReadNextInternalAsync(trace, cancellationToken),
+                                                openTelemetry: (response) => new OpenTelemetryResponse(response),
+                                                traceComponent: TraceComponent.ChangeFeed,
+                                                traceLevel: TraceLevel.Info);
         }
 
         public override async Task<ResponseMessage> ReadNextAsync(ITrace trace, CancellationToken cancellationToken = default)

diff --git a/Microsoft.Azure.Cosmos/src/ChangeFeed/ChangeFeedPartitionKeyResultSetIteratorCore.cs b/Microsoft.Azure.Cosmos/src/ChangeFeed/ChangeFeedPartitionKeyResultSetIteratorCore.cs
@@ -94,6 +94,7 @@ public override Task<ResponseMessage> ReadNextAsync(CancellationToken cancellati
             return this.clientContext.OperationHelperAsync("Change Feed Processor Read Next Async",
                                 requestOptions: this.changeFeedOptions,
                                 task: (trace) => this.ReadNextAsync(trace, cancellationToken),
+                                openTelemetry: (response) => new OpenTelemetryResponse(response),
                                 traceComponent: TraceComponent.ChangeFeed,
                                 traceLevel: TraceLevel.Info);
         }

diff --git a/Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/ChangeFeedEstimatorIterator.cs b/Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/ChangeFeedEstimatorIterator.cs
@@ -114,6 +114,7 @@ public override Task<FeedResponse<ChangeFeedProcessorState>> ReadNextAsync(Cance
             return this.monitoredContainer.ClientContext.OperationHelperAsync("Change Feed Estimator Read Next Async",
                                 requestOptions: null,
                                 task: (trace) => this.ReadNextAsync(trace, cancellationToken),
+                                openTelemetry: (response) => new OpenTelemetryResponse<ChangeFeedProcessorState>(response),
                                 traceComponent: TraceComponent.ChangeFeed,
                                 traceLevel: TraceLevel.Info);
         }

diff --git a/...osoft.Azure.Cosmos/src/ChangeFeedProcessor/Exceptions/ChangeFeedProcessorUserException.cs b/...osoft.Azure.Cosmos/src/ChangeFeedProcessor/Exceptions/ChangeFeedProcessorUserException.cs
@@ -6,6 +6,9 @@ namespace Microsoft.Azure.Cosmos
 {
     using System;
     using System.Runtime.Serialization;
+    using global::Azure.Core.Pipeline;
+    using Microsoft.Azure.Cosmos.Telemetry;
+    using Microsoft.Azure.Cosmos.Telemetry.Diagnostics;
 
     /// <summary>
     /// Exception occurred when an operation in an IChangeFeedObserver is running and throws by user code
@@ -53,5 +56,15 @@ public override void GetObjectData(SerializationInfo info, StreamingContext cont
         {
             base.GetObjectData(info, context);
         }
+
+        /// <summary>
+        /// RecordOtelAttributes
+        /// </summary>
+        /// <param name="exception"></param>
+        /// <param name="scope"></param>
+        internal static void RecordOtelAttributes(ChangeFeedProcessorUserException exception, DiagnosticScope scope)
+        {
+            scope.AddAttribute(OpenTelemetryAttributeKeys.ExceptionMessage, exception.Message);
+        }
     }
 }
diff --git a/Microsoft.Azure.Cosmos/src/CosmosClient.cs b/Microsoft.Azure.Cosmos/src/CosmosClient.cs
@@ -727,7 +727,8 @@ public virtual Task<DatabaseResponse> CreateDatabaseAsync(
                         requestOptions: requestOptions,
                         trace: trace,
                         cancellationToken: cancellationToken);
-                });
+                },
+                (response) => new OpenTelemetryResponse<DatabaseProperties>(response));
         }
 
         /// <summary>
@@ -771,7 +772,8 @@ public virtual Task<DatabaseResponse> CreateDatabaseAsync(
                         requestOptions: requestOptions,
                         trace: trace,
                         cancellationToken: cancellationToken);
-                });
+                },
+                (response) => new OpenTelemetryResponse<DatabaseProperties>(response));
         }
 
         /// <summary>
@@ -863,7 +865,8 @@ public virtual Task<DatabaseResponse> CreateDatabaseIfNotExistsAsync(
 
                     return this.ClientContext.ResponseFactory.CreateDatabaseResponse(this.GetDatabase(databaseProperties.Id), readResponseAfterConflict);
                 }
-            });
+            },
+                (response) => new OpenTelemetryResponse<DatabaseProperties>(response));
         }
 
         /// <summary>
@@ -1163,7 +1166,8 @@ public virtual Task<ResponseMessage> CreateDatabaseStreamAsync(
                          requestOptions,
                          trace,
                          cancellationToken);
-                 });
+                 },
+                 (response) => new OpenTelemetryResponse(response));
         }
 
         /// <summary>
@@ -1257,7 +1261,8 @@ internal virtual Task<ResponseMessage> CreateDatabaseStreamAsync(
                         requestOptions,
                         trace,
                         cancellationToken);
-                });
+                },
+                (response) => new OpenTelemetryResponse(response));
         }
 
         private async Task<DatabaseResponse> CreateDatabaseInternalAsync(

diff --git a/Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs b/Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs
@@ -923,5 +923,15 @@ public override bool CanConvert(Type objectType)
                 return objectType == typeof(DateTime);
             }
         }
+
+        /// <summary>
+        /// Enable OpenTelemetry and start emiting activities for each operations
+        /// </summary>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+            bool EnableOpenTelemetry { get; set; }
     }
-}
+}
diff --git a/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs b/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs
@@ -422,6 +422,21 @@ public CosmosClientBuilder WithConsistencyLevel(Cosmos.ConsistencyLevel consiste
             return this;
         }
 
+        /// <summary>
+        /// Enable OpenTelemetry and start emiting activities for each operations
+        /// </summary>
+        /// <returns>The current <see cref="CosmosClientBuilder"/>.</returns>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+            CosmosClientBuilder EnableOpenTelemetry()
+        {
+            this.clientOptions.EnableOpenTelemetry = true;
+            return this;
+        }
+
         /// <summary>
         /// Sets the connection mode to Gateway. This is used by the client when connecting to the Azure Cosmos DB service.
         /// </summary>

diff --git a/Microsoft.Azure.Cosmos/src/Headers/Headers.cs b/Microsoft.Azure.Cosmos/src/Headers/Headers.cs
@@ -246,6 +246,8 @@ internal virtual string EndEpk
             set => this.CosmosMessageHeaders.EndEpk = value;
         }
 
+        internal virtual string ItemCount => this.CosmosMessageHeaders.Get(HttpConstants.HttpHeaders.ItemCount);
+
         /// <summary>
         /// Creates a new instance of <see cref="Headers"/>.
         /// </summary>

diff --git a/Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs b/Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs
@@ -229,7 +229,7 @@ internal override Task<TResult>
             string operationName,
             RequestOptions requestOptions,
             Func<ITrace, Task<TResult>> task,
-            Func<TResult, OpenTelemetryAttributes> openTelemetry = null,
+            Func<TResult, OpenTelemetryAttributes> openTelemetry,
             TraceComponent traceComponent = TraceComponent.Transport,
             Tracing.TraceLevel traceLevel = Tracing.TraceLevel.Info)
         {
@@ -473,7 +473,10 @@ private async Task<TResult> RunWithDiagnosticsHelperAsync<TResult>(
             Func<TResult, OpenTelemetryAttributes> openTelemetry,
             string operationName)
         {
-            using (OpenTelemetryCoreRecorder recorder = OpenTelemetryRecorderFactory.CreateRecorder(operationName))
+            using (OpenTelemetryCoreRecorder recorder = 
+                                OpenTelemetryRecorderFactory.CreateRecorder(
+                                    operationName: operationName, 
+                                    isFeatureEnabled: this.clientOptions.EnableOpenTelemetry))
             using (new ActivityScope(Guid.NewGuid()))
             {
                 try
@@ -489,23 +492,35 @@ private async Task<TResult> RunWithDiagnosticsHelperAsync<TResult>(
                 }
                 catch (OperationCanceledException oe) when (!(oe is CosmosOperationCanceledException))
                 {
-                    recorder.MarkFailed(oe);
-                    throw new CosmosOperationCanceledException(oe, trace);
+                    CosmosOperationCanceledException operationCancelledException = new CosmosOperationCanceledException(oe, trace);
+                    recorder.MarkFailed(operationCancelledException);
+
+                    throw operationCancelledException;
                 }
                 catch (ObjectDisposedException objectDisposed) when (!(objectDisposed is CosmosObjectDisposedException))
                 {
-                    recorder.MarkFailed(objectDisposed);
-                    throw new CosmosObjectDisposedException(
-                        objectDisposed, 
-                        this.client, 
+                    CosmosObjectDisposedException objectDisposedException = new CosmosObjectDisposedException(
+                        objectDisposed,
+                        this.client,
                         trace);
+                    recorder.MarkFailed(objectDisposedException);
+
+                    throw objectDisposedException;
                 }
                 catch (NullReferenceException nullRefException) when (!(nullRefException is CosmosNullReferenceException))
                 {
-                    recorder.MarkFailed(nullRefException);
-                    throw new CosmosNullReferenceException(
+                    CosmosNullReferenceException nullException = new CosmosNullReferenceException(
                         nullRefException,
                         trace);
+                    recorder.MarkFailed(nullException);
+
+                    throw nullException;
+                }
+                catch (Exception ex)
+                {
+                    recorder.MarkFailed(ex);
+
+                    throw;
                 }
             }
         }

diff --git a/Microsoft.Azure.Cosmos/src/Resource/ClientEncryptionKey/ClientEncryptionKeyInlineCore.cs b/Microsoft.Azure.Cosmos/src/Resource/ClientEncryptionKey/ClientEncryptionKeyInlineCore.cs
@@ -31,7 +31,8 @@ public override Task<ClientEncryptionKeyResponse> ReadAsync(
             return this.ClientContext.OperationHelperAsync(
                 nameof(ReadAsync),
                 requestOptions,
-                (trace) => base.ReadAsync(requestOptions, cancellationToken));
+                (trace) => base.ReadAsync(requestOptions, cancellationToken),
+                (response) => new OpenTelemetryResponse<ClientEncryptionKeyProperties>(response));
         }
 
         public override Task<ClientEncryptionKeyResponse> ReplaceAsync(
@@ -42,7 +43,8 @@ public override Task<ClientEncryptionKeyResponse> ReplaceAsync(
             return this.ClientContext.OperationHelperAsync(
                 nameof(ReplaceAsync),
                 requestOptions,
-                (trace) => base.ReplaceAsync(clientEncryptionKeyProperties, requestOptions, cancellationToken));
+                (trace) => base.ReplaceAsync(clientEncryptionKeyProperties, requestOptions, cancellationToken),
+                (response) => new OpenTelemetryResponse<ClientEncryptionKeyProperties>(response));
         }
     }
 }
diff --git a/Microsoft.Azure.Cosmos/src/Resource/Conflict/ConflictsInlineCore.cs b/Microsoft.Azure.Cosmos/src/Resource/Conflict/ConflictsInlineCore.cs
@@ -27,7 +27,8 @@ public override Task<ResponseMessage> DeleteAsync(
             return this.ClientContext.OperationHelperAsync(
                 operationName: nameof(DeleteAsync),
                 requestOptions: null,
-                task: (trace) => base.DeleteAsync(conflict, partitionKey, trace, cancellationToken));
+                task: (trace) => base.DeleteAsync(conflict, partitionKey, trace, cancellationToken),
+                openTelemetry: (response) => new OpenTelemetryResponse(response));
         }
 
         public override FeedIterator GetConflictQueryStreamIterator(
@@ -86,7 +87,8 @@ public override Task<ItemResponse<T>> ReadCurrentAsync<T>(
             return this.ClientContext.OperationHelperAsync(
                 operationName: nameof(ReadCurrentAsync),
                 requestOptions: null,
-                task: (trace) => base.ReadCurrentAsync<T>(cosmosConflict, partitionKey, trace, cancellationToken));
+                task: (trace) => base.ReadCurrentAsync<T>(cosmosConflict, partitionKey, trace, cancellationToken),
+                openTelemetry: (response) => new OpenTelemetryResponse<T>(response));
         }
 
         public override T ReadConflictContent<T>(ConflictProperties cosmosConflict)