Fixed issue where CosmosDiagnosticsContext is null when diagnostics are sampled out.

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosDiagnosticsE2ETest.java

@@ -39,6 +39,7 @@
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.fail;
@@ -132,6 +133,7 @@ public void onlyDefaultLogger() {
         CosmosContainer container = this.getContainer(builder);
         executeTestCase(container);
 
+
         // no assertions here - invocations for diagnostics handler are validated above
         // log4j event logging isn't validated in general in unit tests because it is too brittle to do so
         // with custom appender
@@ -152,8 +154,33 @@ public void onlyLoggerWithCustomConfig() {
                     .diagnosticsHandler(CosmosDiagnosticsHandler.DEFAULT_LOGGING_HANDLER)
             );
         CosmosContainer container = this.getContainer(builder);
-        executeTestCase(container);
+        CosmosItemResponse<ObjectNode> response = executeTestCase(container);
+        assertThat(response.getDiagnostics()).isNotNull();
+        assertThat(response.getDiagnostics().getDiagnosticsContext()).isNotNull();
+        // no assertions here - invocations for diagnostics handler are validated above
+        // log4j event logging isn't validated in general in unit tests because it is too brittle to do so
+        // with custom appender
+    }
 
+    @Test(groups = { "fast", "emulator" }, timeOut = TIMEOUT)
+    public void onlyLoggerAlwaysSampledOut() {
+        CosmosClientBuilder builder = this
+            .getClientBuilder()
+            .clientTelemetryConfig(
+                new CosmosClientTelemetryConfig()
+                    .diagnosticsThresholds(
+                        new CosmosDiagnosticsThresholds()
+                            .setPointOperationLatencyThreshold(Duration.ofMillis(100))
+                            .setNonPointOperationLatencyThreshold(Duration.ofMillis(2000))
+                            .setRequestChargeThreshold(100)
+                    )
+                    .diagnosticsHandler(CosmosDiagnosticsHandler.DEFAULT_LOGGING_HANDLER)
+                    .sampleDiagnostics(0)
+            );
+        CosmosContainer container = this.getContainer(builder);
+        CosmosItemResponse<ObjectNode> response = executeTestCase(container);
+        assertThat(response.getDiagnostics()).isNotNull();
+        assertThat(response.getDiagnostics().getDiagnosticsContext()).isNotNull();
         // no assertions here - invocations for diagnostics handler are validated above
         // log4j event logging isn't validated in general in unit tests because it is too brittle to do so
         // with custom appender
@@ -177,10 +204,46 @@ public void onlyCustomLoggerWithCustomConfig() {
                     .diagnosticsHandler(capturingLogger)
             );
         CosmosContainer container = this.getContainer(builder);
-        executeTestCase(container);
+        String id = UUID.randomUUID().toString();
+        ObjectNode doc = getDocumentDefinition(id);
+
+        CosmosDiagnostics diagnostics = executeDocumentOperation(container, OperationType.Create, id, doc);
+        assertThat(diagnostics).isNotNull();
+        assertThat(diagnostics.getDiagnosticsContext()).isNotNull();
+
+        diagnostics = executeDocumentOperation(container, OperationType.Query, id, doc);
+        assertThat(diagnostics).isNotNull();
+        assertThat(diagnostics.getDiagnosticsContext()).isNotNull();
 
         assertThat(capturingLogger.getLoggedMessages()).isNotNull();
-        assertThat(capturingLogger.getLoggedMessages()).hasSize(1);
+        assertThat(capturingLogger.getLoggedMessages()).hasSize(2);
+    }
+
+    @Test(groups = { "fast", "emulator" }, timeOut = TIMEOUT)
+    public void onlyCustomLoggerAlwaysSampledOut() {
+
+        CapturingLogger capturingLogger = new CapturingLogger();
+
+        CosmosClientBuilder builder = this
+            .getClientBuilder()
+            .clientTelemetryConfig(
+                new CosmosClientTelemetryConfig()
+                    .diagnosticsThresholds(
+                        new CosmosDiagnosticsThresholds()
+                            .setPointOperationLatencyThreshold(Duration.ofMillis(100))
+                            .setNonPointOperationLatencyThreshold(Duration.ofMillis(2000))
+                            .setRequestChargeThreshold(100)
+                    )
+                    .diagnosticsHandler(capturingLogger)
+                    .sampleDiagnostics(0)
+            );
+        CosmosContainer container = this.getContainer(builder);
+        CosmosItemResponse<ObjectNode> response = executeTestCase(container);
+
+        assertThat(response.getDiagnostics()).isNotNull();
+        assertThat(response.getDiagnostics().getDiagnosticsContext()).isNotNull();
+        assertThat(capturingLogger.getLoggedMessages()).isNotNull();
+        assertThat(capturingLogger.getLoggedMessages()).hasSize(0);
     }
 
     @Test(groups = { "fast", "emulator" }, timeOut = TIMEOUT)
@@ -368,14 +431,16 @@ public void run() {
         assertThat(cosmosDiagnosticsNode.get("jvmFatalErrorMapperExecutionCount").asLong()).isGreaterThan(0);
     }
 
-    private void executeTestCase(CosmosContainer container) {
+    private CosmosItemResponse<ObjectNode> executeTestCase(CosmosContainer container) {
         String id = UUID.randomUUID().toString();
         CosmosItemResponse<ObjectNode> response = container.createItem(
             getDocumentDefinition(id),
             new PartitionKey(id),
-            null);
+            new CosmosItemRequestOptions());
         assertThat(response).isNotNull();
         assertThat(response.getStatusCode()).isEqualTo(201);
+
+        return response;
     }
 
     private ObjectNode getDocumentDefinition(String documentId) {
@@ -407,47 +472,58 @@ private CosmosContainer getContainer(CosmosClientBuilder builder) {
         return this.client.getDatabase(asyncContainer.getDatabase().getId()).getContainer(asyncContainer.getId());
     }
 
-    private void executeDocumentOperation(
+    private CosmosDiagnostics executeDocumentOperation(
         CosmosContainer cosmosContainer,
         OperationType operationType,
         String createdItemId,
         ObjectNode createdItem) {
+
+        final AtomicReference<CosmosDiagnostics> diagnostics = new AtomicReference<>(null);
+
         switch (operationType) {
             case Query:
                 String query = String.format("SELECT * from c");
                 CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions();
                 queryRequestOptions.setFeedRange(FeedRange.forLogicalPartition(new PartitionKey(createdItemId)));
-                Iterable<FeedResponse<JsonNode>> results = cosmosContainer.queryItems(query, queryRequestOptions, JsonNode.class).iterableByPage();
-                results.forEach(t -> {});
+                Iterable<FeedResponse<JsonNode>> queryResults = cosmosContainer.queryItems(query, queryRequestOptions, JsonNode.class).iterableByPage();
+                queryResults.forEach(t -> diagnostics.set(t.getCosmosDiagnostics()));
                 break;
             case ReadFeed:
                 CosmosChangeFeedRequestOptions changeFeedRequestOptions = CosmosChangeFeedRequestOptions
                     .createForProcessingFromBeginning(FeedRange.forFullRange());
-                cosmosContainer.queryChangeFeed(changeFeedRequestOptions, JsonNode.class).iterableByPage();
+                Iterable<FeedResponse<JsonNode>> changeFeedResults =
+                    cosmosContainer.queryChangeFeed(changeFeedRequestOptions, JsonNode.class).iterableByPage();
+                changeFeedResults.forEach(t -> diagnostics.set(t.getCosmosDiagnostics()));
                 break;
             case Read:
-                cosmosContainer
+                CosmosItemResponse<JsonNode> readResponse = cosmosContainer
                     .readItem(createdItemId, new PartitionKey(createdItemId), JsonNode.class);
+                diagnostics.set(readResponse.getDiagnostics());
                 break;
             case Replace:
-                cosmosContainer
+                CosmosItemResponse<JsonNode> replaceResponse = cosmosContainer
                     .replaceItem(createdItem, createdItemId, new PartitionKey(createdItemId), new CosmosItemRequestOptions());
+                diagnostics.set(replaceResponse.getDiagnostics());
                 break;
             case Delete:
                 try {
-                    cosmosContainer.deleteItem(getDocumentDefinition(UUID.randomUUID().toString()), new CosmosItemRequestOptions());
+                    CosmosItemResponse<Object> deleteResponse = cosmosContainer.deleteItem(getDocumentDefinition(UUID.randomUUID().toString()), new CosmosItemRequestOptions());
+                    diagnostics.set(deleteResponse.getDiagnostics());
                 } catch (CosmosException e) {
                     if (!Exceptions.isNotFound(e)) {
                         throw e;
                     }
                 }
                 break;
             case Create:
-                cosmosContainer.createItem(getDocumentDefinition(UUID.randomUUID().toString()));
+                CosmosItemResponse<JsonNode> createResponse = cosmosContainer.createItem(getDocumentDefinition(UUID.randomUUID().toString()));
+                diagnostics.set(createResponse.getDiagnostics());
                 break;
             default:
                 throw new IllegalArgumentException("The operation type is not supported");
         }
+
+        return diagnostics.get();
     }
 
     private static class CapturingDiagnosticsHandler implements CosmosDiagnosticsHandler {

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -9,8 +9,9 @@
 
 #### Bugs Fixed
 * Suppress exceptions when calling diagnostics handlers. - See [PR 39077](https://github.com/Azure/azure-sdk-for-java/pull/39077)
-* Fixed an issue where no cross region retry for write operations due to channel acquisition timeout - See [PR 39255](https://github.com/Azure/azure-sdk-for-java/pull/39255)
-
+* Fixed an issue where no cross region retry for write operations due to channel acquisition timeout. - See [PR 39255](https://github.com/Azure/azure-sdk-for-java/pull/39255)
+* Fixed incorrect container tag value in metrics. - See [PR 39322](https://github.com/Azure/azure-sdk-for-java/pull/39322)
+* Fixed issue where CosmosDiagnosticsContext is null when diagnostics are sampled out. - See [PR 39352](https://github.com/Azure/azure-sdk-for-java/pull/39352)
 #### Other Changes
 * Only call System.exit in `DiagnosticsProvider` for `Error` scenario. Also add `System.err` for `Error` cases. - See [PR 39077](https://github.com/Azure/azure-sdk-for-java/pull/39077)
 

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

@@ -79,6 +79,8 @@ public final class CosmosDiagnosticsContext {
 
     private Double samplingRateSnapshot;
 
+    private boolean isSampledOut;
+
     private ArrayList<CosmosDiagnosticsRequestInfo> requestInfo = null;
 
     private final Integer sequenceNumber;
@@ -128,6 +130,7 @@ public final class CosmosDiagnosticsContext {
         this.userAgent = userAgent;
         this.connectionMode = connectionMode;
         this.sequenceNumber = sequenceNumber;
+        this.isSampledOut = false;
     }
 
     /**
@@ -556,15 +559,20 @@ void recordOperation(int statusCode,
         }
     }
 
-    void setSamplingRateSnapshot(double samplingRate) {
+    void setSamplingRateSnapshot(double samplingRate, boolean isSampledOut) {
         synchronized (this.spanName) {
             this.samplingRateSnapshot = samplingRate;
+            this.isSampledOut = isSampledOut;
             for (CosmosDiagnostics d : this.diagnostics) {
                 diagAccessor.setSamplingRateSnapshot(d, samplingRate);
             }
         }
     }
 
+    boolean isSampledOut() {
+        return this.isSampledOut;
+    }
+
     String getRequestDiagnostics() {
         ObjectNode ctxNode = mapper.createObjectNode();
 
@@ -1014,9 +1022,9 @@ public String getSpanName(CosmosDiagnosticsContext ctx) {
                     }
 
                     @Override
-                    public void setSamplingRateSnapshot(CosmosDiagnosticsContext ctx, double samplingRate) {
+                    public void setSamplingRateSnapshot(CosmosDiagnosticsContext ctx, double samplingRate, boolean isSampledOut) {
                         checkNotNull(ctx, "Argument 'ctx' must not be null.");
-                        ctx.setSamplingRateSnapshot(samplingRate);
+                        ctx.setSamplingRateSnapshot(samplingRate, isSampledOut);
                     }
 
                     @Override

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

@@ -102,9 +102,9 @@ public double getSamplingRateSnapshot() {
         return samplingRateSnapshot;
     }
 
-    public void setSamplingRateSnapshot(double samplingRateSnapshot) {
+    public void setSamplingRateSnapshot(double samplingRateSnapshot, boolean isSampledOut) {
         if (this.operationState != null) {
-            this.operationState.setSamplingRateSnapshot(samplingRateSnapshot);
+            this.operationState.setSamplingRateSnapshot(samplingRateSnapshot, isSampledOut);
         }
     }
 }