Client Level Diagnostics for Spring SDK

sdk/spring/azure-spring-data-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 3.37.0-beta.1 (Unreleased)
 
 #### Features Added
+* Added Diagnostic Threshold configuration support. In order to use diagnostics they must be enabled on the supplied `CosmosClientBuilder` as part of configuration. - See [PR 35546](https://github.com/Azure/azure-sdk-for-java/pull/35546)
 
 #### Breaking Changes
 

sdk/spring/azure-spring-data-cosmos/README.md

@@ -124,6 +124,7 @@ Set `maxDegreeOfParallelism` flag to an integer in application.properties to all
 Set `maxBufferedItemCount` flag to an integer in application.properties to allow the user to set the max number of items that can be buffered during parallel query execution; if set to less than 0, the system automatically decides the number of items to buffer.
 NOTE: Setting this to a very high value can result in high memory consumption.
 Set `responseContinuationTokenLimitInKb` flag to an integer in application.properties to allow the user to limit the length of the continuation token in the query response. The continuation token contains both required and optional fields. The required fields are necessary for resuming the execution from where it was stoped. The optional fields may contain serialized index lookup work that was done but not yet utilized. This avoids redoing the work again in subsequent continuations and hence improve the query performance. Setting the maximum continuation size to 1KB, the Azure Cosmos DB service will only serialize required fields. Starting from 2KB, the Azure Cosmos DB service would serialize as much as it could fit till it reaches the maximum specified size.
+Set `pointOperationLatencyThresholdInMS`, `nonPointOperationLatencyThresholdInMS`, `requestChargeThresholdInRU` and `payloadSizeThresholdInBytes` to enable diagnostics at the client level when these thresholds are exceeded.
 
 ```java readme-sample-AppConfiguration
 @Configuration
@@ -156,6 +157,18 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
     @Value("${azure.cosmos.responseContinuationTokenLimitInKb}")
     private int responseContinuationTokenLimitInKb;
 
+    @Value("${azure.cosmos.diagnosticsThresholds.pointOperationLatencyThresholdInMS}")
+    private int pointOperationLatencyThresholdInMS;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.nonPointOperationLatencyThresholdInMS}")
+    private int nonPointOperationLatencyThresholdInMS;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.requestChargeThresholdInRU}")
+    private int requestChargeThresholdInRU;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.payloadSizeThresholdInBytes}")
+    private int payloadSizeThresholdInBytes;
+
     private AzureKeyCredential azureKeyCredential;
 
     @Bean
@@ -166,7 +179,17 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
         return new CosmosClientBuilder()
             .endpoint(uri)
             .credential(azureKeyCredential)
-            .directMode(directConnectionConfig, gatewayConnectionConfig);
+            .directMode(directConnectionConfig, gatewayConnectionConfig)
+            .clientTelemetryConfig(
+                new CosmosClientTelemetryConfig()
+                    .diagnosticsThresholds(
+                        new CosmosDiagnosticsThresholds()
+                            .setNonPointOperationLatencyThreshold(Duration.ofMillis(nonPointOperationLatencyThresholdInMS))
+                            .setPointOperationLatencyThreshold(Duration.ofMillis(pointOperationLatencyThresholdInMS))
+                            .setPayloadSizeThreshold(payloadSizeThresholdInBytes)
+                            .setRequestChargeThreshold(requestChargeThresholdInRU)
+                    )
+                    .diagnosticsHandler(CosmosDiagnosticsHandler.DEFAULT_LOGGING_HANDLER));
     }
 
     @Override
@@ -201,6 +224,7 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
 ```
 ### Customizing Configuration
 You can customize `DirectConnectionConfig` or `GatewayConnectionConfig` or both and provide them to `CosmosClientBuilder` bean to customize `CosmosAsyncClient`
+You can customize `pointOperationLatencyThresholdInMS`, `nonPointOperationLatencyThresholdInMS`, `requestChargeThresholdInRU` and `payloadSizeThresholdInBytes` to customize the thresholds for diagnostic logging when combined with `CosmosDiagnosticsHandler` which enables diagnostic logging with the default thresholds when added to the `CosmosClientBuilder`.
 
 ```java readme-sample-AppConfigurationCodeSnippet
 @Bean
@@ -210,7 +234,17 @@ public CosmosClientBuilder getCosmosClientBuilder() {
     GatewayConnectionConfig gatewayConnectionConfig = new GatewayConnectionConfig();
     return new CosmosClientBuilder()
         .endpoint(uri)
-        .directMode(directConnectionConfig, gatewayConnectionConfig);
+        .directMode(directConnectionConfig, gatewayConnectionConfig)
+        .clientTelemetryConfig(
+                new CosmosClientTelemetryConfig()
+                    .diagnosticsThresholds(
+                        new CosmosDiagnosticsThresholds()
+                            .setNonPointOperationLatencyThreshold(Duration.ofMillis(nonPointOperationLatencyThresholdInMS))
+                            .setPointOperationLatencyThreshold(Duration.ofMillis(pointOperationLatencyThresholdInMS))
+                            .setPayloadSizeThreshold(payloadSizeThresholdInBytes)
+                            .setRequestChargeThreshold(requestChargeThresholdInRU)
+                    )
+                    .diagnosticsHandler(CosmosDiagnosticsHandler.DEFAULT_LOGGING_HANDLER));
 }
 
 @Override
@@ -228,6 +262,14 @@ public CosmosConfig cosmosConfig() {
 By default, `@EnableCosmosRepositories` will scan the current package for any interfaces that extend one of Spring Data's repository interfaces.
 Use it to annotate your Configuration class to scan a different root package by `@EnableCosmosRepositories(basePackageClass=UserRepository.class)` if your project layout has multiple projects.
 
+#### Enabling logging diagnostics to Azure Application Insights with JavaAgent
+
+Diagnostics can be enabled by passing the JavaAgent with your application like below. This will enable logging with the default thresholds. The '-javaagent' must be passed before the '-jar'.
+
+```commandline
+java -javaagent:"<path-to-applicationinsights-agent-jar>" -jar <myapp.jar>
+```
+
 #### Using database provisioned throughput
 
 Cosmos supports both [container](https://docs.microsoft.com/azure/cosmos-db/sql/how-to-provision-container-throughput)
@@ -689,7 +731,7 @@ public class SecondaryDatasourceConfiguration {
     public CosmosAsyncClient getCosmosAsyncClient(@Qualifier("secondary") CosmosProperties secondaryProperties) {
         return CosmosFactory.createCosmosAsyncClient(new CosmosClientBuilder()
             .key(secondaryProperties.getKey())
-            .endpoint(secondaryProperties.getUri()));
+            .endpoint(secondaryProperties.getUri());
     }
 
     @Bean("secondaryCosmosConfig")
@@ -727,7 +769,7 @@ public class SecondaryDatasourceConfiguration {
 public CosmosAsyncClient getCosmosAsyncClient(@Qualifier("secondary") CosmosProperties secondaryProperties) {
     return CosmosFactory.createCosmosAsyncClient(new CosmosClientBuilder()
         .key(secondaryProperties.getKey())
-        .endpoint(secondaryProperties.getUri()));
+        .endpoint(secondaryProperties.getUri());
 }
 
 @Bean("secondaryCosmosConfig")

sdk/spring/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/AppConfiguration.java

@@ -4,8 +4,11 @@
 
 import com.azure.core.credential.AzureKeyCredential;
 import com.azure.cosmos.CosmosClientBuilder;
+import com.azure.cosmos.CosmosDiagnosticsHandler;
+import com.azure.cosmos.CosmosDiagnosticsThresholds;
 import com.azure.cosmos.DirectConnectionConfig;
 import com.azure.cosmos.GatewayConnectionConfig;
+import com.azure.cosmos.models.CosmosClientTelemetryConfig;
 import com.azure.spring.data.cosmos.config.AbstractCosmosConfiguration;
 import com.azure.spring.data.cosmos.config.CosmosConfig;
 import com.azure.spring.data.cosmos.core.ResponseDiagnostics;
@@ -18,6 +21,8 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.lang.Nullable;
 
+import java.time.Duration;
+
 // BEGIN: readme-sample-AppConfiguration
 @Configuration
 @EnableCosmosRepositories
@@ -49,6 +54,19 @@ public class AppConfiguration extends AbstractCosmosConfiguration {
     @Value("${azure.cosmos.responseContinuationTokenLimitInKb}")
     private int responseContinuationTokenLimitInKb;
 
+    @Value("${azure.cosmos.diagnosticsThresholds.pointOperationLatencyThresholdInMS}")
+    private int pointOperationLatencyThresholdInMS;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.nonPointOperationLatencyThresholdInMS}")
+    private int nonPointOperationLatencyThresholdInMS;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.requestChargeThresholdInRU}")
+    private int requestChargeThresholdInRU;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.payloadSizeThresholdInBytes}")
+    private int payloadSizeThresholdInBytes;
+
+
     private AzureKeyCredential azureKeyCredential;
 
     @Bean
@@ -59,7 +77,17 @@ public CosmosClientBuilder getCosmosClientBuilder() {
         return new CosmosClientBuilder()
             .endpoint(uri)
             .credential(azureKeyCredential)
-            .directMode(directConnectionConfig, gatewayConnectionConfig);
+            .directMode(directConnectionConfig, gatewayConnectionConfig)
+            .clientTelemetryConfig(
+                new CosmosClientTelemetryConfig()
+                    .diagnosticsThresholds(
+                        new CosmosDiagnosticsThresholds()
+                            .setNonPointOperationLatencyThreshold(Duration.ofMillis(nonPointOperationLatencyThresholdInMS))
+                            .setPointOperationLatencyThreshold(Duration.ofMillis(pointOperationLatencyThresholdInMS))
+                            .setPayloadSizeThreshold(payloadSizeThresholdInBytes)
+                            .setRequestChargeThreshold(requestChargeThresholdInRU)
+                    )
+                    .diagnosticsHandler(CosmosDiagnosticsHandler.DEFAULT_LOGGING_HANDLER));
     }
 
     @Override

sdk/spring/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/AppConfigurationCodeSnippet.java

@@ -3,8 +3,11 @@
 package com.azure.spring.data.cosmos;
 
 import com.azure.cosmos.CosmosClientBuilder;
+import com.azure.cosmos.CosmosDiagnosticsHandler;
+import com.azure.cosmos.CosmosDiagnosticsThresholds;
 import com.azure.cosmos.DirectConnectionConfig;
 import com.azure.cosmos.GatewayConnectionConfig;
+import com.azure.cosmos.models.CosmosClientTelemetryConfig;
 import com.azure.spring.data.cosmos.config.AbstractCosmosConfiguration;
 import com.azure.spring.data.cosmos.config.CosmosConfig;
 import com.azure.spring.data.cosmos.core.ResponseDiagnostics;
@@ -17,6 +20,8 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.lang.Nullable;
 
+import java.time.Duration;
+
 @Configuration
 @EnableCosmosRepositories
 public class AppConfigurationCodeSnippet extends AbstractCosmosConfiguration {
@@ -48,6 +53,18 @@ public class AppConfigurationCodeSnippet extends AbstractCosmosConfiguration {
     @Value("${azure.cosmos.responseContinuationTokenLimitInKb}")
     private int responseContinuationTokenLimitInKb;
 
+    @Value("${azure.cosmos.diagnosticsThresholds.pointOperationLatencyThresholdInMS}")
+    private int pointOperationLatencyThresholdInMS;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.nonPointOperationLatencyThresholdInMS}")
+    private int nonPointOperationLatencyThresholdInMS;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.requestChargeThresholdInRU}")
+    private int requestChargeThresholdInRU;
+
+    @Value("${azure.cosmos.diagnosticsThresholds.payloadSizeThresholdInBytes}")
+    private int payloadSizeThresholdInBytes;
+
     // BEGIN: readme-sample-AppConfigurationCodeSnippet
     @Bean
     public CosmosClientBuilder getCosmosClientBuilder() {
@@ -56,7 +73,17 @@ public CosmosClientBuilder getCosmosClientBuilder() {
         GatewayConnectionConfig gatewayConnectionConfig = new GatewayConnectionConfig();
         return new CosmosClientBuilder()
             .endpoint(uri)
-            .directMode(directConnectionConfig, gatewayConnectionConfig);
+            .directMode(directConnectionConfig, gatewayConnectionConfig)
+            .clientTelemetryConfig(
+                new CosmosClientTelemetryConfig()
+                    .diagnosticsThresholds(
+                        new CosmosDiagnosticsThresholds()
+                            .setNonPointOperationLatencyThreshold(Duration.ofMillis(nonPointOperationLatencyThresholdInMS))
+                            .setPointOperationLatencyThreshold(Duration.ofMillis(pointOperationLatencyThresholdInMS))
+                            .setPayloadSizeThreshold(payloadSizeThresholdInBytes)
+                            .setRequestChargeThreshold(requestChargeThresholdInRU)
+                    )
+                    .diagnosticsHandler(CosmosDiagnosticsHandler.DEFAULT_LOGGING_HANDLER));
     }
 
     @Override

sdk/spring/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/CosmosProperties.java

@@ -28,6 +28,14 @@ public class CosmosProperties {
 
     private int responseContinuationTokenLimitInKb;
 
+    private int pointOperationLatencyThresholdInMS;
+
+    private int nonPointOperationLatencyThresholdInMS;
+
+    private int requestChargeThresholdInRU;
+
+    private int payloadSizeThresholdInBytes;
+
     public String getUri() {
         return uri;
     }
@@ -72,23 +80,55 @@ public int getMaxDegreeOfParallelism() {
         return maxDegreeOfParallelism;
     }
 
+    public void setMaxDegreeOfParallelism(int maxDegreeOfParallelism) {
+        this.maxDegreeOfParallelism = maxDegreeOfParallelism;
+    }
+
     public int getMaxBufferedItemCount() {
         return maxBufferedItemCount;
     }
 
+    public void setMaxBufferedItemCount(int maxBufferedItemCount) {
+        this.maxBufferedItemCount = maxBufferedItemCount;
+    }
+
     public int getResponseContinuationTokenLimitInKb() {
         return responseContinuationTokenLimitInKb;
     }
 
-    public void setMaxDegreeOfParallelism(int maxDegreeOfParallelism) {
-        this.maxDegreeOfParallelism = maxDegreeOfParallelism;
+    public void setResponseContinuationTokenLimitInKb(int responseContinuationTokenLimitInKb) {
+        this.responseContinuationTokenLimitInKb = responseContinuationTokenLimitInKb;
     }
 
-    public void setMaxBufferedItemCount(int maxBufferedItemCount) {
-        this.maxBufferedItemCount = maxBufferedItemCount;
+    public int getPointOperationLatencyThresholdInMS() {
+        return pointOperationLatencyThresholdInMS;
     }
 
-    public void setResponseContinuationTokenLimitInKb(int responseContinuationTokenLimitInKb) {
-        this.responseContinuationTokenLimitInKb = responseContinuationTokenLimitInKb;
+    public void setPointOperationLatencyThresholdInMS(int pointOperationLatencyThresholdInMS) {
+        this.pointOperationLatencyThresholdInMS = pointOperationLatencyThresholdInMS;
+    }
+
+    public int getNonPointOperationLatencyThresholdInMS() {
+        return nonPointOperationLatencyThresholdInMS;
+    }
+
+    public void setNonPointOperationLatencyThresholdInMS(int nonPointOperationLatencyThresholdInMS) {
+        this.nonPointOperationLatencyThresholdInMS = nonPointOperationLatencyThresholdInMS;
+    }
+
+    public int getRequestChargeThresholdInRU() {
+        return requestChargeThresholdInRU;
+    }
+
+    public void setRequestChargeThresholdInRU(int requestChargeThresholdInRU) {
+        this.requestChargeThresholdInRU = requestChargeThresholdInRU;
+    }
+
+    public int getPayloadSizeThresholdInBytes() {
+        return payloadSizeThresholdInBytes;
+    }
+
+    public void setPayloadSizeThresholdInBytes(int payloadSizeThresholdInBytes) {
+        this.payloadSizeThresholdInBytes = payloadSizeThresholdInBytes;
     }
 }

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplateIT.java

@@ -59,6 +59,7 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
@@ -90,6 +91,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 @RunWith(SpringJUnit4ClassRunner.class)
@@ -138,6 +140,8 @@ public class CosmosTemplateIT {
     private static CosmosEntityInformation<Person, String> personInfo;
     private static String containerName;
 
+    private MappingCosmosConverter cosmosConverter;
+
     private Person insertedPerson;
 
     private BasicItem pointReadItem;
@@ -174,7 +178,7 @@ private CosmosTemplate createCosmosTemplate(CosmosConfig config, String dbName)
         final CosmosFactory cosmosFactory = new CosmosFactory(client, dbName);
         final CosmosMappingContext mappingContext = new CosmosMappingContext();
         mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class));
-        final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null);
+        cosmosConverter = new MappingCosmosConverter(mappingContext, null);
         return new CosmosTemplate(cosmosFactory, config, cosmosConverter);
     }
 
@@ -222,6 +226,17 @@ public void testFindAll() {
     }
 
     @Test
+    public void testDiagnosticsLogged() {
+        TestRepositoryConfig.capturingLogger.loggedMessages = new ArrayList<>();
+        cosmosTemplate.insert(TEST_PERSON_2,
+            new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2)));
+        cosmosTemplate.insert(TEST_PERSON_3,
+            new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_3)));
+        final List<Person> result = TestUtils.toList(cosmosTemplate.findAll(Person.class.getSimpleName(),
+            Person.class));
+        assertTrue(TestRepositoryConfig.capturingLogger.getLoggedMessages().size() > 0);
+    }
+
     public void testFindByIdPointRead() {
         final BasicItem result = cosmosTemplate.findById(BasicItem.class.getSimpleName(),
             BASIC_ITEM.getId(), BasicItem.class);