Fixes bug resetting client-level diagnostic thresholds for point read operations (#45488)

Author: FabianMeiswinkel

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

@@ -16,6 +16,7 @@
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import com.azure.cosmos.implementation.apachecommons.lang.tuple.ImmutablePair;
+import com.azure.cosmos.models.CosmosClientTelemetryConfig;
 import com.azure.cosmos.models.CosmosContainerProperties;
 import com.azure.cosmos.models.CosmosItemIdentity;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
@@ -195,6 +196,76 @@ public void readItem() throws Exception {
 
     }
 
+    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    public void readItemWithTimeout() throws Exception {
+        String id = UUID.randomUUID().toString();
+        ObjectNode document = getDocumentDefinition(id, id);
+        container.createItem(document);
+
+        CosmosClientTelemetryConfig clientTelemetryConfig = new CosmosClientTelemetryConfig()
+            .diagnosticsThresholds(
+                new CosmosDiagnosticsThresholds()
+                    .setPointOperationLatencyThreshold(Duration.ofMillis(100))
+            );
+
+        FaultInjectionRule transitTimeout = null;
+        try (CosmosAsyncClient clientWithCustomDiagnosticThresholds = copyCosmosClientBuilder(getClientBuilder())
+            .clientTelemetryConfig(clientTelemetryConfig)
+            .buildAsyncClient()) {
+
+            CosmosAsyncContainer containerWithClientLevelThresholds = clientWithCustomDiagnosticThresholds
+                .getDatabase(container.asyncContainer.getDatabase().getId())
+                .getContainer(container.getId());
+
+            FaultInjectionRuleBuilder ruleBuilder = new FaultInjectionRuleBuilder("extremelyLongResponseDelayRead");
+            FaultInjectionConditionBuilder conditionBuilder = new FaultInjectionConditionBuilder()
+                .operationType(FaultInjectionOperationType.READ_ITEM);
+
+            if (client.asyncClient().getConnectionPolicy().getConnectionMode() != ConnectionMode.DIRECT) {
+                conditionBuilder = conditionBuilder.connectionType(FaultInjectionConnectionType.GATEWAY);
+            } else {
+                conditionBuilder = conditionBuilder.connectionType(FaultInjectionConnectionType.DIRECT);
+            }
+
+            conditionBuilder = conditionBuilder.endpoints(
+                new FaultInjectionEndpointBuilder(FeedRange.forLogicalPartition(new PartitionKey(id)))
+                    .replicaCount(4)
+                    .includePrimary(true)
+                    .build()
+            );
+            FaultInjectionCondition faultInjectionCondition = conditionBuilder.build();
+            FaultInjectionServerErrorResult retryWithResult = FaultInjectionResultBuilders
+                .getResultBuilder(FaultInjectionServerErrorType.RESPONSE_DELAY)
+                .times(1)
+                .delay(Duration.ofMillis(200))
+                .build();
+            transitTimeout = ruleBuilder
+                .condition(faultInjectionCondition)
+                .result(retryWithResult)
+                .duration(Duration.ofSeconds(240))
+                .build();
+
+            CosmosFaultInjectionHelper
+                .configureFaultInjectionRules(containerWithClientLevelThresholds, Arrays.asList(transitTimeout))
+                .block();
+
+            CosmosItemResponse<ObjectNode> itemResponse = containerWithClientLevelThresholds
+                .readItem(id, new PartitionKey(id), ObjectNode.class)
+                .block();
+
+            assertThat(itemResponse.getStatusCode()).isEqualTo(200);
+            assertThat(itemResponse.getDiagnostics()).isNotNull();
+            assertThat(itemResponse.getDiagnostics().getDiagnosticsContext()).isNotNull();
+            assertThat(itemResponse.getDiagnostics().getDiagnosticsContext().isThresholdViolated())
+                .as("Threshold not met - " + itemResponse.getDiagnostics().getDiagnosticsContext().toJson())
+                .isEqualTo(true);
+        } finally {
+            if (transitTimeout != null) {
+                transitTimeout.disable();
+            }
+        }
+    }
+
     @Test(groups = { "fast" }, timeOut = TIMEOUT)
     public void readMany() throws Exception {
         List<CosmosItemIdentity> cosmosItemIdentities = new ArrayList<>();

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

@@ -633,6 +633,13 @@ String getRequestDiagnostics() {
         if (this.subStatusCode != 0) {
             ctxNode.put("subStatus", this.subStatusCode);
         }
+
+        if (this.duration != null) {
+            ctxNode.put("durationInMs",  this.duration.toNanos() / 1_000_000d);
+        } else {
+            ctxNode.put("durationInMs", (Double)null);
+        }
+
         ctxNode.put("RUs", this.totalRequestCharge);
         ctxNode.put("maxRequestSizeInBytes", this.maxRequestSize);
         ctxNode.put("maxResponseSizeInBytes", this.maxResponseSize);

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosItemRequestOptions.java

@@ -98,7 +98,6 @@ public CosmosItemRequestOptions() {
         super();
 
         setPartitionKey(partitionKey);
-        this.thresholds = new CosmosDiagnosticsThresholds();
     }
 
     /**
@@ -540,6 +539,9 @@ public void setThroughputControlGroupName(String throughputControlGroupName) {
      * @return  thresholdForDiagnosticsOnTracerInMS the latency threshold for diagnostics on tracer.
      */
     public Duration getThresholdForDiagnosticsOnTracer() {
+        if (this.thresholds == null) {
+            return Duration.ofMillis(100);
+        }
 
         return thresholdsAccessor.getPointReadLatencyThreshold(this.thresholds);
     }
@@ -554,6 +556,10 @@ public Duration getThresholdForDiagnosticsOnTracer() {
      * @return the CosmosItemRequestOptions
      */
     public CosmosItemRequestOptions setThresholdForDiagnosticsOnTracer(Duration thresholdForDiagnosticsOnTracer) {
+        if (this.thresholds == null) {
+            this.thresholds = new CosmosDiagnosticsThresholds();
+        }
+
         this.thresholds.setPointOperationLatencyThreshold(thresholdForDiagnosticsOnTracer);
 
         return this;