Integrating cosmos diagnostics with open telemetry tracer (Azure#22202)

* Diagnostic integration with open telemetry

* adding test case for query metrics in tracer

* incremental change

* Adding api to pass threshhold for diagnostics

* Correcting Api name

* Documentation improvement

* resolving comments

* Adding pkrange and pkrange id both for querymetrics key

* resolving comments

* Using new core addEvent api with context param

* build error fix

* resolving comments

* fixing spot bugs

* fixing test case

* fixing test for emulator

Author: simplynaveen20

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java

@@ -28,7 +28,6 @@
 import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupFactory;
 import com.azure.cosmos.implementation.throughputControl.config.GlobalThroughputControlGroup;
 import com.azure.cosmos.implementation.throughputControl.config.LocalThroughputControlGroup;
-import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupFactory;
 import com.azure.cosmos.models.CosmosChangeFeedRequestOptions;
 import com.azure.cosmos.models.CosmosConflictProperties;
 import com.azure.cosmos.models.CosmosContainerProperties;
@@ -302,7 +301,8 @@ private <T> Mono<CosmosItemResponse<T>> createItemInternal(T item, CosmosItemReq
                 database.getClient(),
                 ModelBridgeInternal.getConsistencyLevel(options),
                 OperationType.Create,
-                ResourceType.Document);
+                ResourceType.Document,
+                options.getThresholdForDiagnosticsOnTracer());
     }
 
     private <T> Mono<CosmosItemResponse<T>> createItemInternal(T item, CosmosItemRequestOptions options) {
@@ -403,6 +403,7 @@ <T> CosmosPagedFlux<T> readAllItems(CosmosQueryRequestOptions options, Class<T>
             pagedFluxOptions.setTracerAndTelemetryInformation(this.readAllItemsSpanName, database.getId(),
                 this.getId(), OperationType.ReadFeed, ResourceType.Document, this.getDatabase().getClient());
             setContinuationTokenAndMaxItemCount(pagedFluxOptions, options);
+            pagedFluxOptions.setThresholdForDiagnosticsOnTracer(options.getThresholdForDiagnosticsOnTracer());
             return getDatabase().getDocClientWrapper().readDocuments(getLink(), options).map(
                 response -> prepareFeedResponse(response, false, classType));
         });
@@ -546,6 +547,7 @@ <T> Function<CosmosPagedFluxOptions, Flux<FeedResponse<T>>> queryItemsInternalFu
             pagedFluxOptions.setTracerAndTelemetryInformation(spanName, database.getId(),
                 this.getId(), OperationType.Query, ResourceType.Document, this.getDatabase().getClient());
             setContinuationTokenAndMaxItemCount(pagedFluxOptions, cosmosQueryRequestOptions);
+            pagedFluxOptions.setThresholdForDiagnosticsOnTracer(cosmosQueryRequestOptions.getThresholdForDiagnosticsOnTracer());
 
                 return getDatabase().getDocClientWrapper()
                              .queryDocuments(CosmosAsyncContainer.this.getLink(), sqlQuerySpec, cosmosQueryRequestOptions)
@@ -562,6 +564,7 @@ <T> Function<CosmosPagedFluxOptions, Flux<FeedResponse<T>>> queryItemsInternalFu
             pagedFluxOptions.setTracerAndTelemetryInformation(spanName, database.getId(),
                 this.getId(), OperationType.Query, ResourceType.Document, this.getDatabase().getClient());
             setContinuationTokenAndMaxItemCount(pagedFluxOptions, cosmosQueryRequestOptions);
+            pagedFluxOptions.setThresholdForDiagnosticsOnTracer(cosmosQueryRequestOptions.getThresholdForDiagnosticsOnTracer());
 
             return sqlQuerySpecMono.flux()
                 .flatMap(sqlQuerySpec -> getDatabase().getDocClientWrapper()
@@ -973,7 +976,7 @@ public <T> CosmosPagedFlux<T> readAllItems(
         Class<T> classType) {
         final CosmosQueryRequestOptions requestOptions = options == null ? new CosmosQueryRequestOptions() : options;
         requestOptions.setPartitionKey(partitionKey);
-        
+
         return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> {
             pagedFluxOptions.setTracerAndTelemetryInformation(this.readAllItemsSpanName, database.getId(),
                 this.getId(), OperationType.ReadFeed, ResourceType.Document, this.getDatabase().getClient());
@@ -1301,7 +1304,8 @@ private Mono<CosmosItemResponse<Object>> deleteItemInternal(
                 database.getClient(),
                 requestOptions.getConsistencyLevel(),
                 OperationType.Delete,
-                ResourceType.Document);
+                ResourceType.Document,
+                requestOptions.getThresholdForDiagnosticsOnTracer());
     }
 
     private <T> Mono<CosmosItemResponse<T>> replaceItemInternal(
@@ -1327,7 +1331,8 @@ private <T> Mono<CosmosItemResponse<T>> replaceItemInternal(
                 database.getClient(),
                 ModelBridgeInternal.getConsistencyLevel(options),
                 OperationType.Replace,
-                ResourceType.Document);
+                ResourceType.Document,
+                options.getThresholdForDiagnosticsOnTracer());
     }
 
     private <T> Mono<CosmosItemResponse<T>> patchItemInternal(
@@ -1354,7 +1359,8 @@ private <T> Mono<CosmosItemResponse<T>> patchItemInternal(
                 database.getClient(),
                 ModelBridgeInternal.getConsistencyLevel(options),
                 OperationType.Patch,
-                ResourceType.Document);
+                ResourceType.Document,
+                options.getThresholdForDiagnosticsOnTracer());
     }
 
     private <T> Mono<CosmosItemResponse<T>> upsertItemInternal(T item, CosmosItemRequestOptions options, Context context) {
@@ -1378,7 +1384,8 @@ private <T> Mono<CosmosItemResponse<T>> upsertItemInternal(T item, CosmosItemReq
                 database.getClient(),
                 ModelBridgeInternal.getConsistencyLevel(options),
                 OperationType.Upsert,
-                ResourceType.Document);
+                ResourceType.Document,
+                options.getThresholdForDiagnosticsOnTracer());
     }
 
     private <T> Mono<CosmosItemResponse<T>> readItemInternal(
@@ -1401,7 +1408,8 @@ private <T> Mono<CosmosItemResponse<T>> readItemInternal(
                 database.getClient(),
                 requestOptions.getConsistencyLevel(),
                 OperationType.Read,
-                ResourceType.Document);
+                ResourceType.Document,
+                requestOptions.getThresholdForDiagnosticsOnTracer());
     }
 
     Mono<CosmosContainerResponse> read(CosmosContainerRequestOptions options, Context context) {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosDiagnostics.java

@@ -80,6 +80,7 @@ public Duration getDuration() {
 
     /**
      * Regions contacted for this request
+     *
      * @return set of regions contacted for this request
      */
     @Beta(value = Beta.SinceVersion.V4_9_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
@@ -102,7 +103,7 @@ void fillCosmosDiagnostics(ObjectNode parentNode, StringBuilder stringBuilder) {
             }
 
             if (stringBuilder != null) {
-                stringBuilder.append(USER_AGENT_KEY +"=").append(USER_AGENT).append(System.lineSeparator());
+                stringBuilder.append(USER_AGENT_KEY + "=").append(USER_AGENT).append(System.lineSeparator());
                 stringBuilder.append(feedResponseDiagnostics);
             }
         } else {
@@ -135,6 +136,15 @@ private AtomicBoolean isDiagnosticsCapturedInPagedFlux(){
     static {
         ImplementationBridgeHelpers.CosmosDiagnosticsHelper.setCosmosDiagnosticsAccessor(
             new ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnosticsAccessor() {
+                @Override
+                public FeedResponseDiagnostics getFeedResponseDiagnostics(CosmosDiagnostics cosmosDiagnostics) {
+                    if (cosmosDiagnostics != null) {
+                        return cosmosDiagnostics.getFeedResponseDiagnostics();
+                    }
+
+                    return null;
+                }
+
                 @Override
                 public AtomicBoolean isDiagnosticsCapturedInPagedFlux(CosmosDiagnostics cosmosDiagnostics) {
                     return cosmosDiagnostics.isDiagnosticsCapturedInPagedFlux();

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java

@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.implementation;
 
-import com.azure.cosmos.ConnectionMode;
 import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import com.azure.cosmos.implementation.cpu.CpuMemoryMonitor;
@@ -30,7 +29,6 @@
 @JsonSerialize(using = ClientSideRequestStatistics.ClientSideRequestStatisticsSerializer.class)
 public class ClientSideRequestStatistics {
     private static final int MAX_SUPPLEMENTAL_REQUESTS_FOR_TO_STRING = 10;
-    private final DiagnosticsClientContext clientContext;
     private final DiagnosticsClientContext diagnosticsClientContext;
 
     private List<StoreResponseStatistics> responseStatisticsList;
@@ -50,7 +48,6 @@ public class ClientSideRequestStatistics {
 
     public ClientSideRequestStatistics(DiagnosticsClientContext diagnosticsClientContext) {
         this.diagnosticsClientContext = diagnosticsClientContext;
-        this.clientContext = null;
         this.requestStartTimeUTC = Instant.now();
         this.requestEndTimeUTC = Instant.now();
         this.responseStatisticsList = new ArrayList<>();
@@ -68,6 +65,14 @@ public Duration getDuration() {
         return Duration.between(requestStartTimeUTC, requestEndTimeUTC);
     }
 
+    public Instant getRequestStartTimeUTC() {
+        return requestStartTimeUTC;
+    }
+
+    public DiagnosticsClientContext getDiagnosticsClientContext() {
+        return diagnosticsClientContext;
+    }
+
     public void recordResponse(RxDocumentServiceRequest request, StoreResult storeResult) {
         Objects.requireNonNull(request, "request is required and cannot be null.");
         Instant responseTime = Instant.now();
@@ -135,6 +140,7 @@ public void recordGatewayResponse(
                 this.gatewayStatistics.requestCharge = storeResponse
                                                            .getHeaderValue(HttpConstants.HttpHeaders.REQUEST_CHARGE);
                 this.gatewayStatistics.requestTimeline = DirectBridgeInternal.getRequestTimeline(storeResponse);
+                this.gatewayStatistics.partitionKeyRangeId = storeResponse.getPartitionKeyRangeId();
             } else if (exception != null) {
                 this.gatewayStatistics.statusCode = exception.getStatusCode();
                 this.gatewayStatistics.subStatusCode = exception.getSubStatusCode();
@@ -225,6 +231,22 @@ public RetryContext getRetryContext() {
         return retryContext;
     }
 
+    public List<StoreResponseStatistics> getResponseStatisticsList() {
+        return responseStatisticsList;
+    }
+
+    public List<StoreResponseStatistics> getSupplementalResponseStatisticsList() {
+        return supplementalResponseStatisticsList;
+    }
+
+    public Map<String, AddressResolutionStatistics> getAddressResolutionStatistics() {
+        return addressResolutionStatistics;
+    }
+
+    public GatewayStatistics getGatewayStatistics() {
+        return gatewayStatistics;
+    }
+
     public static class StoreResponseStatistics {
         @JsonSerialize(using = StoreResult.StoreResultSerializer.class)
         StoreResult storeResult;
@@ -234,6 +256,14 @@ public static class StoreResponseStatistics {
         ResourceType requestResourceType;
         @JsonSerialize
         OperationType requestOperationType;
+
+        public StoreResult getStoreResult() {
+            return storeResult;
+        }
+
+        public Instant getRequestResponseTimeUTC() {
+            return requestResponseTimeUTC;
+        }
     }
 
     private static class SystemInformation {
@@ -278,20 +308,7 @@ public void serialize(
             generator.writeStringField("requestStartTimeUTC", DiagnosticsInstantSerializer.fromInstant(statistics.requestStartTimeUTC));
             generator.writeStringField("requestEndTimeUTC", DiagnosticsInstantSerializer.fromInstant(statistics.requestEndTimeUTC));
             generator.writeObjectField("responseStatisticsList", statistics.responseStatisticsList);
-            int supplementalResponseStatisticsListCount = statistics.supplementalResponseStatisticsList.size();
-            int initialIndex =
-                Math.max(supplementalResponseStatisticsListCount - MAX_SUPPLEMENTAL_REQUESTS_FOR_TO_STRING, 0);
-            if (initialIndex != 0) {
-                List<StoreResponseStatistics> subList = statistics.supplementalResponseStatisticsList
-                                                            .subList(initialIndex,
-                                                                     supplementalResponseStatisticsListCount);
-                generator.writeObjectField("supplementalResponseStatisticsList", subList);
-            } else {
-                generator
-                    .writeObjectField("supplementalResponseStatisticsList",
-                                      statistics.supplementalResponseStatisticsList);
-            }
-
+            generator.writeObjectField("supplementalResponseStatisticsList", getCappedSupplementalResponseStatisticsList(statistics.supplementalResponseStatisticsList));
             generator.writeObjectField("addressResolutionStatistics", statistics.addressResolutionStatistics);
             generator.writeObjectField("regionsContacted", statistics.regionsContacted);
             generator.writeObjectField("retryContext", statistics.retryContext);
@@ -300,17 +317,7 @@ public void serialize(
             generator.writeObjectField("gatewayStatistics", statistics.gatewayStatistics);
 
             try {
-                SystemInformation systemInformation = new SystemInformation();
-                Runtime runtime = Runtime.getRuntime();
-                long totalMemory = runtime.totalMemory() / 1024;
-                long freeMemory = runtime.freeMemory() / 1024;
-                long maxMemory = runtime.maxMemory() / 1024;
-                systemInformation.usedMemory = totalMemory - freeMemory + " KB";
-                systemInformation.availableMemory = (maxMemory - (totalMemory - freeMemory)) + " KB";
-                systemInformation.availableProcessors = runtime.availableProcessors();
-
-                // TODO: other system related info also can be captured using a similar approach
-                systemInformation.systemCpuLoad = CpuMemoryMonitor.getCpuLoad().toString();
+                SystemInformation systemInformation = fetchSystemInformation();
                 generator.writeObjectField("systemInformation", systemInformation);
             } catch (Exception e) {
                 // Error while evaluating system information, do nothing
@@ -321,7 +328,20 @@ public void serialize(
         }
     }
 
-    private static class AddressResolutionStatistics {
+    public static List<StoreResponseStatistics> getCappedSupplementalResponseStatisticsList(List<StoreResponseStatistics> supplementalResponseStatisticsList) {
+        int supplementalResponseStatisticsListCount = supplementalResponseStatisticsList.size();
+        int initialIndex =
+            Math.max(supplementalResponseStatisticsListCount - MAX_SUPPLEMENTAL_REQUESTS_FOR_TO_STRING, 0);
+        if (initialIndex != 0) {
+            List<StoreResponseStatistics> subList = supplementalResponseStatisticsList
+                .subList(initialIndex,
+                    supplementalResponseStatisticsListCount);
+            return subList;
+        }
+        return supplementalResponseStatisticsList;
+    }
+
+    public static class AddressResolutionStatistics {
         @JsonSerialize(using = DiagnosticsInstantSerializer.class)
         Instant startTimeUTC;
         @JsonSerialize(using = DiagnosticsInstantSerializer.class)
@@ -336,16 +356,21 @@ private static class AddressResolutionStatistics {
         // indicating background addressResolution is still inflight
         @JsonSerialize
         boolean inflightRequest = true;
+
+        public Instant getStartTimeUTC() {
+            return startTimeUTC;
+        }
     }
 
-    private static class GatewayStatistics {
+    public static class GatewayStatistics {
         String sessionToken;
         OperationType operationType;
         ResourceType resourceType;
         int statusCode;
         int subStatusCode;
         String requestCharge;
         RequestTimeline requestTimeline;
+        String partitionKeyRangeId;
 
         public String getSessionToken() {
             return sessionToken;
@@ -374,5 +399,24 @@ public RequestTimeline getRequestTimeline() {
         public ResourceType getResourceType() {
             return resourceType;
         }
+
+        public String getPartitionKeyRangeId() {
+            return partitionKeyRangeId;
+        }
+    }
+
+    public static SystemInformation fetchSystemInformation() {
+        SystemInformation systemInformation = new SystemInformation();
+        Runtime runtime = Runtime.getRuntime();
+        long totalMemory = runtime.totalMemory() / 1024;
+        long freeMemory = runtime.freeMemory() / 1024;
+        long maxMemory = runtime.maxMemory() / 1024;
+        systemInformation.usedMemory = totalMemory - freeMemory + " KB";
+        systemInformation.availableMemory = (maxMemory - (totalMemory - freeMemory)) + " KB";
+        systemInformation.availableProcessors = runtime.availableProcessors();
+
+        // TODO: other system related info also can be captured using a similar approach
+        systemInformation.systemCpuLoad = CpuMemoryMonitor.getCpuLoad().toString();
+        return systemInformation;
     }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java

@@ -7,6 +7,8 @@
 import com.azure.cosmos.CosmosAsyncClient;
 import com.azure.cosmos.util.CosmosPagedFlux;
 
+import java.time.Duration;
+
 /**
  * Specifies paging options for Cosmos Paged Flux implementation.
  * @see CosmosPagedFlux
@@ -23,6 +25,7 @@ public class CosmosPagedFluxOptions {
     private ResourceType resourceType;
     private String serviceEndpoint;
     private CosmosAsyncClient cosmosAsyncClient;
+    private Duration thresholdForDiagnosticsOnTracer;
 
     public CosmosPagedFluxOptions() {}
 
@@ -132,6 +135,30 @@ public String getServiceEndpoint() {
         return serviceEndpoint;
     }
 
+    /**
+     * Gets the thresholdForDiagnosticsOnTracer, if latency on query operation is greater than this
+     * diagnostics will be send to open telemetry exporter as events in tracer span of end to end CRUD api.
+     *
+     * Default is 500 ms.
+     *
+     * @return  thresholdForDiagnosticsOnTracer the latency threshold for diagnostics on tracer.
+     */
+    public Duration getThresholdForDiagnosticsOnTracer() {
+        return thresholdForDiagnosticsOnTracer;
+    }
+
+    /**
+     * Sets the thresholdForDiagnosticsOnTracer, if latency on query operation is greater than this
+     * diagnostics will be send to open telemetry exporter as events in tracer span of end to end CRUD api.
+     *
+     * Default is 500 ms.
+     *
+     * @param thresholdForDiagnosticsOnTracer the latency threshold for diagnostics on tracer.
+     */
+    public void setThresholdForDiagnosticsOnTracer(Duration thresholdForDiagnosticsOnTracer) {
+        this.thresholdForDiagnosticsOnTracer = thresholdForDiagnosticsOnTracer;
+    }
+
     public void setTracerInformation(TracerProvider tracerProvider, String tracerSpanName, String serviceEndpoint, String databaseId) {
         this.databaseId = databaseId;
         this.serviceEndpoint = serviceEndpoint;

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/FeedResponseDiagnostics.java

@@ -33,7 +33,7 @@ public FeedResponseDiagnostics(Map<String, QueryMetrics> queryMetricsMap) {
         this.clientSideRequestStatisticsList = Collections.synchronizedList(new ArrayList<>());
     }
 
-    Map<String, QueryMetrics> getQueryMetricsMap() {
+    public Map<String, QueryMetrics> getQueryMetricsMap() {
         return queryMetricsMap;
     }
 
@@ -88,6 +88,10 @@ public void setDiagnosticsContext(QueryInfo.QueryPlanDiagnosticsContext diagnost
         this.diagnosticsContext = diagnosticsContext;
     }
 
+    public QueryInfo.QueryPlanDiagnosticsContext getQueryPlanDiagnosticsContext() {
+        return diagnosticsContext;
+    }
+
     /**
      * Getter for property 'clientSideRequestStatisticsList'.
      *