[DRAFT] - Thinclient Endpoint Discovery + E2E Flow

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java

@@ -168,12 +168,12 @@ public void http2MaxConcurrentStreams() {
     @Test(groups = { "unit" })
     public void thinClientEnabledTest() {
         Configs config = new Configs();
-        assertThat(config.getThinclientEnabled()).isFalse();
+        assertThat(config.isThinClientEnabled()).isFalse();
 
         System.clearProperty("COSMOS.THINCLIENT_ENABLED");
         System.setProperty("COSMOS.THINCLIENT_ENABLED", "true");
         try {
-            assertThat(config.getThinclientEnabled()).isTrue();
+            assertThat(config.isThinClientEnabled()).isTrue();
         } finally {
             System.clearProperty("COSMOS.THINCLIENT_ENABLED");
         }

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientTest.java

@@ -0,0 +1,45 @@
+package com.azure.cosmos.implementation;
+
+import com.azure.cosmos.ConsistencyLevel;
+import com.azure.cosmos.CosmosAsyncClient;
+import com.azure.cosmos.CosmosAsyncContainer;
+import com.azure.cosmos.CosmosClientBuilder;
+import com.azure.cosmos.implementation.throughputControl.TestItem;
+import com.azure.cosmos.models.CosmosItemResponse;
+import com.azure.cosmos.models.PartitionKey;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.testng.annotations.Test;
+
+import java.util.UUID;
+
+import static java.lang.Thread.sleep;
+
+public class ThinClientTest {
+
+    @Test
+    public void testThinclientHttp2() {
+        try {
+            System.setProperty(Configs.THINCLIENT_ENABLED, "true");
+            System.setProperty(Configs.HTTP2_ENABLED, "true");
+
+            CosmosAsyncClient client  = new CosmosClientBuilder()
+                    .key(TestConfigurations.MASTER_KEY)
+                    .endpoint(TestConfigurations.HOST)
+                    .gatewayMode()
+                    .consistencyLevel(ConsistencyLevel.SESSION)
+                    .buildAsyncClient();
+
+            CosmosAsyncContainer container = client.getDatabase("updatedd-thin-client-test-db").getContainer("thin-client-test-container-1");
+            TestItem testItem = TestItem.createNewItem();
+            System.out.println(testItem.getId());
+            container.createItem(testItem).block();
+            container.readItem(testItem.getId(), new PartitionKey(testItem.getPk()), TestItem.class).block();
+            container.deleteItem(testItem.getId(), new PartitionKey(testItem.getPk())).block();
+        } finally {
+            System.clearProperty(Configs.THINCLIENT_ENABLED);
+            System.clearProperty(Configs.HTTP2_ENABLED);
+        }
+
+    }
+}

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/throughputControl/TestItem.java

@@ -7,6 +7,7 @@
 
 public class TestItem {
     private String id;
+    private String pk;
     private String mypk;
     private String prop;
 
@@ -15,6 +16,7 @@ public TestItem() {
 
     public TestItem(String id, String mypk, String prop) {
         this.id = id;
+        this.pk = mypk;
         this.mypk = mypk;
         this.prop = prop;
     }
@@ -31,6 +33,8 @@ public void setId(String id) {
         this.id = id;
     }
 
+    public String getPk() { return pk; }
+
     public String getMypk() {
         return mypk;
     }

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

@@ -46,10 +46,10 @@ public class Configs {
     private static final String UNAVAILABLE_LOCATIONS_EXPIRATION_TIME_IN_SECONDS = "COSMOS.UNAVAILABLE_LOCATIONS_EXPIRATION_TIME_IN_SECONDS";
     private static final String GLOBAL_ENDPOINT_MANAGER_INITIALIZATION_TIME_IN_SECONDS = "COSMOS.GLOBAL_ENDPOINT_MANAGER_MAX_INIT_TIME_IN_SECONDS";
     private static final String DEFAULT_THINCLIENT_ENDPOINT = "";
-    private static final String THINCLIENT_ENDPOINT = "COSMOS.THINCLIENT_ENDPOINT";
+    public static final String THINCLIENT_ENDPOINT = "COSMOS.THINCLIENT_ENDPOINT";
     private static final String THINCLIENT_ENDPOINT_VARIABLE = "COSMOS_THINCLIENT_ENDPOINT";
     private static final boolean DEFAULT_THINCLIENT_ENABLED = false;
-    private static final String THINCLIENT_ENABLED = "COSMOS.THINCLIENT_ENABLED";
+    public static final String THINCLIENT_ENABLED = "COSMOS.THINCLIENT_ENABLED";
     private static final String THINCLIENT_ENABLED_VARIABLE = "COSMOS_THINCLIENT_ENABLED";
 
     private static final String MAX_HTTP_BODY_LENGTH_IN_BYTES = "COSMOS.MAX_HTTP_BODY_LENGTH_IN_BYTES";
@@ -312,7 +312,7 @@ public class Configs {
 
     // Flag to indicate whether enabled http2 for gateway
     private static final boolean DEFAULT_HTTP2_ENABLED = false;
-    private static final String HTTP2_ENABLED = "COSMOS.HTTP2_ENABLED";
+    public static final String HTTP2_ENABLED = "COSMOS.HTTP2_ENABLED";
     private static final String HTTP2_ENABLED_VARIABLE = "COSMOS_HTTP2_ENABLED";
 
     // Config to indicate the maximum number of live connections to keep in the pool for http2
@@ -434,7 +434,7 @@ public URI getThinclientEndpoint() {
         return URI.create(DEFAULT_THINCLIENT_ENDPOINT);
     }
 
-    public static boolean getThinclientEnabled() {
+    public static boolean isThinClientEnabled() {
         String valueFromSystemProperty = System.getProperty(THINCLIENT_ENABLED);
         if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) {
             return Boolean.parseBoolean(valueFromSystemProperty);

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

@@ -199,6 +199,8 @@ public static final class Properties {
         public static final String Name = "name";
         public static final String WRITABLE_LOCATIONS = "writableLocations";
         public static final String READABLE_LOCATIONS = "readableLocations";
+        public static final String THINCLIENT_WRITABLE_LOCATIONS = "thinClientWritableLocations";
+        public static final String THINCLIENT_READABLE_LOCATIONS = "thinClientReadableLocations";
         public static final String DATABASE_ACCOUNT_ENDPOINT = "databaseAccountEndpoint";
 
       //Authorization

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

@@ -251,6 +251,46 @@ public void setReadableLocations(Iterable<DatabaseAccountLocation> locations) {
         this.set(Constants.Properties.READABLE_LOCATIONS, locations);
     }
 
+    /**
+     * Gets the list of thin client readable locations for this database account.
+     *
+     * @return the list of thin client readable locations.
+     */
+    public Iterable<DatabaseAccountLocation> getThinClientReadableLocations() {
+        return super.getCollection(Constants.Properties.THINCLIENT_READABLE_LOCATIONS, DatabaseAccountLocation.class);
+    }
+
+    /**
+     * Sets the list of thin client readable locations for this database account.
+     * <p>
+     * The list of thin client readable locations are returned by the service.
+     *
+     * @param locations the list of thin client readable locations.
+     */
+    public void setThinClientReadableLocations(Iterable<DatabaseAccountLocation> locations) {
+        this.set(Constants.Properties.THINCLIENT_READABLE_LOCATIONS, locations);
+    }
+
+    /**
+     * Gets the list of thin client writable locations for this database account.
+     *
+     * @return the list of thin client writable locations.
+     */
+    public Iterable<DatabaseAccountLocation> getThinClientWritableLocations() {
+        return super.getCollection(Constants.Properties.THINCLIENT_WRITABLE_LOCATIONS, DatabaseAccountLocation.class);
+    }
+
+    /**
+     * Sets the list of thin client writable locations for this database account.
+     * <p>
+     * The list of thin client writable locations are returned by the service.
+     *
+     * @param locations the list of thin clientwritable locations.
+     */
+    public void setThinClientWritableLocations(Iterable<DatabaseAccountLocation> locations) {
+        this.set(Constants.Properties.THINCLIENT_WRITABLE_LOCATIONS, locations);
+    }
+
     /**
      * Gets if enable multiple write locations is set.
      *

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

@@ -330,8 +330,10 @@ private Mono<Void> startRefreshLocationTimerAsync(boolean initialization) {
                             }
 
                             logger.debug("startRefreshLocationTimerAsync() - Invoking refresh, I was registered on [{}]", now);
-                            Mono<DatabaseAccount> databaseAccountObs = GlobalEndpointManager.getDatabaseAccountFromAnyLocationsAsync(this.defaultEndpoint, new ArrayList<>(this.getEffectivePreferredRegions()),
-                                    this::getDatabaseAccountAsync);
+                            Mono<DatabaseAccount> databaseAccountObs = GlobalEndpointManager.getDatabaseAccountFromAnyLocationsAsync(
+                                this.defaultEndpoint,
+                                new ArrayList<>(this.getEffectivePreferredRegions()),
+                                this::getDatabaseAccountAsync);
 
                             return databaseAccountObs.flatMap(dbAccount -> {
                                 logger.info("db account retrieved {}", dbAccount);

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

@@ -255,6 +255,8 @@ public static class HttpHeaders {
         public static final String POPULATE_INDEX_METRICS = "x-ms-cosmos-populateindexmetrics";
         public static final String INDEX_UTILIZATION = "x-ms-cosmos-index-utilization";
         public static final String QUERY_EXECUTION_INFO = "x-ms-cosmos-query-execution-info";
+        public static final String START_EPK_HASH = "x-ms-cosmos-start-epk-hash";
+        public static final String END_EPK_HASH = "x-ms-cosmos-end-epk-hash";
 
         // Batch operations
         public static final String IS_BATCH_ATOMIC = "x-ms-cosmos-batch-atomic";
@@ -283,11 +285,12 @@ public static class HttpHeaders {
         // Priority Level for throttling
         public static final String PRIORITY_LEVEL = "x-ms-cosmos-priority-level";
 
+        public static final String GLOBAL_DATABASE_ACCOUNT_NAME = "GlobalDatabaseAccountName";
+
         // Thinclient headers
         public static final String THINCLIENT_PROXY_OPERATION_TYPE = "x-ms-thinclient-proxy-operation-type";
         public static final String THINCLIENT_PROXY_RESOURCE_TYPE = "x-ms-thinclient-proxy-resource-type";
-
-        public static final String GLOBAL_DATABASE_ACCOUNT_NAME = "GlobalDatabaseAccountName";
+        public static final String THINCLIENT_OPT_IN = "x-ms-cosmos-use-thinclient";
     }
 
     public static class A_IMHeaderValues {

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

@@ -56,6 +56,7 @@
 import com.azure.cosmos.implementation.query.PipelinedQueryExecutionContextBase;
 import com.azure.cosmos.implementation.query.QueryInfo;
 import com.azure.cosmos.implementation.routing.CollectionRoutingMap;
+import com.azure.cosmos.implementation.routing.LocationCache;
 import com.azure.cosmos.implementation.routing.PartitionKeyAndResourceTokenPair;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper;
@@ -211,6 +212,7 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization
     private String firstResourceTokenFromPermissionFeed = StringUtils.EMPTY;
     private RxClientCollectionCache collectionCache;
     private RxGatewayStoreModel gatewayProxy;
+    private RxGatewayStoreModel thinProxy;
     private RxStoreModel storeModel;
     private GlobalAddressResolver addressResolver;
     private RxPartitionKeyRangeCache partitionKeyRangeCache;
@@ -261,6 +263,7 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization
     private final boolean isRegionScopedSessionCapturingEnabledOnClientOrSystemConfig;
     private List<CosmosOperationPolicy> operationPolicies;
     private AtomicReference<CosmosAsyncClient> cachedCosmosAsyncClientSnapshot;
+    private final boolean isThinClientEnabled;
 
     public RxDocumentClientImpl(URI serviceEndpoint,
                                 String masterKeyOrResourceToken,
@@ -489,6 +492,7 @@ private RxDocumentClientImpl(URI serviceEndpoint,
             this.connectionSharingAcrossClientsEnabled = connectionSharingAcrossClientsEnabled;
             this.configs = configs;
             this.masterKeyOrResourceToken = masterKeyOrResourceToken;
+            this.isThinClientEnabled = Configs.isThinClientEnabled();
             this.serviceEndpoint = serviceEndpoint;
             this.credential = credential;
             this.tokenCredential = tokenCredential;
@@ -669,6 +673,14 @@ private void updateGatewayProxy() {
         (this.gatewayProxy).setSessionContainer(this.sessionContainer);
     }
 
+    private void updateThinProxy() {
+        (this.thinProxy).setGatewayServiceConfigurationReader(this.gatewayConfigurationReader);
+        (this.thinProxy).setCollectionCache(this.collectionCache);
+        (this.thinProxy).setPartitionKeyRangeCache(this.partitionKeyRangeCache);
+        (this.thinProxy).setUseMultipleWriteLocations(this.useMultipleWriteLocations);
+        (this.thinProxy).setSessionContainer(this.sessionContainer);
+    }
+
     public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Function<HttpClient, HttpClient> httpClientInterceptor) {
         try {
 
@@ -685,6 +697,12 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func
                 this.reactorHttpClient,
                 this.apiType);
 
+            this.thinProxy = createThinProxy(this.sessionContainer,
+                this.consistencyLevel,
+                this.userAgentContainer,
+                this.globalEndpointManager,
+                this.reactorHttpClient);
+
             this.globalEndpointManager.init();
 
             DatabaseAccount databaseAccountSnapshot = this.initializeGatewayConfigurationReader();
@@ -712,6 +730,7 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func
                 collectionCache);
 
             updateGatewayProxy();
+            updateThinProxy();
             clientTelemetry = new ClientTelemetry(
                     this,
                     null,
@@ -827,6 +846,20 @@ RxGatewayStoreModel createRxGatewayProxy(ISessionContainer sessionContainer,
                 apiType);
     }
 
+    ThinClientStoreModel createThinProxy(ISessionContainer sessionContainer,
+                                             ConsistencyLevel consistencyLevel,
+                                             UserAgentContainer userAgentContainer,
+                                             GlobalEndpointManager globalEndpointManager,
+                                             HttpClient httpClient) {
+        return new ThinClientStoreModel(
+            this,
+            sessionContainer,
+            consistencyLevel,
+            userAgentContainer,
+            globalEndpointManager,
+            httpClient);
+    }
+
     private HttpClient httpClient() {
         HttpClientConfig httpClientConfig = new HttpClientConfig(this.configs)
             .withMaxIdleConnectionTimeout(this.connectionPolicy.getIdleHttpConnectionTimeout())
@@ -1870,6 +1903,8 @@ private Mono<Tuple2<RxDocumentServiceRequest, Utils.ValueHolder<DocumentCollecti
             getEffectiveClientContext(clientContextOverride),
             operationType, ResourceType.Document, path, requestHeaders, options, content);
 
+        request.useThinProxy = Configs.isThinClientEnabled() &&  this.connectionPolicy.getConnectionMode() == ConnectionMode.GATEWAY ? true : false;
+
         if (operationType.isWriteOperation() &&  options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled()) {
             request.setNonIdempotentWriteRetriesEnabled(true);
         }
@@ -2033,6 +2068,10 @@ private RxDocumentServiceRequest addBatchHeaders(RxDocumentServiceRequest reques
      * @return Mono, which on subscription will populate the headers in the request passed in the argument.
      */
     public Mono<RxDocumentServiceRequest> populateHeadersAsync(RxDocumentServiceRequest request, RequestVerb httpMethod) {
+        // if thin client enabled, populate thin client header so we can get thin client read and writeable locations
+        if (request.useThinProxy) {
+            request.getHeaders().put(HttpConstants.HttpHeaders.THINCLIENT_OPT_IN, "true");
+        }
         request.getHeaders().put(HttpConstants.HttpHeaders.X_DATE, Utils.nowAsRFC1123());
         if (this.masterKeyOrResourceToken != null || this.resourceTokensMap != null
             || this.cosmosAuthorizationTokenResolver != null || this.credential != null) {
@@ -3405,6 +3444,7 @@ private Mono<ResourceResponse<Document>> readDocumentInternal(
                 getEffectiveClientContext(clientContextOverride),
                 OperationType.Read, ResourceType.Document, path, requestHeaders, options);
 
+            request.useThinProxy = Configs.isThinClientEnabled() && this.connectionPolicy.getConnectionMode() == ConnectionMode.GATEWAY ? true : false;
             DocumentServiceRequestContext requestContext = request.requestContext;
 
             options.getMarkE2ETimeoutInRequestContextCallbackHook().set(
@@ -5696,6 +5736,7 @@ public Flux<DatabaseAccount> getDatabaseAccountFromEndpoint(URI endpoint) {
         return Flux.defer(() -> {
             RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
                 OperationType.Read, ResourceType.DatabaseAccount, "", null, (Object) null);
+            request.useThinProxy = this.isThinClientEnabled && this.connectionPolicy.getConnectionMode() == ConnectionMode.GATEWAY ? true : false;
             return this.populateHeadersAsync(request, RequestVerb.GET)
                 .flatMap(requestPopulated -> {
 
@@ -5721,6 +5762,15 @@ public Flux<DatabaseAccount> getDatabaseAccountFromEndpoint(URI endpoint) {
      * @return RxStoreModel
      */
     private RxStoreModel getStoreProxy(RxDocumentServiceRequest request) {
+        // TODO: remove, for testing
+        if (request.isMetadataRequest()) {
+            return this.gatewayProxy;
+        }
+
+        if (request.useThinProxy) {
+            return this.thinProxy;
+        }
+
         // If a request is configured to always use GATEWAY mode(in some cases when targeting .NET Core)
         // we return the GATEWAY store model
         if (request.useGatewayMode) {

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

@@ -86,6 +86,7 @@ public class RxDocumentServiceRequest implements Cloneable {
     public String throughputControlGroupName;
     public volatile boolean intendedCollectionRidPassedIntoSDK = false;
     private volatile Duration responseTimeout;
+    public volatile boolean useThinProxy = false;
 
     private volatile boolean nonIdempotentWriteRetriesEnabled = false;
 
@@ -1185,9 +1186,11 @@ public void setEffectivePartitionKey(String effectivePartitionKey) {
         this.effectivePartitionKey = effectivePartitionKey;
     }
 
-    public void setThinclientHeaders(String operationType, String resourceType) {
+    public void setThinclientHeaders(String operationType, String resourceType, String globalDatabaseAccountName, String resourceId) {
         this.headers.put(HttpConstants.HttpHeaders.THINCLIENT_PROXY_OPERATION_TYPE, operationType);
         this.headers.put(HttpConstants.HttpHeaders.THINCLIENT_PROXY_RESOURCE_TYPE, resourceType);
+        this.headers.put(HttpConstants.HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME, globalDatabaseAccountName);
+        this.headers.put(WFConstants.BackendHeaders.COLLECTION_RID, resourceId);
     }
 
     public RxDocumentServiceRequest setHttpTransportSerializer(HttpTransportSerializer transportSerializer) {