faultInjectionOnGateway (Azure#35378)

* faultInjectionOnGateway

---------

Co-authored-by: annie-mac <[email protected]>

Author: xinlian12

sdk/cosmos/azure-cosmos-test/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 1.0.0-beta.4 (Unreleased)
 
 #### Features Added
+* Added fault injection support for Gateway connection mode - See [PR 35378](https://github.com/Azure/azure-sdk-for-java/pull/35378)
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/faultinjection/FaultInjectionCondition.java

@@ -3,6 +3,8 @@
 
 package com.azure.cosmos.test.faultinjection;
 
+import com.azure.cosmos.test.implementation.ImplementationBridgeHelpers;
+
 /***
  * Fault injection condition.
  * A fault injection rule will not be applicable if the condition mismatches.
@@ -19,7 +21,7 @@ public final class FaultInjectionCondition {
         String region,
         FaultInjectionEndpoints endpoints) {
         this.operationType = operationType;
-        this.connectionType = connectionType;
+        this.connectionType = this.isMetadataOperationType() ? FaultInjectionConnectionType.GATEWAY : connectionType;
         this.region = region;
         this.endpoints = endpoints;
     }
@@ -60,6 +62,18 @@ public String getRegion() {
         return this.region;
     }
 
+    boolean isMetadataOperationType() {
+        if (this.operationType == null) {
+            return false;
+        }
+
+        return this.operationType == FaultInjectionOperationType.METADATA_REQUEST_PARTITION_KEY_RANGES
+            || this.operationType == FaultInjectionOperationType.METADATA_REQUEST_ADDRESS_REFRESH
+            || this.operationType == FaultInjectionOperationType.METADATA_REQUEST_CONTAINER
+            || this.operationType == FaultInjectionOperationType.METADATA_REQUEST_QUERY_PLAN
+            || this.operationType == FaultInjectionOperationType.METADATA_REQUEST_DATABASE_ACCOUNT;
+    }
+
     @Override
     public String toString() {
         return String.format(
@@ -69,4 +83,17 @@ public String toString() {
             this.connectionType,
             this.region);
     }
+
+    ///////////////////////////////////////////////////////////////////////////////////////////
+    // the following helper/accessor only helps to access this class outside of this package.//
+    ///////////////////////////////////////////////////////////////////////////////////////////
+    static void initialize() {
+        ImplementationBridgeHelpers.FaultInjectionConditionHelper.setFaultInjectionConditionAccessor(
+            faultInjectionCondition -> faultInjectionCondition.isMetadataOperationType()
+        );
+    }
+
+    static {
+        initialize();
+    }
 }

sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/faultinjection/FaultInjectionConnectionType.java

@@ -10,5 +10,9 @@ public enum FaultInjectionConnectionType {
     /***
      * Direct connection type.
      */
-    DIRECT
+    DIRECT,
+    /***
+     * Gateway connection type.
+     */
+    GATEWAY
 }

sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/faultinjection/FaultInjectionOperationType.java

@@ -34,7 +34,25 @@ public enum FaultInjectionOperationType {
     /**
      * Patch item.
      */
-    PATCH_ITEM
-
-    // Add support for metadata request type
+    PATCH_ITEM,
+    /**
+     * Read container request.
+     */
+    METADATA_REQUEST_CONTAINER,
+    /**
+     * Read database account request.
+     */
+    METADATA_REQUEST_DATABASE_ACCOUNT,
+    /**
+     * Query query plan request.
+     */
+    METADATA_REQUEST_QUERY_PLAN,
+    /**
+     * Partition key ranges request.
+     */
+    METADATA_REQUEST_PARTITION_KEY_RANGES,
+    /**
+     * Address refresh request.
+     */
+    METADATA_REQUEST_ADDRESS_REFRESH;
 }

sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/faultinjection/FaultInjectionRuleBuilder.java

@@ -4,6 +4,7 @@
 package com.azure.cosmos.test.faultinjection;
 
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
+import com.azure.cosmos.test.implementation.ImplementationBridgeHelpers;
 
 import java.time.Duration;
 
@@ -117,10 +118,21 @@ public FaultInjectionRuleBuilder enabled(boolean enabled) {
      * Create a new fault injection rule.
      *
      * @return the {@link FaultInjectionRule}.
+     * @throws IllegalArgumentException if condition is null.
+     * @throws IllegalArgumentException if result is null.
      */
     public FaultInjectionRule build() {
-        checkNotNull(this.condition, "Argument 'condition' can not be null");
-        checkNotNull(this.result, "Argument 'result' can not be null");
+        if (this.condition == null) {
+            throw new IllegalArgumentException("Argument 'condition' can not be null");
+        }
+
+        if (this.result == null) {
+            throw new IllegalArgumentException("Argument 'result' can not be null");
+        }
+
+        if (this.condition.getConnectionType() == FaultInjectionConnectionType.GATEWAY) {
+            this.validateRuleOnGatewayConnection();
+        }
 
         return new FaultInjectionRule(
             this.id,
@@ -131,4 +143,38 @@ public FaultInjectionRule build() {
             this.condition,
             this.result);
     }
+
+    private void validateRuleOnGatewayConnection() {
+        if (this.result == null) {
+            throw new IllegalArgumentException("Argument 'result' can not be null");
+        }
+
+        if (this.result instanceof FaultInjectionConnectionErrorResult) {
+            throw new IllegalArgumentException("FaultInjectionConnectionError result can not be configured for rule with gateway connection type.");
+        }
+
+        FaultInjectionServerErrorResult serverErrorResult = (FaultInjectionServerErrorResult) this.result;
+
+        // Gateway service internally will retry for 410/0, so the eventual exceptions being returned to SDK is 503(SERVICE_UNAVAILABLE) instead
+        if (serverErrorResult.getServerErrorType() == FaultInjectionServerErrorType.GONE) {
+            throw new IllegalArgumentException("Gone exception can not be injected for rule with gateway connection type");
+        }
+
+        if (serverErrorResult.getServerErrorType() == FaultInjectionServerErrorType.STALED_ADDRESSES_SERVER_GONE) {
+            throw new IllegalArgumentException("STALED_ADDRESSES exception can not be injected for rule with gateway connection type");
+        }
+
+        // for metadata request related rule, only CONNECTION_DELAY, RESPONSE_DELAY, TOO_MANY_REQUEST error can be injected
+        if (ImplementationBridgeHelpers
+                .FaultInjectionConditionHelper
+                .getFaultInjectionConditionAccessor()
+                .isMetadataOperationType(this.condition)) {
+            if (serverErrorResult.getServerErrorType() != FaultInjectionServerErrorType.TOO_MANY_REQUEST
+                && serverErrorResult.getServerErrorType() != FaultInjectionServerErrorType.RESPONSE_DELAY
+                && serverErrorResult.getServerErrorType() != FaultInjectionServerErrorType.CONNECTION_DELAY) {
+
+                throw new IllegalArgumentException("Error type " + serverErrorResult.getServerErrorType() + " is not supported for rule with metadata request");
+            }
+        }
+    }
 }

sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/faultinjection/FaultInjectionServerErrorResultBuilder.java

@@ -83,6 +83,12 @@ public FaultInjectionServerErrorResult build() {
             throw new IllegalArgumentException("Argument 'delay' is required for server error type " + this.serverErrorType);
         }
 
+        if (this.serverErrorType == FaultInjectionServerErrorType.STALED_ADDRESSES_SERVER_GONE) {
+            // for staled addresses errors, the error can only be cleared if forceRefresh address refresh request happened
+            // so default the times to max value
+            this.times = Integer.MAX_VALUE;
+        }
+
         return new FaultInjectionServerErrorResult(
             this.serverErrorType,
             this.times,

sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/faultinjection/FaultInjectionServerErrorType.java

@@ -8,7 +8,7 @@
  */
 public enum FaultInjectionServerErrorType {
 
-    /** 410 from server */
+    /** 410 from server. Only applicable for direct connection type. */
     GONE,
 
     /** 449 from server */
@@ -36,5 +36,13 @@ public enum FaultInjectionServerErrorType {
     RESPONSE_DELAY,
 
     /** simulate high channel acquisition, when it is over connection timeout, can simulate connectionTimeoutException */
-    CONNECTION_DELAY
+    CONNECTION_DELAY,
+    /**
+     * Simulate service unavailable(503)
+     */
+    SERVICE_UNAVAILABLE,
+    /**
+     * simulate 410-0 due to staled addresses. The exception will only be cleared if a forceRefresh address refresh happened.
+     */
+    STALED_ADDRESSES_SERVER_GONE
 }

sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/implementation/ImplementationBridgeHelpers.java

@@ -3,8 +3,9 @@
 
 package com.azure.cosmos.test.implementation;
 
-import com.azure.cosmos.test.implementation.faultinjection.IFaultInjectionRuleInternal;
+import com.azure.cosmos.test.faultinjection.FaultInjectionCondition;
 import com.azure.cosmos.test.faultinjection.FaultInjectionRule;
+import com.azure.cosmos.test.implementation.faultinjection.IFaultInjectionRuleInternal;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -52,4 +53,42 @@ public interface FaultInjectionRuleAccessor {
             void setEffectiveFaultInjectionRule(FaultInjectionRule rule, IFaultInjectionRuleInternal ruleInternal);
         }
     }
+
+    public static final class FaultInjectionConditionHelper {
+        private static final AtomicBoolean faultInjectionConditionClassLoaded = new AtomicBoolean(false);
+        private static final AtomicReference<FaultInjectionConditionAccessor> accessor = new AtomicReference<>();
+
+        private FaultInjectionConditionHelper() {
+        }
+
+        public static FaultInjectionConditionAccessor getFaultInjectionConditionAccessor() {
+            if (!faultInjectionConditionClassLoaded.get()) {
+                logger.debug("Initializing FaultInjectionConditionAccessor...");
+            }
+
+            FaultInjectionConditionAccessor snapshot = accessor.get();
+            if (snapshot == null) {
+                logger.error("FaultInjectionConditionAccessor is not initialized yet!");
+                System.exit(8701); // Using a unique status code here to help debug the issue.
+            }
+
+            return snapshot;
+        }
+
+        public static void setFaultInjectionConditionAccessor(final FaultInjectionConditionAccessor newAccessor) {
+
+            assert (newAccessor != null);
+
+            if (!accessor.compareAndSet(null, newAccessor)) {
+                logger.debug("FaultInjectionConditionAccessor already initialized!");
+            } else {
+                logger.debug("Setting FaultInjectionConditionAccessor...");
+                faultInjectionConditionClassLoaded.set(true);
+            }
+        }
+
+        public interface FaultInjectionConditionAccessor {
+            boolean isMetadataOperationType(FaultInjectionCondition faultInjectionCondition);
+        }
+    }
 }