Update spring cloud azure cosmos proxy and retry configuration (#25766)

Azure · Dec 3, 2021 · 97ee351 · 97ee351
1 parent afca541
commit 97ee351
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 5 deletions.
diff --git a/...in/java/com/azure/spring/cloud/autoconfigure/cosmos/properties/AzureCosmosProperties.java b/...in/java/com/azure/spring/cloud/autoconfigure/cosmos/properties/AzureCosmosProperties.java
@@ -54,11 +54,16 @@ public class AzureCosmosProperties extends AbstractAzureServiceCP implements Cos
 
     private final List<String> preferredRegions = new ArrayList<>();
 
+    @NestedConfigurationProperty
     private final ThrottlingRetryOptions throttlingRetryOptions = new ThrottlingRetryOptions();
 
     private ConsistencyLevel consistencyLevel;
     private ConnectionMode connectionMode;
+
+    @NestedConfigurationProperty
     private final GatewayConnectionConfig gatewayConnection = new GatewayConnectionConfig();
+
+    @NestedConfigurationProperty
     private final DirectConnectionConfig directConnection = new DirectConnectionConfig();
 
     /**

diff --git a/...ure-service/src/main/java/com/azure/spring/service/cosmos/CosmosClientBuilderFactory.java b/...ure-service/src/main/java/com/azure/spring/service/cosmos/CosmosClientBuilderFactory.java
@@ -4,9 +4,14 @@
 package com.azure.spring.service.cosmos;
 
 import com.azure.core.credential.TokenCredential;
+import com.azure.core.http.ProxyOptions;
 import com.azure.core.util.Configuration;
 import com.azure.cosmos.ConnectionMode;
 import com.azure.cosmos.CosmosClientBuilder;
+import com.azure.cosmos.GatewayConnectionConfig;
+import com.azure.cosmos.ThrottlingRetryOptions;
+import com.azure.spring.core.aware.ProxyAware;
+import com.azure.spring.core.aware.RetryAware;
 import com.azure.spring.core.credential.descriptor.AuthenticationDescriptor;
 import com.azure.spring.core.credential.descriptor.KeyAuthenticationDescriptor;
 import com.azure.spring.core.credential.descriptor.TokenAuthenticationDescriptor;
@@ -20,6 +25,8 @@
 import java.util.List;
 import java.util.function.BiConsumer;
 
+import static com.azure.spring.core.converter.AzureHttpProxyOptionsConverter.HTTP_PROXY_CONVERTER;
+
 /**
  * Cosmos client builder factory, it builds the {@link CosmosClientBuilder} according the configuration context and
  * blob properties.
@@ -30,6 +37,9 @@ public class CosmosClientBuilderFactory extends AbstractAzureServiceClientBuilde
 
     private final CosmosProperties cosmosProperties;
 
+    private ProxyOptions proxyOptions;
+    private ThrottlingRetryOptions throttlingRetryOptions;
+
     public CosmosClientBuilderFactory(CosmosProperties cosmosProperties) {
         this.cosmosProperties = cosmosProperties;
     }
@@ -54,12 +64,23 @@ protected List<AuthenticationDescriptor<?>> getAuthenticationDescriptors(CosmosC
 
     @Override
     protected void configureProxy(CosmosClientBuilder builder) {
-        LOGGER.debug("No configureProxy for CosmosClientBuilder.");
+        ProxyAware.Proxy proxy = this.cosmosProperties.getProxy();
+        this.proxyOptions = HTTP_PROXY_CONVERTER.convert(proxy);
+        if (this.proxyOptions == null) {
+            LOGGER.debug("No proxy properties available.");
+        }
     }
 
     @Override
     protected void configureRetry(CosmosClientBuilder builder) {
-        LOGGER.debug("No configureRetry for CosmosClientBuilder.");
+        RetryAware.Retry retry = this.cosmosProperties.getRetry();
+        if (isInvalidRetry(retry)) {
+            return;
+        }
+
+        this.throttlingRetryOptions = new ThrottlingRetryOptions();
+        this.throttlingRetryOptions.setMaxRetryWaitTime(retry.getTimeout());
+        this.throttlingRetryOptions.setMaxRetryAttemptsOnThrottledRequests(retry.getMaxAttempts());
     }
 
     @Override
@@ -76,18 +97,70 @@ protected void configureService(CosmosClientBuilder builder) {
         map.from(this.cosmosProperties.getReadRequestsFallbackEnabled()).to(builder::readRequestsFallbackEnabled);
         map.from(this.cosmosProperties.getSessionCapturingOverrideEnabled()).to(builder::sessionCapturingOverrideEnabled);
         map.from(this.cosmosProperties.getPreferredRegions()).whenNot(List::isEmpty).to(builder::preferredRegions);
-        map.from(this.cosmosProperties.getThrottlingRetryOptions()).to(builder::throttlingRetryOptions);
+        configureThrottlingRetryOptions(builder, map);
+        configureConnection(builder, map);
+    }
 
+    /**
+     * Configure Cosmos connection.
+     * If not configured the proxy of gateway connection, then will try to use the root proxy of Cosmos properties.
+     * @param builder Cosmos client builder
+     * @param map Property mapper
+     */
+    private void configureConnection(CosmosClientBuilder builder, PropertyMapper map) {
         // TODO (xiada): should we count this as authentication
         map.from(this.cosmosProperties.getResourceToken()).to(builder::resourceToken);
         map.from(this.cosmosProperties.getPermissions()).whenNot(List::isEmpty).to(builder::permissions);
+        GatewayConnectionConfig gatewayConnection = this.cosmosProperties.getGatewayConnection();
+        if (proxyOptions != null && gatewayConnection.getProxy() == null) {
+            gatewayConnection.setProxy(proxyOptions);
+            LOGGER.debug("The proxy of the Gateway connection is not configured, "
+                + "then the Azure Spring Proxy configuration will be applied to Cosmos gateway connection.");
+        }
         if (ConnectionMode.DIRECT.equals(this.cosmosProperties.getConnectionMode())) {
-            builder.directMode(this.cosmosProperties.getDirectConnection(), this.cosmosProperties.getGatewayConnection());
+            builder.directMode(this.cosmosProperties.getDirectConnection(), gatewayConnection);
         } else if (ConnectionMode.GATEWAY.equals(this.cosmosProperties.getConnectionMode())) {
-            builder.gatewayMode(this.cosmosProperties.getGatewayConnection());
+            builder.gatewayMode(gatewayConnection);
         }
     }
 
+    /**
+     * Configure ThrottlingRetryOptions.
+     * If not configured the retry options of ThrottlingRetryOptions, then will try to use the root retry options of Cosmos properties.
+     * @param builder Cosmos client builder
+     * @param map Property mapper
+     */
+    private void configureThrottlingRetryOptions(CosmosClientBuilder builder, PropertyMapper map) {
+        ThrottlingRetryOptions retryOptions = this.cosmosProperties.getThrottlingRetryOptions();
+        if (this.throttlingRetryOptions != null && isDefaultThrottlingRetryOptions(retryOptions)) {
+            map.from(this.throttlingRetryOptions).to(builder::throttlingRetryOptions);
+            LOGGER.debug("The throttling retry options is not configured, "
+                + "then the Azure Spring Retry configuration will be applied to Cosmos service builder.");
+        } else {
+            map.from(retryOptions).to(builder::throttlingRetryOptions);
+        }
+    }
+
+    /**
+     * Check if the retry option is the default value, which is defined in azure-cosmos SDK.
+     * @param retryOptions retry options to be checked
+     * @return result
+     */
+    private boolean isDefaultThrottlingRetryOptions(ThrottlingRetryOptions retryOptions) {
+        ThrottlingRetryOptions defaultOptions = new ThrottlingRetryOptions();
+        return defaultOptions.getMaxRetryAttemptsOnThrottledRequests() == retryOptions.getMaxRetryAttemptsOnThrottledRequests()
+            && defaultOptions.getMaxRetryWaitTime().equals(retryOptions.getMaxRetryWaitTime());
+    }
+
+    /**
+     * Check if the properties of the retry is invalid value.
+     * @param retry retry options to be checked
+     * @return result
+     */
+    private boolean isInvalidRetry(RetryAware.Retry retry) {
+        return retry.getMaxAttempts() == null || retry.getTimeout() == null;
+    }
+
     @Override
     protected BiConsumer<CosmosClientBuilder, String> consumeApplicationId() {
         return CosmosClientBuilder::userAgentSuffix;