[FEATURE] Cosmos | Connection endpoint rediscovery (connection state listener)

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncBenchmark.java

@@ -89,8 +89,9 @@ abstract class AsyncBenchmark<T> {
         configuration = cfg;
         logger = LoggerFactory.getLogger(this.getClass());
 
+        cosmosAsyncDatabase = cosmosClient.getDatabase(this.configuration.getDatabaseId());
+
         try {
-            cosmosAsyncDatabase = cosmosClient.getDatabase(this.configuration.getDatabaseId());
             cosmosAsyncDatabase.read().block();
         } catch (CosmosException e) {
             if (e.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
@@ -103,25 +104,25 @@ abstract class AsyncBenchmark<T> {
             }
         }
 
-        try {
-            cosmosAsyncContainer = cosmosAsyncDatabase.getContainer(this.configuration.getCollectionId());
+        cosmosAsyncContainer = cosmosAsyncDatabase.getContainer(this.configuration.getCollectionId());
 
+        try {
             cosmosAsyncContainer.read().block();
-
         } catch (CosmosException e) {
-            if (e.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
-                cosmosAsyncDatabase.createContainer(
-                    this.configuration.getCollectionId(),
-                    Configuration.DEFAULT_PARTITION_KEY_PATH,
-                    ThroughputProperties.createManualThroughput(this.configuration.getThroughput())
-                ).block();
-
-                cosmosAsyncContainer = cosmosAsyncDatabase.getContainer(this.configuration.getCollectionId());
-                logger.info("Collection {} is created for this test", this.configuration.getCollectionId());
-                collectionCreated = true;
-            } else {
+
+            if (e.getStatusCode() != HttpConstants.StatusCodes.NOTFOUND) {
                 throw e;
             }
+
+            cosmosAsyncDatabase.createContainer(
+                this.configuration.getCollectionId(),
+                Configuration.DEFAULT_PARTITION_KEY_PATH,
+                ThroughputProperties.createManualThroughput(this.configuration.getThroughput())
+            ).block();
+
+            cosmosAsyncContainer = cosmosAsyncDatabase.getContainer(this.configuration.getCollectionId());
+            logger.info("Collection {} is created for this test", this.configuration.getCollectionId());
+            collectionCreated = true;
         }
 
         partitionKey = cosmosAsyncContainer.read().block().getProperties().getPartitionKeyDefinition()

sdk/cosmos/azure-cosmos-benchmark/src/test/java/com/azure/cosmos/benchmark/ReadMyWritesConsistencyTest.java

@@ -196,7 +196,6 @@ private void scheduleScaleUp(int delayStartInSeconds, int newThroughput) {
                 },
                 () -> {
                     logger.info("Collection Scale up request sent to the service");
-
                 }
             );
     }

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

@@ -8,10 +8,11 @@
 import java.time.Duration;
 
 /**
- * Represents the connection config with {@link ConnectionMode#DIRECT} associated with Cosmos Client in the Azure Cosmos DB database service.
- * For performance tips on how to optimize Direct connection configuration,
- * refer to performance tips guide:
- * <a href="https://docs.microsoft.com/en-us/azure/cosmos-db/performance-tips-java-sdk-v4-sql?tabs=api-async">Performance tips guide</a>
+ * Represents the connection config with {@link ConnectionMode#DIRECT} associated with Cosmos Client in the Azure Cosmos
+ * DB database service. For performance tips on how to optimize Direct connection configuration, refer to performance
+ * tips guide:
+ * <a href="https://docs.microsoft.com/en-us/azure/cosmos-db/performance-tips-java-sdk-v4-sql?tabs=api-async">Performance
+ * tips guide</a>
  */
 public final class DirectConnectionConfig {
     //  Constants
@@ -22,6 +23,7 @@ public final class DirectConnectionConfig {
     private static final int DEFAULT_MAX_REQUESTS_PER_CONNECTION = 30;
 
     private Duration connectTimeout;
+    private boolean enableConnectionEndpointRediscovery;
     private Duration idleConnectionTimeout;
     private Duration idleEndpointTimeout;
     private Duration requestTimeout;
@@ -33,6 +35,7 @@ public final class DirectConnectionConfig {
      */
     public DirectConnectionConfig() {
         this.connectTimeout = DEFAULT_CONNECT_TIMEOUT;
+        this.enableConnectionEndpointRediscovery = false;
         this.idleConnectionTimeout = Duration.ZERO;
         this.idleEndpointTimeout = DEFAULT_IDLE_ENDPOINT_TIMEOUT;
         this.maxConnectionsPerEndpoint = DEFAULT_MAX_CONNECTIONS_PER_ENDPOINT;
@@ -79,6 +82,42 @@ public DirectConnectionConfig setConnectTimeout(Duration connectTimeout) {
         return this;
     }
 
+    /**
+     * Gets a value indicating whether Direct TCP connection endpoint rediscovery is enabled.
+     * <p>
+     * The connection endpoint rediscovery feature is designed to reduce and spread-out latency spikes that are likely
+     * to occur:
+     * <ul>
+     * <li>During rolling upgrades of a Cosmos instance or
+     * <li>When a backend node is being decommissioned or restarted (e.g., to restart or remove an unhealthy replica.)
+     * </ul>
+     *
+     * @return {@code true} if Direct TCP connection endpoint rediscovery is enabled; {@code false} otherwise.
+     */
+    public boolean isConnectionEndpointRediscoveryEnabled() {
+        return this.enableConnectionEndpointRediscovery;
+    }
+
+    /**
+     * Sets a value indicating whether Direct TCP connection endpoint rediscovery should be enabled.
+     * <p>
+     * The connection endpoint rediscovery feature is designed to reduce and spread-out latency spikes that are likely
+     * to occur:
+     * <ul>
+     * <li>During rolling upgrades of a Cosmos instance or
+     * <li>When a backend node is being decommissioned or restarted (e.g., to restart or remove an unhealthy replica.)
+     * </ul>
+     *
+     * @param connectionEndpointRediscoveryEnabled {@code true} if connection endpoint rediscovery is enabled; {@code
+     *                                             false} otherwise.
+     *
+     * @return the {@linkplain DirectConnectionConfig}.
+     */
+    public DirectConnectionConfig setConnectionEndpointRediscoveryEnabled(boolean connectionEndpointRediscoveryEnabled) {
+        this.enableConnectionEndpointRediscovery = connectionEndpointRediscoveryEnabled;
+        return this;
+    }
+
     /**
      * Gets the idle connection timeout for direct client
      *

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

@@ -41,16 +41,18 @@ public final class ConnectionPolicy {
 
     //  Direct connection config properties
     private Duration connectTimeout;
-    private Duration idleEndpointTimeout;
+    private boolean tcpConnectionEndpointRediscoveryEnabled;
+    private Duration idleTcpConnectionTimeout;
+    private Duration idleTcpEndpointTimeout;
     private int maxConnectionsPerEndpoint;
     private int maxRequestsPerConnection;
-    private Duration idleTcpConnectionTimeout;
 
     /**
      * Constructor.
      */
     public ConnectionPolicy(GatewayConnectionConfig gatewayConnectionConfig) {
         this(ConnectionMode.GATEWAY);
+        this.tcpConnectionEndpointRediscoveryEnabled = true;
         this.idleHttpConnectionTimeout = gatewayConnectionConfig.getIdleConnectionTimeout();
         this.maxConnectionPoolSize = gatewayConnectionConfig.getMaxConnectionPoolSize();
         this.requestTimeout = BridgeInternal.getRequestTimeoutFromGatewayConnectionConfig(gatewayConnectionConfig);
@@ -60,8 +62,9 @@ public ConnectionPolicy(GatewayConnectionConfig gatewayConnectionConfig) {
     public ConnectionPolicy(DirectConnectionConfig directConnectionConfig) {
         this(ConnectionMode.DIRECT);
         this.connectTimeout = directConnectionConfig.getConnectTimeout();
+        this.tcpConnectionEndpointRediscoveryEnabled = directConnectionConfig.isConnectionEndpointRediscoveryEnabled();
         this.idleTcpConnectionTimeout = directConnectionConfig.getIdleConnectionTimeout();
-        this.idleEndpointTimeout = directConnectionConfig.getIdleEndpointTimeout();
+        this.idleTcpEndpointTimeout = directConnectionConfig.getIdleEndpointTimeout();
         this.maxConnectionsPerEndpoint = directConnectionConfig.getMaxConnectionsPerEndpoint();
         this.maxRequestsPerConnection = directConnectionConfig.getMaxRequestsPerConnection();
         this.requestTimeout = BridgeInternal.getRequestTimeoutFromDirectConnectionConfig(directConnectionConfig);
@@ -87,6 +90,25 @@ public static ConnectionPolicy getDefaultPolicy() {
         return ConnectionPolicy.defaultPolicy;
     }
 
+    /**
+     * Gets a value that indicates whether Direct TCP connection endpoint rediscovery should is enabled.
+     *
+     * @return {@code true} if Direct TCP connection endpoint rediscovery should is enabled; {@code false} otherwise.
+     */
+    public boolean isTcpConnectionEndpointRediscoveryEnabled() {
+        return this.tcpConnectionEndpointRediscoveryEnabled;
+    }
+
+    /**
+     * Sets a value that indicates whether Direct TCP connection endpoint rediscovery should is enabled.
+     *
+     * @return the {@linkplain ConnectionPolicy}.
+     */
+    public ConnectionPolicy setTcpConnectionEndpointRediscoveryEnabled(boolean tcpConnectionEndpointRediscoveryEnabled) {
+        this.tcpConnectionEndpointRediscoveryEnabled = tcpConnectionEndpointRediscoveryEnabled;
+        return this;
+    }
+
     /**
      * Gets the request timeout (time to wait for response from network peer).
      *
@@ -426,17 +448,17 @@ public ConnectionPolicy setConnectTimeout(Duration connectTimeout) {
      * Gets the idle endpoint timeout
      * @return the idle endpoint timeout
      */
-    public Duration getIdleEndpointTimeout() {
-        return idleEndpointTimeout;
+    public Duration getIdleTcpEndpointTimeout() {
+        return this.idleTcpEndpointTimeout;
     }
 
     /**
      * Sets the idle endpoint timeout
-     * @param idleEndpointTimeout the idle endpoint timeout
+     * @param idleTcpEndpointTimeout the idle endpoint timeout
      * @return the {@link ConnectionPolicy}
      */
-    public ConnectionPolicy setIdleEndpointTimeout(Duration idleEndpointTimeout) {
-        this.idleEndpointTimeout = idleEndpointTimeout;
+    public ConnectionPolicy setIdleTcpEndpointTimeout(Duration idleTcpEndpointTimeout) {
+        this.idleTcpEndpointTimeout = idleTcpEndpointTimeout;
         return this;
     }
 
@@ -493,9 +515,10 @@ public String toString() {
             ", inetSocketProxyAddress=" + (proxy != null ? proxy.getAddress() : null) +
             ", readRequestsFallbackEnabled=" + readRequestsFallbackEnabled +
             ", connectTimeout=" + connectTimeout +
-            ", idleEndpointTimeout=" + idleEndpointTimeout +
+            ", idleTcpEndpointTimeout=" + idleTcpEndpointTimeout +
             ", maxConnectionsPerEndpoint=" + maxConnectionsPerEndpoint +
             ", maxRequestsPerConnection=" + maxRequestsPerConnection +
+            ", tcpConnectionEndpointRediscoveryEnabled=" + tcpConnectionEndpointRediscoveryEnabled +
             '}';
     }
 }

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

@@ -58,6 +58,17 @@ public GoneException(String message, String requestUri) {
         this(message, null, new HashMap<>(), requestUri);
     }
 
+    /**
+     * Instantiates a new {@link GoneException Gone exception}.
+     *
+     * @param message    the message
+     * @param requestUri the request uri
+     * @param cause      the cause of this (client-side) {@link GoneException}
+     */
+    public GoneException(String message, URI requestUri, Exception cause) {
+        this(message, cause, requestUri, (String) null);
+    }
+
     GoneException(String message,
                   Exception innerException,
                   URI requestUri,

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

@@ -288,25 +288,37 @@ public static class StatusCodes {
     }
 
     public static class SubStatusCodes {
+
         // Unknown SubStatus Code
         public static final int UNKNOWN = 0;
 
-        // 400: Bad Request substatus
+        // 400: Bad Request sub-status
         public static final int PARTITION_KEY_MISMATCH = 1001;
         public static final int CROSS_PARTITION_QUERY_NOT_SERVABLE = 1004;
 
-        // 410: StatusCodeType_Gone: substatus
+        // 410: StatusCodeType_Gone: sub-status
         public static final int NAME_CACHE_IS_STALE = 1000;
         public static final int PARTITION_KEY_RANGE_GONE = 1002;
         public static final int COMPLETING_SPLIT = 1007;
         public static final int COMPLETING_PARTITION_MIGRATION = 1008;
 
-        // 403: Forbidden substatus
+        // 403: Forbidden sub-status
         public static final int FORBIDDEN_WRITEFORBIDDEN = 3;
         public static final int DATABASE_ACCOUNT_NOTFOUND = 1008;
 
         // 404: LSN in session token is higher
         public static final int READ_SESSION_NOT_AVAILABLE = 1002;
+
+        // 410: Gone sub-status
+        // Sub-status code indicating that a GoneException was instantiated by the client; not as a result of a response
+        // message provide by a Cosmos instance.
+        public static final int CLIENT_GENERATED = 10_000;
+
+        // 410: Gone sub-status
+        // Sub-status code zero in a response from a service endpoint indicates that a replica is being discontinued or
+        // reconfigured. When endpoint rediscovery is enabled the RntbdTransportClient converts sub-status code zero to
+        // this sub-status code value.
+        public static final int DISCONTINUING_SERVICE = CLIENT_GENERATED + 2;
     }
 
     public static class HeaderValues {

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

@@ -156,13 +156,10 @@ public static final class Event {
         @JsonIgnore
         private final Duration duration;
 
-        @JsonSerialize(using = ToStringSerializer.class)
-        private final long durationInMicroSec;
-
         @JsonProperty("eventName")
         private final String name;
 
-        @JsonSerialize(using = ToStringSerializer.class)
+        @JsonProperty("startTimeUTC")
         private final Instant startTime;
 
         public Event(final String name, final Instant from, final Instant to) {
@@ -172,11 +169,12 @@ public Event(final String name, final Instant from, final Instant to) {
             this.name = name;
             this.startTime = from;
 
-            this.duration = from == null ? null : to == null ? Duration.ZERO : Duration.between(from, to);
-            if(this.duration != null) {
-                this.durationInMicroSec = duration.toNanos()/1000L;
+            if (from == null) {
+                this.duration = null;
+            } else if (to == null) {
+                this.duration = Duration.ZERO;
             } else {
-                this.durationInMicroSec = 0;
+                this.duration = Duration.between(from, to);
             }
         }