[DRAFT] - Thinclient Endpoint Discovery + E2E Flow

sdk/cosmos/azure-cosmos-dotnet-benchmark/pom.xml

@@ -410,7 +410,7 @@ Licensed under the MIT License.
                   </descriptorRefs>
                   <archive>
                     <manifest>
-                      <mainClass>com.azure.cosmos.dotnet.benchmark.Main</mainClass>
+                      <mainClass>com.azure.cosmos.dotnet.benchmark.ThinClientReproMain</mainClass>
                     </manifest>
                   </archive>
                 </configuration>

sdk/cosmos/azure-cosmos-dotnet-benchmark/src/main/java/com/azure/cosmos/dotnet/benchmark/ThinClientReproMain.java

@@ -0,0 +1,63 @@
+package com.azure.cosmos.dotnet.benchmark;
+
+import com.azure.cosmos.ConsistencyLevel;
+import com.azure.cosmos.CosmosAsyncClient;
+import com.azure.cosmos.CosmosAsyncContainer;
+import com.azure.cosmos.CosmosClientBuilder;
+import com.azure.cosmos.CosmosException;
+import com.azure.cosmos.models.CosmosContainerResponse;
+import com.azure.cosmos.models.CosmosItemResponse;
+import com.azure.cosmos.models.PartitionKey;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+
+import java.util.UUID;
+
+public class ThinClientReproMain {
+    public static void main(String[] args) {
+        try {
+            System.setProperty("COSMOS.THINCLIENT_ENABLED", "true");
+            System.setProperty("COSMOS.HTTP2_ENABLED", "true");
+
+            CosmosAsyncClient client = new CosmosClientBuilder()
+                .key(System.getProperty("COSMOS.KEY"))
+                .endpoint(System.getProperty("COSMOS.ENDPOINT"))
+                .gatewayMode()
+                .consistencyLevel(ConsistencyLevel.SESSION)
+                .userAgentSuffix("fabianmThinClientProxyTest")
+                .buildAsyncClient();
+
+            CosmosAsyncContainer container = client.getDatabase("HashV2Small1").getContainer("HashV2Small1");
+            CosmosContainerResponse containerResponse = container.read().block();
+            System.out.println("Container RID: " + containerResponse.getProperties().getResourceId());
+            ObjectMapper mapper = new ObjectMapper();
+            ObjectNode doc = mapper.createObjectNode();
+            String idValue = UUID.randomUUID().toString();
+            doc.put("id", idValue);
+            System.out.println("Document to be ingested - " + doc.toPrettyString());
+
+            while (true) {
+                try {
+                    // container.readItem(
+                    //                        "HelloWorld",
+                    //                        new PartitionKey("HelloWorld"),
+                    //                        ObjectNode.class)
+                    CosmosItemResponse<ObjectNode> createResponse = container.createItem(doc).block();
+                    System.out.println("CREATE DIAGNOSTICS: " + createResponse.getDiagnostics());
+                    break;
+                } catch (CosmosException cosmosError) {
+                    System.out.println("COSMOS ERROR: " + cosmosError.getStatusCode() + "/" + cosmosError.getShortMessage());
+                    Thread.sleep(10_000);
+                }
+            }
+
+            CosmosItemResponse<ObjectNode> response = container.readItem(idValue, new PartitionKey(idValue), ObjectNode.class).block();
+            System.out.println("READ DIAGNOSTICS: " + response.getDiagnostics());
+            ObjectNode readDoc = response.getItem();
+
+            System.out.println("Document read - " + readDoc.toPrettyString());
+        } catch (CosmosException | InterruptedException cosmosException) {
+            System.out.println("COSMOS ERROR: " + cosmosException);
+        }
+    }
+}

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

@@ -4,24 +4,20 @@
 import com.azure.cosmos.CosmosAsyncClient;
 import com.azure.cosmos.CosmosAsyncContainer;
 import com.azure.cosmos.CosmosClientBuilder;
-import com.azure.cosmos.GatewayConnectionConfig;
 import com.azure.cosmos.implementation.throughputControl.TestItem;
-import com.azure.cosmos.models.CosmosDatabaseProperties;
-import com.azure.cosmos.models.CosmosDatabaseResponse;
 import com.azure.cosmos.models.PartitionKey;
-import com.azure.cosmos.util.CosmosPagedFlux;
 import com.fasterxml.jackson.databind.JsonNode;
 import org.testng.annotations.Test;
-import reactor.core.publisher.Mono;
 
 public class ThinClientTest {
+
     @Test
     public void testThinclientHttp2() {
         try {
             //String thinclientEndpoint = "https://cdb-ms-stage-eastus2-fe2.eastus2.cloudapp.azure.com:10650";
-            String thinclientEndpoint = "https://chukangzhongstagesignoff.documents-staging.windows-ppe.net:443/";
+            //String thinclientEndpoint = "https://chukangzhongstagesignoff.documents-staging.windows-ppe.net:443/";
             System.setProperty(Configs.THINCLIENT_ENABLED, "true");
-            System.setProperty(Configs.THINCLIENT_ENDPOINT, thinclientEndpoint);
+            //System.setProperty(Configs.THINCLIENT_ENDPOINT, thinclientEndpoint);
             System.setProperty(Configs.HTTP2_ENABLED, "true");
 
             CosmosAsyncClient client  = new CosmosClientBuilder()

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdTokenStreamTests.java

@@ -0,0 +1,84 @@
+package com.azure.cosmos.implementation.directconnectivity.rntbd;
+
+import com.azure.cosmos.implementation.HttpConstants;
+import com.azure.cosmos.implementation.RxDocumentServiceRequest;
+import com.azure.cosmos.implementation.ThinClientStoreModel;
+import com.azure.cosmos.implementation.http.HttpHeaders;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import org.apache.commons.io.FileUtils;
+import org.mockito.Mockito;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.annotations.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
+
+import static com.azure.cosmos.implementation.directconnectivity.WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY;
+
+public class RntbdTokenStreamTests {
+    // Created this test for thin client testing, to make sure thin client headers are always special cased
+    private static final Logger logger = LoggerFactory.getLogger(RntbdTokenStreamTests.class);
+    @Test
+    public void parseDotNet() {
+        //dumpFile("E:\\Temp\\dotnetRead.bin");
+        //dumpFile("E:\\Temp\\javad6bd7634-35ed-44d3-849f-364456be3001.bin");
+
+        dumpFile("E:\\Temp\\dotnet_proxyInputRequest_Read_18.bin");
+        dumpFile("E:\\Temp\\java_proxyInputRequest_Read_29.bin");
+    }
+
+    private void dumpFile(String fileName) {
+        logger.error("FILENAME: {}", fileName);
+        File file = new File(fileName);
+        byte[] byteArray = null;
+        try {
+            byteArray = FileUtils.readFileToByteArray(file);
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+
+        ByteBuf content = Unpooled.wrappedBuffer(byteArray);
+
+        if (RntbdFramer.canDecodeHead(content)) {
+
+            final RntbdRequest request = RntbdRequest.decode(content);
+
+            if (request != null) {
+                logger.error("HEADERS: {}", request.getHeaders().dumpTokens());
+            }
+
+            logger.error("RNTBD REQUEST empty");
+        }
+    }
+
+    @Test(groups = { "unit" })
+    public void testThinClientSpecialCasing() {
+        RxDocumentServiceRequest mockRequest = Mockito.mock(RxDocumentServiceRequest.class);
+        Map<String, String> headers = new HashMap<>();
+        headers.put(EFFECTIVE_PARTITION_KEY, "13A141365AE34002732EE6DD02677CFC");
+        headers.put(HttpConstants.HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME, "globalDatabaseAccountName");
+        Mockito.doReturn(headers).when(mockRequest).getHeaders();
+        RntbdRequestArgs mockRntbdRequestArgs = Mockito.mock(RntbdRequestArgs.class);
+        Mockito.doReturn(mockRequest).when(mockRntbdRequestArgs).serviceRequest();
+        Mockito.doReturn("").when(mockRntbdRequestArgs).replicaPath();
+        Mockito.doReturn(0L).when(mockRntbdRequestArgs).transportRequestId();
+
+        RntbdRequestFrame mockRntbdRequestFrame = Mockito.mock(RntbdRequestFrame.class);
+        Mockito.doReturn(RntbdConstants.RntbdOperationType.Connection).when(mockRntbdRequestFrame).getOperationType();
+        RntbdRequestHeaders rntbdRequestHeaders = new RntbdRequestHeaders(mockRntbdRequestArgs, mockRntbdRequestFrame);
+
+        final ByteBuf out = Unpooled.buffer();
+        rntbdRequestHeaders.encode(out, false);
+    }
+
+    /*final class TestRntbdTokenStream extends RntbdTokenStream<RntbdConstants.RntbdRequestHeader> {
+        TestRntbdTokenStream(EnumSet<RntbdConstants.RntbdRequestHeader> headers, Map<Short, RntbdConstants.RntbdRequestHeader> ids, ByteBuf in, Class<RntbdConstants.RntbdRequestHeader> classType) {
+            super(headers, ids, in, classType);
+        }
+    }*/
+}

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

@@ -333,9 +333,7 @@ private SslContext sslContextInit(boolean serverCertVerificationDisabled, boolea
                     .forClient()
                     .sslProvider(SslContext.defaultClientProvider());
 
-            if (serverCertVerificationDisabled) {
-                sslContextBuilder.trustManager(InsecureTrustManagerFactory.INSTANCE); // disable cert verification
-            }
+            sslContextBuilder.trustManager(InsecureTrustManagerFactory.INSTANCE); // disable cert verification
 
             if (http2Enabled) {
                 sslContextBuilder

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,6 +285,8 @@ 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";

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

@@ -3436,6 +3436,8 @@ private Mono<ResourceResponse<Document>> readDocumentInternal(
                 getEffectiveClientContext(clientContextOverride),
                 OperationType.Read, ResourceType.Document, path, requestHeaders, options);
 
+            request.useThinProxy = Configs.isThinClientEnabled() && request.useGatewayMode ? true : false;
+            request.useThinProxy = true;
             DocumentServiceRequestContext requestContext = request.requestContext;
 
             options.getMarkE2ETimeoutInRequestContextCallbackHook().set(

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

@@ -3,23 +3,35 @@
 package com.azure.cosmos.implementation;
 
 import com.azure.cosmos.ConsistencyLevel;
+import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
+import com.azure.cosmos.implementation.directconnectivity.WFConstants;
+import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdConstants;
+import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdFramer;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdRequest;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdRequestArgs;
+import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdResponse;
 import com.azure.cosmos.implementation.http.HttpClient;
 import com.azure.cosmos.implementation.http.HttpHeaders;
 import com.azure.cosmos.implementation.http.HttpRequest;
-import com.azure.cosmos.models.FeedRange;
+import com.azure.cosmos.implementation.routing.HexConvert;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.HttpMethod;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import java.io.IOException;
 import java.net.URI;
+import java.nio.file.Files;
+import java.time.Instant;
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 
-import static com.azure.cosmos.implementation.directconnectivity.WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY;
 import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
 
 /**
@@ -30,6 +42,22 @@
  */
 public class ThinClientStoreModel extends RxGatewayStoreModel {
 
+    private static final Logger logger = LoggerFactory.getLogger(ThinClientStoreModel.class);
+
+    private static final List<RntbdConstants.RntbdRequestHeader> thinClientHeadersInOrder = Arrays.asList(
+        RntbdConstants.RntbdRequestHeader.EffectivePartitionKey,
+        RntbdConstants.RntbdRequestHeader.GlobalDatabaseAccountName,
+        RntbdConstants.RntbdRequestHeader.DatabaseName,
+        RntbdConstants.RntbdRequestHeader.CollectionName,
+        RntbdConstants.RntbdRequestHeader.CollectionRid,
+        //RntbdConstants.RntbdRequestHeader.ResourceId,
+        RntbdConstants.RntbdRequestHeader.PayloadPresent,
+        RntbdConstants.RntbdRequestHeader.DocumentName,
+        RntbdConstants.RntbdRequestHeader.AuthorizationToken,
+        RntbdConstants.RntbdRequestHeader.Date);
+
+
+
     public ThinClientStoreModel(
         DiagnosticsClientContext clientContext,
         ISessionContainer sessionContainer,
@@ -80,7 +108,37 @@ protected Map<String, String> getDefaultHeaders(
     @Override
     public URI getRootUri(RxDocumentServiceRequest request) {
         //var uri = this.globalEndpointManager.resolveServiceEndpoint(request).getThinClientLocationEndpoint();
-        return URI.create("https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/");
+        return URI.create("https://57.155.105.105:10650/"); // https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/
+    }
+
+    @Override
+    public StoreResponse unwrapToStoreResponse(RxDocumentServiceRequest request, int statusCode, HttpHeaders headers, ByteBuf content) {
+        if (content == null || content.readableBytes() == 0) {
+            return super.unwrapToStoreResponse(request, statusCode, headers, Unpooled.EMPTY_BUFFER);
+        }
+
+        Instant decodeStartTime = Instant.now();
+
+        if (RntbdFramer.canDecodeHead(content)) {
+
+            final RntbdResponse response = RntbdResponse.decode(content);
+
+            if (response != null) {
+                response.setDecodeEndTime(Instant.now());
+                response.setDecodeStartTime(decodeStartTime);
+
+                return super.unwrapToStoreResponse(
+                    request,
+                    response.getStatus().code(),
+                    new HttpHeaders(response.getHeaders().asMap(request.getActivityId())),
+                    response.getContent()
+                );
+            }
+
+            return super.unwrapToStoreResponse(request, statusCode, headers, null);
+        }
+
+        throw new IllegalStateException("Invalid rntbd response");
     }
 
     @Override
@@ -89,12 +147,14 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque
         // todo - neharao1 - validate b/w name() v/s toString()
         request.setThinclientHeaders(request.getOperationType().name(), request.getResourceType().name());
 
-        String epk = request.getPartitionKeyInternal().getEffectivePartitionKeyString(request.getPartitionKeyInternal(), request.getPartitionKeyDefinition());
+        byte[] epk = request.getPartitionKeyInternal().getEffectivePartitionKeyBytes(request.getPartitionKeyInternal(), request.getPartitionKeyDefinition());
         if (request.properties == null) {
             request.properties = new HashMap<>();
         }
-        request.properties.put(EFFECTIVE_PARTITION_KEY, epk);
-        //request.getHeaders().put(EFFECTIVE_PARTITION_KEY, epk);
+        //request.properties.put(EFFECTIVE_PARTITION_KEY, epk);
+        //request.properties.put(HttpConstants.HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME, "chukangzhongstagesignoff");
+        request.getHeaders().put(HttpConstants.HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME, "tiagonapoli-cdb-test"); // "chukangzhongstagesignoff"
+        request.getHeaders().put(WFConstants.BackendHeaders.COLLECTION_RID, "cLklAJU8SN0=");
         // todo - neharao1: no concept of a replica / service endpoint that can be passed
         RntbdRequestArgs rntbdRequestArgs = new RntbdRequestArgs(request);
 
@@ -103,24 +163,46 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque
         HttpHeaders headers = this.getHttpHeaders();
 
         RntbdRequest rntbdRequest = RntbdRequest.from(rntbdRequestArgs);
-
+        boolean success = rntbdRequest.setHeaderValue(
+            RntbdConstants.RntbdRequestHeader.EffectivePartitionKey,
+            epk);
+        if (!success) {
+            logger.error("Failed to update EPK to value {}", HexConvert.bytesToHex(epk));
+        } else {
+            logger.error("Updated EPK to value {}", HexConvert.bytesToHex(epk));
+        }
         // todo: neharao1 - validate whether Java heap buffer is okay v/s Direct buffer
         // todo: eventually need to use pooled buffer
         ByteBuf byteBuf = Unpooled.buffer();
 
+        logger.error("HEADERS: {}", rntbdRequest.getHeaders().dumpTokens());
+
         // todo: lifting the logic from there to encode the RntbdRequest instance into a ByteBuf (ByteBuf is a network compatible format)
         // todo: double-check with fabianm to see if RntbdRequest across RNTBD over TCP (Direct connectivity mode) is same as that when using ThinClient proxy
         // todo: need to conditionally add some headers (userAgent, replicaId/endpoint, etc)
-        rntbdRequest.encode(byteBuf);
+        rntbdRequest.encode(byteBuf, true);
+
+        byte[] contentAsByteArray = new byte[byteBuf.writerIndex()];
+        byteBuf.getBytes(0, contentAsByteArray, 0, byteBuf.writerIndex());
+
+        try {
+            Files.write(java.nio.file.Paths.get("E:\\Temp\\java" + UUID.randomUUID() + ".bin"), contentAsByteArray);
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
 
         return new HttpRequest(
             HttpMethod.POST,
             //requestUri,
-            URI.create("https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/"),
+            //https://thinclient-performancetests-eastus2.documents-staging.windows-ppe.net:10650
+            //https://cdb-ms-stage-eastus2-fe2-sql.eastus2.cloudapp.azure.com:10650
+            //https://57.155.105.105:10650/
+            // https://tiagonapoli-cdb-test-westus3.documents.azure.com:10650
+            URI.create("https://57.155.105.105:10650/"), // https://127.0.0.1:10650/ //https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/ // thinclient-performancetests-eastus2.documents-staging.windows-ppe.net  cdb-ms-stage-eastus2-fe2-sql.eastus2.cloudapp.azure.com
             //requestUri.getPort(),
             10650,
             headers,
-            Flux.just(byteBuf.array()));
+            Flux.just(contentAsByteArray));
     }
 
     private HttpHeaders getHttpHeaders() {