[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

@@ -22,6 +22,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 +34,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 +81,28 @@ public DirectConnectionConfig setConnectTimeout(Duration connectTimeout) {
         return this;
     }
 
+    /**
+     * 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 getEnableConnectionEndpointRediscovery() {
+        return this.enableConnectionEndpointRediscovery;
+    }
+
+    /**
+     * Sets a value that indicates whether Direct TCP connection endpoint rediscovery should is enabled.
+     *
+     * @param enableConnectionEndpointRediscovery {@code true} if connection endpoint rediscovery is enabled; {@code
+     *                                            false} otherwise.
+     *
+     * @return the {@linkplain DirectConnectionConfig}.
+     */
+    public DirectConnectionConfig setEnableConnectionEndpointRediscovery(boolean enableConnectionEndpointRediscovery) {
+        this.enableConnectionEndpointRediscovery = enableConnectionEndpointRediscovery;
+        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 enableTcpConnectionEndpointRediscovery;
+    private Duration idleTcpConnectionTimeout;
+    private Duration idleTcpEndpointTimeout;
     private int maxConnectionsPerEndpoint;
     private int maxRequestsPerConnection;
-    private Duration idleTcpConnectionTimeout;
 
     /**
      * Constructor.
      */
     public ConnectionPolicy(GatewayConnectionConfig gatewayConnectionConfig) {
         this(ConnectionMode.GATEWAY);
+        this.enableTcpConnectionEndpointRediscovery = false;
         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.enableTcpConnectionEndpointRediscovery = directConnectionConfig.getEnableConnectionEndpointRediscovery();
         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 getEnableTcpConnectionEndpointRediscovery() {
+        return this.enableTcpConnectionEndpointRediscovery;
+    }
+
+    /**
+     * Sets a value that indicates whether Direct TCP connection endpoint rediscovery should is enabled.
+     *
+     * @return the {@linkplain ConnectionPolicy}.
+     */
+    public ConnectionPolicy getEnableTcpConnectionEndpointRediscovery(boolean enableTcpConnectionEndpointRediscovery) {
+        this.enableTcpConnectionEndpointRediscovery = enableTcpConnectionEndpointRediscovery;
+        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,7 +515,7 @@ public String toString() {
             ", inetSocketProxyAddress=" + (proxy != null ? proxy.getAddress() : null) +
             ", readRequestsFallbackEnabled=" + readRequestsFallbackEnabled +
             ", connectTimeout=" + connectTimeout +
-            ", idleEndpointTimeout=" + idleEndpointTimeout +
+            ", idleEndpointTimeout=" + idleTcpEndpointTimeout +
             ", maxConnectionsPerEndpoint=" + maxConnectionsPerEndpoint +
             ", maxRequestsPerConnection=" + maxRequestsPerConnection +
             '}';

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);
             }
         }