Add opt-in Cosmos Binary encoding for Direct item, query, and batch operations

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

@@ -0,0 +1,113 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation;
+
+import com.azure.cosmos.ConnectionMode;
+import com.azure.cosmos.CosmosItemSerializer;
+import com.azure.cosmos.implementation.json.CosmosBinaryJacksonCodec;
+import com.fasterxml.jackson.databind.JsonNode;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Collections;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public final class BinaryEncodingHelperTest {
+    @AfterMethod(alwaysRun = true)
+    public void clearConfiguration() {
+        System.clearProperty(Configs.BINARY_ENCODING_ENABLED);
+    }
+
+    @Test(groups = "unit")
+    public void enablesSupportedDirectPointOperations() {
+        System.setProperty(Configs.BINARY_ENCODING_ENABLED, "true");
+
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Create, new RequestOptions())).isTrue();
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Upsert, new RequestOptions())).isTrue();
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Replace, new RequestOptions())).isTrue();
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Read, new RequestOptions())).isTrue();
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Delete, new RequestOptions())).isTrue();
+    }
+
+    @Test(groups = "unit")
+    public void serializesPojoByteArrayAsNativeBinaryWhenEnabled() throws Exception {
+        PayloadItem item = new PayloadItem();
+        item.id = "one";
+        item.payload = new byte[] { 0, 1, 2, 3, (byte) 0xFF };
+
+        ByteBuffer buffer = Utils.serializeJsonToByteBuffer(
+            CosmosItemSerializer.DEFAULT_SERIALIZER, item, null, true, true);
+        byte[] encoded = new byte[buffer.remaining()];
+        buffer.get(encoded);
+        JsonNode decoded = CosmosBinaryJacksonCodec.decode(encoded);
+
+        assertThat(encoded[0] & 0xFF).isEqualTo(0x80);
+        assertThat(decoded.path("payload").isBinary()).isTrue();
+        assertThat(decoded.path("payload").binaryValue()).isEqualTo(item.payload);
+    }
+
+    @Test(groups = "unit")
+    public void enablesOnlyDirectDocumentQueries() {
+        System.setProperty(Configs.BINARY_ENCODING_ENABLED, "true");
+        RxDocumentServiceRequest query = RxDocumentServiceRequest.create(
+            null, OperationType.Query, ResourceType.Document);
+        RxDocumentServiceRequest sqlQuery = RxDocumentServiceRequest.create(
+            null, OperationType.SqlQuery, ResourceType.Document);
+        RxDocumentServiceRequest readFeed = RxDocumentServiceRequest.create(
+            null, OperationType.ReadFeed, ResourceType.Document);
+        RxDocumentServiceRequest changeFeed = RxDocumentServiceRequest.create(
+            null, OperationType.ReadFeed, ResourceType.Document);
+        changeFeed.getHeaders().put(
+            HttpConstants.HttpHeaders.A_IM, HttpConstants.A_IMHeaderValues.INCREMENTAL_FEED);
+        RxDocumentServiceRequest avadChangeFeed = RxDocumentServiceRequest.create(
+            null, OperationType.ReadFeed, ResourceType.Document);
+        avadChangeFeed.getHeaders().put(
+            HttpConstants.HttpHeaders.A_IM, HttpConstants.A_IMHeaderValues.FULL_FIDELITY_FEED);
+        RxDocumentServiceRequest thinQuery = RxDocumentServiceRequest.create(
+            null, OperationType.Query, ResourceType.Document);
+        thinQuery.useThinClientMode = true;
+
+        assertThat(BinaryEncodingHelper.canUseBinaryQueryResponse(query)).isTrue();
+        assertThat(BinaryEncodingHelper.canUseBinaryQueryResponse(sqlQuery)).isTrue();
+        assertThat(BinaryEncodingHelper.canUseBinaryQueryResponse(readFeed)).isFalse();
+        assertThat(BinaryEncodingHelper.canUseBinaryQueryResponse(thinQuery)).isFalse();
+        assertThat(BinaryEncodingHelper.canUseBinaryChangeFeedResponse(changeFeed)).isTrue();
+        assertThat(BinaryEncodingHelper.canUseBinaryChangeFeedResponse(avadChangeFeed)).isTrue();
+        assertThat(BinaryEncodingHelper.canUseBinaryChangeFeedResponse(readFeed)).isFalse();
+    }
+
+    @Test(groups = "unit")
+    public void remainsOffByDefault() {
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Read, new RequestOptions())).isFalse();
+    }
+
+    @Test(groups = "unit")
+    public void excludesGatewayQueriesBatchesPatchesAndTriggers() {
+        System.setProperty(Configs.BINARY_ENCODING_ENABLED, "true");
+        RequestOptions triggered = new RequestOptions();
+        triggered.setPreTriggerInclude(Collections.singletonList("pre"));
+
+        assertThat(canUse(ConnectionMode.GATEWAY, OperationType.Read, new RequestOptions())).isFalse();
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Query, new RequestOptions())).isFalse();
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Batch, new RequestOptions())).isFalse();
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Patch, new RequestOptions())).isFalse();
+        assertThat(canUse(ConnectionMode.DIRECT, OperationType.Create, triggered)).isFalse();
+    }
+
+    static final class PayloadItem {
+        public String id;
+        public byte[] payload;
+    }
+
+    private static boolean canUse(
+        ConnectionMode connectionMode,
+        OperationType operationType,
+        RequestOptions options) {
+
+        return BinaryEncodingHelper.canUseBinaryEncoding(
+            connectionMode, ResourceType.Document, operationType, options);
+    }
+}

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

@@ -5,6 +5,7 @@
 
 import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.TestPojo;
+import com.azure.cosmos.implementation.json.CosmosBinaryJacksonCodec;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import org.testng.annotations.Test;
@@ -41,6 +42,20 @@ public void ByteArrayWithTrackingId() throws IOException {
         validateTrackingId(blob, expectedTrackingId);
     }
 
+    @Test(groups = {"unit"})
+    public void preSerializedJsonIsConvertedWhenBinaryEncodingIsEnabled() throws IOException {
+        byte[] json = "{\"id\":\"binary\",\"payload\":\"value\"}"
+            .getBytes(java.nio.charset.StandardCharsets.UTF_8);
+
+        ByteBuffer buffer = InternalObjectNode.serializeJsonToByteBuffer(
+            json, CosmosItemSerializer.DEFAULT_SERIALIZER, null, false, true);
+        byte[] encoded = new byte[buffer.remaining()];
+        buffer.get(encoded);
+
+        assertThat(CosmosBinaryJacksonCodec.isBinaryFormat(encoded)).isTrue();
+        assertThat(CosmosBinaryJacksonCodec.decode(encoded).path("id").textValue()).isEqualTo("binary");
+    }
+
     @Test(groups = {"unit"})
     public void internalObjectNodeWithTrackingId() throws IOException {
         String expectedTrackingId = UUID.randomUUID().toString();

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

@@ -3,6 +3,7 @@
 
 package com.azure.cosmos.implementation.batch;
 
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import com.azure.cosmos.models.CosmosItemOperation;
@@ -11,7 +12,10 @@
 import org.testng.annotations.Test;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -99,6 +103,51 @@ public void overflowsBasedOnCountWithOffset() {
         assertThat(serverOperationBatchRequest.getBatchPendingOperations().get(0).getId()).isEqualTo(operations.get(2).getId());
     }
 
+    @Test(groups = {"unit"}, timeOut = TIMEOUT)
+    public void hybridRowSerializesEachItemOnce() {
+        AtomicInteger serializationCount = new AtomicInteger();
+        CosmosItemSerializer serializer = new CosmosItemSerializer() {
+            @Override
+            public <T> Map<String, Object> serialize(T item) {
+                serializationCount.incrementAndGet();
+                return Collections.singletonMap("id", item);
+            }
+
+            @Override
+            public <T> T deserialize(Map<String, Object> jsonNodeMap, Class<T> classType) {
+                throw new UnsupportedOperationException();
+            }
+        };
+        List<CosmosItemOperation> operations = Collections.singletonList(new ItemBulkOperation<>(
+            CosmosItemOperationType.CREATE, "one", PartitionKey.NONE, null, "one", null));
+
+        PartitionKeyRangeServerBatchRequest.createBatchRequest(
+            "0", operations, Integer.MAX_VALUE, 100, serializer, true);
+
+        assertThat(serializationCount).hasValue(1);
+    }
+
+    @Test(groups = {"unit"}, timeOut = TIMEOUT)
+    public void hybridRowUsesExactEncodedSizeForOverflow() {
+        JsonSerializable item = new JsonSerializable();
+        item.set("payload", StringUtils.repeat("x", 128));
+        List<CosmosItemOperation> operations = new ArrayList<>();
+        operations.add(new ItemBulkOperation<>(
+            CosmosItemOperationType.CREATE, "one", PartitionKey.NONE, null, item, null));
+        operations.add(new ItemBulkOperation<>(
+            CosmosItemOperationType.CREATE, "two", PartitionKey.NONE, null, item, null));
+
+        ServerOperationBatchRequest single = PartitionKeyRangeServerBatchRequest.createBatchRequest(
+            "0", operations.subList(0, 1), Integer.MAX_VALUE, 100, null, true);
+        int exactOneOperationLength = single.getBatchRequest().getRequestBody().length;
+        ServerOperationBatchRequest split = PartitionKeyRangeServerBatchRequest.createBatchRequest(
+            "0", operations, exactOneOperationLength, 100, null, true);
+
+        assertThat(split.getBatchRequest().getOperations()).hasSize(1);
+        assertThat(split.getBatchPendingOperations()).hasSize(1);
+        assertThat(split.getBatchRequest().getRequestBody()).hasSize(exactOneOperationLength);
+    }
+
     @Test(groups = {"unit"}, timeOut = TIMEOUT * 10)
     public void partitionKeyRangeServerBatchRequestSizeTests() {
 

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchCodecTest.java

@@ -0,0 +1,110 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation.batch.hybridrow;
+
+import com.azure.cosmos.models.CosmosItemOperationType;
+import io.netty.buffer.Unpooled;
+import io.netty.handler.codec.CorruptedFrameException;
+import org.testng.annotations.Test;
+
+import java.util.Collections;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public final class HybridRowBatchCodecTest {
+    private static final String OPERATION =
+        "817054e17f4300000000020000004a80ea472c95687962726964726f772d676f6c64656e2d6974656d"
+            + "8573686172649a687962726964726f772d676f6c64656e2d706172746974696f6e3c86676f6c64656e"
+            + "8374746cc92c01";
+    private static final String RECORD_IO =
+        "81f0d8ff7f010a00000081f1d8ff7f590000009d24a5cc" + OPERATION;
+
+    @Test(groups = "unit")
+    public void matchesDotNetOperationAndRecordIo() {
+        byte[] body = hex("80ea472c95687962726964726f772d676f6c64656e2d6974656d8573686172649a687962"
+            + "726964726f772d676f6c64656e2d706172746974696f6e3c86676f6c64656e8374746cc92c01");
+
+        byte[] operation = BatchOperationCodec.encode(
+            new BatchOperationCodec.Operation(CosmosItemOperationType.CREATE).resourceBody(body));
+
+        assertThat(operation).isEqualTo(hex(OPERATION));
+        assertThat(RecordIoCodec.encode(Collections.singletonList(operation))).isEqualTo(hex(RECORD_IO));
+        assertThat(RecordIoCodec.decode(hex(RECORD_IO), 1)).singleElement().isEqualTo(operation);
+    }
+
+    @Test(groups = "unit")
+    public void matchesDotNetReadAndSparseOptions() {
+        String expected = "817054e17f1302000000020000001d687962726964726f772d6f7074696f6e732d676f6c64656e"
+            + "2d6974656d140a09726561642d65746167";
+        BatchOperationCodec.Options options = new BatchOperationCodec.Options(null, null, "read-etag", false);
+
+        byte[] operation = BatchOperationCodec.encode(new BatchOperationCodec.Operation(CosmosItemOperationType.READ)
+            .id("hybridrow-options-golden-item").options(options));
+
+        assertThat(operation).isEqualTo(hex(expected));
+    }
+
+    @Test(groups = "unit")
+    public void decodesDotNetBatchResult() {
+        String response = "81f0d8ff7f010a00000081f1d8ff7f3f000000d5437d92817154e17f07c800000000000000"
+            + "262230303030613034312d303030302d303830302d303030302d366135373631623130303030221006"
+            + "5555555555552540";
+
+        List<byte[]> rows = RecordIoCodec.decode(hex(response), 1);
+        BatchResultCodec.Result result = BatchResultCodec.decode(rows.get(0));
+
+        assertThat(result.getStatusCode()).isEqualTo(200);
+        assertThat(result.getSubStatusCode()).isZero();
+        assertThat(result.getETag()).isEqualTo("\"0000a041-0000-0800-0000-6a5761b10000\"");
+        assertThat(result.getRequestCharge()).isEqualTo(10.666666666666666d);
+    }
+
+    @Test(groups = "unit")
+    public void rejectsMoreRecordsThanRequestedOperations() {
+        byte[] payload = RecordIoCodec.encode(java.util.Arrays.asList(new byte[0], new byte[0]));
+
+        assertThatThrownBy(() -> RecordIoCodec.decode(payload, 1))
+            .isInstanceOf(CorruptedFrameException.class)
+            .hasMessageContaining("record count");
+    }
+
+    @Test(groups = "unit")
+    public void rejectsDuplicateSparseResultFields() {
+        byte[] row = hex("817154e17f03c8000000000000000b05010000000b0502000000");
+
+        assertThatThrownBy(() -> BatchResultCodec.decode(row))
+            .isInstanceOf(CorruptedFrameException.class)
+            .hasMessageContaining("Duplicate");
+    }
+
+    @Test(groups = "unit")
+    public void rejectsMalformedUtf8InVariableStrings() {
+        HybridRowWireReader reader = new HybridRowWireReader(
+            Unpooled.wrappedBuffer(new byte[] { 0x01, (byte) 0x80 }));
+
+        assertThatThrownBy(reader::readVariableString)
+            .isInstanceOf(CorruptedFrameException.class)
+            .hasMessageContaining("UTF-8");
+    }
+
+    @Test(groups = "unit")
+    public void rejectsTruncatedAndCorruptRecordIo() {
+        byte[] valid = hex(RECORD_IO);
+        assertThatThrownBy(() -> RecordIoCodec.decode(java.util.Arrays.copyOf(valid, valid.length - 1), 1))
+            .isInstanceOf(CorruptedFrameException.class);
+        valid[19] ^= 1;
+        assertThatThrownBy(() -> RecordIoCodec.decode(valid, 1)).isInstanceOf(CorruptedFrameException.class)
+            .hasMessageContaining("CRC");
+    }
+
+    private static byte[] hex(String value) {
+        byte[] bytes = new byte[value.length() / 2];
+        for (int index = 0; index < bytes.length; index++) {
+            bytes[index] = (byte) Integer.parseInt(value.substring(index * 2, index * 2 + 2), 16);
+        }
+        return bytes;
+    }
+}