Fixed small perf overhead due to NPE for readItem returning 404 (Azure#44008)

* Addressing NPE for readItem returning 404

* Update CHANGELOG.md

* Update SparkE2EGatewayChangeFeedITest.scala

Author: FabianMeiswinkel

sdk/cosmos/azure-cosmos-spark_3_2-12/src/test/scala/com/azure/cosmos/spark/SparkE2EGatewayChangeFeedITest.scala

@@ -116,7 +116,6 @@ class SparkE2EGatewayChangeFeedITest
     assertMetrics(meterRegistry, "cosmos.client.op.latency", expectedToFind = true)
     assertMetrics(meterRegistry, "cosmos.client.system.avgCpuLoad", expectedToFind = true)
     assertMetrics(meterRegistry, "cosmos.client.req.gw", expectedToFind = true)
-    assertMetrics(meterRegistry, "cosmos.client.req.rntbd", expectedToFind = false)
   }
   //scalastyle:on magic.number
   //scalastyle:on multiple.string.literals

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

@@ -9,6 +9,7 @@
 import com.azure.cosmos.implementation.AsyncDocumentClient;
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.ConsoleLoggingRegistryFactory;
+import com.azure.cosmos.implementation.DiagnosticsProvider;
 import com.azure.cosmos.implementation.GlobalEndpointManager;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.InternalObjectNode;
@@ -88,6 +89,8 @@ public class ClientMetricsTest extends BatchTestBase {
     private CosmosClientTelemetryConfig inputClientTelemetryConfig;
     private CosmosMicrometerMetricsOptions inputMetricsOptions;
     private Tag clientCorrelationTag;
+    private long diagnosticHandlerFailuresBaseline;
+    private DiagnosticsProvider diagnosticsProvider;
 
     @Factory(dataProvider = "clientBuildersWithDirectTcpSession")
     public ClientMetricsTest(CosmosClientBuilder clientBuilder) {
@@ -113,7 +116,6 @@ public void beforeTest(
         assertThat(this.meterRegistry).isNull();
 
         this.meterRegistry = ConsoleLoggingRegistryFactory.create(1);
-
         this.inputMetricsOptions = new CosmosMicrometerMetricsOptions()
             .meterRegistry(this.meterRegistry)
             .setMetricCategories(metricCategories)
@@ -137,6 +139,9 @@ public void beforeTest(
             .clientTelemetryConfig(inputClientTelemetryConfig)
             .buildClient();
 
+        this.diagnosticsProvider = ReflectionUtils.getDiagnosticsProvider(client.asyncClient());
+        this.diagnosticHandlerFailuresBaseline = diagnosticsProvider.getDiagnosticHandlerFailuresSnapshot();
+
         assertThat(
             ImplementationBridgeHelpers
                 .CosmosAsyncClientHelper
@@ -162,6 +167,10 @@ public void beforeTest(
     }
 
     public void afterTest() {
+        if (this.diagnosticsProvider != null) {
+            assertThat(this.diagnosticsProvider.getDiagnosticHandlerFailuresSnapshot())
+                .isEqualTo(this.diagnosticHandlerFailuresBaseline);
+        }
         this.container = null;
         CosmosClient clientSnapshot = this.client;
         if (clientSnapshot != null) {
@@ -414,6 +423,38 @@ public void readItem() throws Exception {
         }
     }
 
+    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    public void readNonExistingItem() throws Exception {
+        this.beforeTest(CosmosMetricCategory.DEFAULT);
+        try {
+
+            try {
+                container.readItem(
+                    UUID.randomUUID().toString(),
+                    new PartitionKey(UUID.randomUUID().toString()),
+                    new CosmosItemRequestOptions(),
+                    InternalObjectNode.class);
+            } catch (CosmosException expectedError) {
+                assertThat(expectedError.getStatusCode()).isEqualTo(404);
+                assertThat(expectedError.getSubStatusCode()).isEqualTo(0);
+            }
+
+            this.validateMetrics(
+                Tag.of(
+                    TagName.Operation.toString(), "Document/Read"),
+                Tag.of(TagName.RequestOperationType.toString(), "Document/Read"),
+                0,
+                500
+            );
+
+            Tag queryPlanTag = Tag.of(TagName.RequestOperationType.toString(), "DocumentCollection_QueryPlan");
+            this.assertMetrics("cosmos.client.req.gw", false, queryPlanTag);
+            this.assertMetrics("cosmos.client.req.rntbd", false, queryPlanTag);
+        } finally {
+            this.afterTest();
+        }
+    }
+
     @Test(groups = { "fast" }, timeOut = TIMEOUT)
     public void readManySingleItem() throws Exception {
         this.beforeTest(CosmosMetricCategory.DEFAULT);

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

@@ -911,6 +911,36 @@ public void cosmosAsyncContainer(
         mockTracer.reset();
     }
 
+    @Test(groups = { "fast", "simple" }, timeOut = 10 * TIMEOUT)
+    public void readItemWith404() {
+        ITEM_ID =  "tracerDoc_" + testCaseCount.incrementAndGet();
+        TracerUnderTest mockTracer = Mockito.spy(new TracerUnderTest());
+
+        DiagnosticsProvider provider = createAndInitializeDiagnosticsProvider(
+            mockTracer,
+            false,
+            true,
+            ShowQueryMode.NONE,
+            false,
+            1);
+
+        long diagnosticHandlerFailuresBaseline = provider.getDiagnosticHandlerFailuresSnapshot();
+        try {
+            cosmosAsyncContainer
+                .readItem(ITEM_ID, new PartitionKey(ITEM_ID), null, ObjectNode.class)
+                .block();
+
+            fail("404 Expected");
+        } catch (CosmosException cosmosException) {
+            assertThat(cosmosException.getStatusCode()).isEqualTo(404);
+        }
+
+        assertThat(provider.getDiagnosticHandlerFailuresSnapshot())
+            .isEqualTo(diagnosticHandlerFailuresBaseline);
+
+        mockTracer.reset();
+    }
+
     @Test(groups = { "fast", "simple" }, dataProvider = "traceTestCaseProvider", timeOut = TIMEOUT)
     public void cosmosAsyncScripts(
         boolean useLegacyTracing,

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java

@@ -187,6 +187,13 @@ public static void setDiagnosticsProvider(CosmosAsyncClient cosmosAsyncClient, D
         set(cosmosAsyncClient, tracerProvider, "diagnosticsProvider");
     }
 
+    public static DiagnosticsProvider getDiagnosticsProvider(CosmosAsyncClient cosmosAsyncClient){
+        return get(
+            DiagnosticsProvider.class,
+            cosmosAsyncClient,
+            "diagnosticsProvider");
+    }
+
     public static void setClientTelemetryConfig(CosmosAsyncClient cosmosAsyncClient, CosmosClientTelemetryConfig cfg){
         set(cosmosAsyncClient, cfg, "clientTelemetryConfig");
         AsyncDocumentClient asyncClient = get(

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -7,7 +7,8 @@
 #### Breaking Changes
 
 #### Bugs Fixed
-* Fixes an issue in change feed processor where records are skipped and excessive requests are prefetched. - See [PR 43788](https://github.com/Azure/azure-sdk-for-java/pull/43788)
+* Fixed an issue in change feed processor where records are skipped and excessive requests are prefetched. - See [PR 43788](https://github.com/Azure/azure-sdk-for-java/pull/43788)
+* Fixed small perf overhead due to NPE for readItem returning 404. - See [PR 44008](https://github.com/Azure/azure-sdk-for-java/pull/44008)
 * Perform cross-region retry for `Document` reads when enclosing address requests hit request timeouts (408:10002). - See [PR 43937](https://github.com/Azure/azure-sdk-for-java/pull/43937)
 
 #### Other Changes

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

@@ -107,6 +107,8 @@ public final class DiagnosticsProvider {
     private final CosmosClientTelemetryConfig telemetryConfig;
     private final boolean shouldSystemExitOnError;
 
+    private final AtomicLong diagnosticsHandlerFailures = new AtomicLong(0);
+
     final Supplier<Double> samplingRateSnapshotSupplier;
 
 
@@ -210,6 +212,10 @@ public static Context getContextFromReactorOrNull(ContextView reactorContext) {
         return null;
     }
 
+    public long getDiagnosticHandlerFailuresSnapshot() {
+        return this.diagnosticsHandlerFailures.get();
+    }
+
     /**
      * Stores {@link Context} in Reactor {@link reactor.util.context.Context}.
      *
@@ -538,6 +544,7 @@ private void handleDiagnostics(Context context, CosmosDiagnosticsContext cosmosC
                 try {
                     handler.handleDiagnostics(cosmosCtx, context);
                 } catch (Exception e) {
+                    this.diagnosticsHandlerFailures.incrementAndGet();
                     LOGGER.error("HandledDiagnostics failed. ", e);
                 }
             }
@@ -1698,6 +1705,7 @@ private void handleErrors(Throwable throwable, int systemExitCode) {
         if (throwable instanceof Error) {
             handleFatalError((Error) throwable, systemExitCode);
         } else {
+            this.diagnosticsHandlerFailures.incrementAndGet();
             LOGGER.error("Unexpected exception in DiagnosticsProvider.endSpan. ",  throwable);
             throw new RuntimeException(throwable);
         }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/clienttelemetry/ClientTelemetryMetrics.java

@@ -269,10 +269,10 @@ private static void recordOperation(
             latency,
             maxItemCount == null ? -1 : maxItemCount,
             actualItemCount == null ? -1: actualItemCount,
-            opCountPerEvaluation,
-            opRetriedCountPerEvaluation,
-            globalOpCount,
-            targetMaxMicroBatchSize,
+            opCountPerEvaluation == null ? 0 : opCountPerEvaluation,
+            opRetriedCountPerEvaluation == null ? 0 : opRetriedCountPerEvaluation,
+            globalOpCount == null ? 0 : globalOpCount,
+            targetMaxMicroBatchSize == null ? 0 : targetMaxMicroBatchSize,
             diagnosticsContext,
             contactedRegions
         );