[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 <sourabhjain@microsoft.com>
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)