diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/BinaryEncodingHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/BinaryEncodingHelperTest.java new file mode 100644 index 000000000000..13d757144317 --- /dev/null +++ b/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); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/InternalObjectNodeTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/InternalObjectNodeTest.java index 171c01e712a1..064be36915f9 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/InternalObjectNodeTest.java +++ b/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(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyRangeServerBatchRequestTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyRangeServerBatchRequestTests.java index 19e67f4fbee6..14301165adf8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyRangeServerBatchRequestTests.java +++ b/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 Map serialize(T item) { + serializationCount.incrementAndGet(); + return Collections.singletonMap("id", item); + } + + @Override + public T deserialize(Map jsonNodeMap, Class classType) { + throw new UnsupportedOperationException(); + } + }; + List 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 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() { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchCodecTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchCodecTest.java new file mode 100644 index 000000000000..ff3b704aca9e --- /dev/null +++ b/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 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; + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTest.java new file mode 100644 index 000000000000..903da1a5e1e5 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTest.java @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.directconnectivity; + +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.json.CosmosBinaryJacksonCodec; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.netty.buffer.ByteBufInputStream; +import io.netty.buffer.Unpooled; +import org.testng.annotations.Test; + +import java.nio.charset.StandardCharsets; +import java.util.Collections; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public final class JsonNodeStorePayloadTest { + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test(groups = "unit") + public void parsesCosmosBinaryPayload() throws Exception { + JsonNode source = MAPPER.readTree("{\"id\":\"one\",\"value\":42}"); + byte[] content = CosmosBinaryJacksonCodec.encode(source); + + JsonNodeStorePayload payload = new JsonNodeStorePayload( + new ByteBufInputStream(Unpooled.wrappedBuffer(content)), + content.length, + Collections.emptyMap()); + + assertThat(payload.getPayload().path("id").textValue()).isEqualTo("one"); + assertThat(payload.getPayload().path("value").longValue()).isEqualTo(42L); + assertThat(payload.getResponsePayloadSize()).isEqualTo(content.length); + } + + @Test(groups = "unit") + public void mapsCorruptCosmosBinaryToStructuredParseFailure() { + byte[] content = new byte[] { (byte) 0x80, (byte) 0x82, (byte) 0xC3, 0x28 }; + + assertThatThrownBy(() -> new JsonNodeStorePayload( + new ByteBufInputStream(Unpooled.wrappedBuffer(content)), + content.length, + Collections.emptyMap())) + .isInstanceOf(CosmosException.class) + .satisfies(error -> assertThat(((CosmosException) error).getSubStatusCode()) + .isEqualTo(HttpConstants.SubStatusCodes.FAILED_TO_PARSE_SERVER_RESPONSE)); + } + + @Test(groups = "unit") + public void preservesRawRecordIoWithoutJsonParsing() { + byte[] content = new byte[] { (byte) 0x81, 1, 2, 3 }; + + JsonNodeStorePayload payload = new JsonNodeStorePayload( + new ByteBufInputStream(Unpooled.wrappedBuffer(content)), + content.length, + Collections.emptyMap()); + + assertThat(payload.getPayload()).isNull(); + assertThat(payload.getRawPayload()).isEqualTo(content).isNotSameAs(content); + } + + @Test(groups = "unit") + public void preservesJsonTextPayloadPath() { + byte[] content = "{\"id\":\"text\"}".getBytes(StandardCharsets.UTF_8); + + JsonNodeStorePayload payload = new JsonNodeStorePayload( + new ByteBufInputStream(Unpooled.wrappedBuffer(content)), + content.length, + Collections.emptyMap()); + + assertThat(payload.getPayload().path("id").textValue()).isEqualTo("text"); + assertThat(payload.getRawPayload()).isNull(); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdBinaryEncodingHeaderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdBinaryEncodingHeaderTest.java new file mode 100644 index 000000000000..2df57414d370 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdBinaryEncodingHeaderTest.java @@ -0,0 +1,112 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.directconnectivity.rntbd; + +import com.azure.cosmos.implementation.ContentSerializationFormat; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.directconnectivity.Uri; +import org.testng.annotations.Test; + +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; + +public final class RntbdBinaryEncodingHeaderTest { + private static final String BINARY_ENCODING_PROPERTY = "azure.cosmos.binaryEncodingEnabled"; + + @Test(groups = "unit") + public void mapsCosmosBinaryHeaderToRntbdToken() { + RxDocumentServiceRequest request = RxDocumentServiceRequest.create( + null, OperationType.Read, ResourceType.Document); + request.getHeaders().put( + HttpConstants.HttpHeaders.CONTENT_SERIALIZATION_FORMAT, + ContentSerializationFormat.CosmosBinary.toString()); + RntbdRequestFrame frame = new RntbdRequestFrame( + UUID.randomUUID(), + RntbdConstants.RntbdOperationType.Read, + RntbdConstants.RntbdResourceType.Document); + + RntbdRequestHeaders headers = new RntbdRequestHeaders( + new RntbdRequestArgs(request, new Uri("rntbd://localhost:10255")), frame); + RntbdToken token = headers.get(RntbdConstants.RntbdRequestHeader.ContentSerializationFormat); + + assertThat(token.isPresent()).isTrue(); + assertThat(token.getValue(Byte.class)) + .isEqualTo(RntbdConstants.RntbdContentSerializationFormat.CosmosBinary.id()); + } + + @Test(groups = "unit") + public void negotiatesBinaryQueryResponsesWithDotNetToken() { + System.setProperty(BINARY_ENCODING_PROPERTY, "true"); + try { + RxDocumentServiceRequest request = RxDocumentServiceRequest.create( + null, OperationType.Query, ResourceType.Document); + + RntbdToken token = createHeaders(request) + .get(RntbdConstants.RntbdRequestHeader.SupportedSerializationFormats); + + assertThat(RntbdConstants.RntbdRequestHeader.SupportedSerializationFormats.id()) + .isEqualTo((short) 0x00C4); + assertThat(token.getTokenType()).isEqualTo(RntbdTokenType.Byte); + assertThat(token.isPresent()).isTrue(); + assertThat(token.getValue(Byte.class)).isEqualTo((byte) 0x03); + } finally { + System.clearProperty(BINARY_ENCODING_PROPERTY); + } + } + + @Test(groups = "unit") + public void selectsBinaryForChangeFeedButNotGenericReadFeed() { + System.setProperty(BINARY_ENCODING_PROPERTY, "true"); + try { + RxDocumentServiceRequest changeFeed = RxDocumentServiceRequest.create( + null, OperationType.ReadFeed, ResourceType.Document); + changeFeed.getHeaders().put( + HttpConstants.HttpHeaders.A_IM, HttpConstants.A_IMHeaderValues.INCREMENTAL_FEED); + RxDocumentServiceRequest readFeed = RxDocumentServiceRequest.create( + null, OperationType.ReadFeed, ResourceType.Document); + + RntbdToken changeFeedFormat = createHeaders(changeFeed) + .get(RntbdConstants.RntbdRequestHeader.ContentSerializationFormat); + assertThat(changeFeedFormat.isPresent()).isTrue(); + assertThat(changeFeedFormat.getValue(Byte.class)) + .isEqualTo(RntbdConstants.RntbdContentSerializationFormat.CosmosBinary.id()); + assertThat(createHeaders(readFeed) + .get(RntbdConstants.RntbdRequestHeader.ContentSerializationFormat).isPresent()).isFalse(); + } finally { + System.clearProperty(BINARY_ENCODING_PROPERTY); + } + } + + @Test(groups = "unit") + public void doesNotNegotiateBinaryQueryForReadFeedOrThinClient() { + System.setProperty(BINARY_ENCODING_PROPERTY, "true"); + try { + RxDocumentServiceRequest readFeed = RxDocumentServiceRequest.create( + null, OperationType.ReadFeed, ResourceType.Document); + RxDocumentServiceRequest thinQuery = RxDocumentServiceRequest.create( + null, OperationType.Query, ResourceType.Document); + thinQuery.useThinClientMode = true; + + assertThat(createHeaders(readFeed) + .get(RntbdConstants.RntbdRequestHeader.SupportedSerializationFormats).isPresent()).isFalse(); + assertThat(createHeaders(thinQuery) + .get(RntbdConstants.RntbdRequestHeader.SupportedSerializationFormats).isPresent()).isFalse(); + } finally { + System.clearProperty(BINARY_ENCODING_PROPERTY); + } + } + + private static RntbdRequestHeaders createHeaders(RxDocumentServiceRequest request) { + RntbdRequestFrame frame = new RntbdRequestFrame( + UUID.randomUUID(), + RntbdConstants.RntbdOperationType.Query, + RntbdConstants.RntbdResourceType.Document); + return new RntbdRequestHeaders( + new RntbdRequestArgs(request, new Uri("rntbd://localhost:10255")), frame); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/json/CosmosBinaryCodecTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/json/CosmosBinaryCodecTest.java new file mode 100644 index 000000000000..b4ccb8b0d5ff --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/json/CosmosBinaryCodecTest.java @@ -0,0 +1,300 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.json; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.testng.annotations.Test; + +import java.math.BigInteger; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public final class CosmosBinaryCodecTest { + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test(groups = "unit") + public void roundTripsCoreValuesAndNativeBinary() { + byte[] payload = sequence(300); + Map value = new LinkedHashMap<>(); + value.put("empty", null); + value.put("enabled", true); + value.put("negative", -42L); + value.put("ratio", 1.25d); + value.put("values", Arrays.asList("one", false, 9L)); + value.put("payload", CosmosBinary.fromBytes(payload)); + + Object decoded = CosmosBinaryReader.decode(CosmosBinaryWriter.encode(value)); + + assertThat(decoded).isInstanceOf(Map.class); + Map result = (Map) decoded; + assertThat(result.get("empty")).isNull(); + assertThat(result.get("enabled")).isEqualTo(true); + assertThat(result.get("negative")).isEqualTo(-42L); + assertThat(result.get("ratio")).isEqualTo(1.25d); + assertThat(result.get("values")).isEqualTo(Arrays.asList("one", false, 9L)); + assertThat(result.get("payload")).isEqualTo(CosmosBinary.fromBytes(payload)); + } + + @Test(groups = "unit") + public void roundTripsLargeNestedScopesWithFourByteLengths() { + byte[] payload = sequence(70_000); + Map nested = new LinkedHashMap<>(); + nested.put("payload", CosmosBinary.fromBytes(payload)); + Map root = new LinkedHashMap<>(); + root.put("items", Arrays.asList(nested, "tail")); + root.put("enabled", true); + + Object decoded = CosmosBinaryReader.decode(CosmosBinaryWriter.encode(root)); + + Map decodedRoot = (Map) decoded; + List items = (List) decodedRoot.get("items"); + Map decodedNested = (Map) items.get(0); + assertThat(((CosmosBinary) decodedNested.get("payload")).toByteArray()).isEqualTo(payload); + assertThat(items.get(1)).isEqualTo("tail"); + } + + @Test(groups = "unit") + public void pojoByteArrayBecomesNativeBinaryInsteadOfBase64String() throws Exception { + PayloadItem item = new PayloadItem(); + item.id = "jackson-probe"; + item.payload = new byte[] { 0, 1, 2, 3, (byte) 0xFF }; + JsonNode itemTree = MAPPER.valueToTree(item); + + byte[] encoded = CosmosBinaryJacksonCodec.encode(itemTree); + JsonNode decoded = CosmosBinaryJacksonCodec.decode(encoded); + + assertThat(itemTree.path("payload").isBinary()).isTrue(); + assertThat(decoded.path("payload").isBinary()).isTrue(); + assertThat(decoded.path("payload").binaryValue()).isEqualTo(item.payload); + assertThat(decoded.path("id").textValue()).isEqualTo(item.id); + assertThat(encoded).contains((byte) 0xDD); + assertThat(new String(encoded, StandardCharsets.ISO_8859_1)) + .doesNotContain(java.util.Base64.getEncoder().encodeToString(item.payload)); + } + + @Test(groups = "unit") + public void roundTripsUnsigned64BitValues() { + BigInteger value = new BigInteger("18446744073709551615"); + + byte[] encoded = CosmosBinaryWriter.encode(value); + + assertThat(encoded[1] & 0xFF).isEqualTo(0xC7); + assertThat(CosmosBinaryReader.decode(encoded)).isEqualTo(value); + assertThat(CosmosBinaryJacksonCodec.decode(encoded).bigIntegerValue()).isEqualTo(value); + } + + @Test(groups = "unit") + public void decodesFloat16AndGuidScalars() { + assertThat(CosmosBinaryReader.decode(new byte[] { (byte) 0x80, (byte) 0xCF, 0x00, 0x3C })) + .isEqualTo(1.0d); + assertThat(CosmosBinaryReader.decode(new byte[] { (byte) 0x80, (byte) 0xCF, 0x00, (byte) 0xC0 })) + .isEqualTo(-2.0d); + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0xD3, + 0x33, 0x22, 0x11, 0x00, 0x55, 0x44, 0x77, 0x66, + (byte) 0x88, (byte) 0x99, (byte) 0xAA, (byte) 0xBB, + (byte) 0xCC, (byte) 0xDD, (byte) 0xEE, (byte) 0xFF + })).isEqualTo("00112233-4455-6677-8899-aabbccddeeff"); + } + + @Test(groups = "unit") + public void rejectsExcessiveNestingAndMalformedUtf8() { + byte[] nested = new byte[131]; + nested[0] = (byte) 0x80; + java.util.Arrays.fill(nested, 1, 130, (byte) 0xE1); + nested[130] = (byte) 0xD0; + assertThatThrownBy(() -> CosmosBinaryReader.decode(nested)) + .hasMessageContaining("Maximum nesting depth"); + + List value = new java.util.ArrayList<>(); + List root = value; + for (int index = 0; index < 130; index++) { + List child = new java.util.ArrayList<>(); + value.add(child); + value = child; + } + assertThatThrownBy(() -> CosmosBinaryWriter.encode(root)) + .hasMessageContaining("Maximum nesting depth"); + + assertThatThrownBy(() -> CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0x82, (byte) 0xC3, 0x28 + })).hasMessageContaining("Invalid UTF-8"); + } + + @Test(groups = "unit") + public void rejectsCountsThatExceedAvailablePayloadBeforeAllocating() { + assertThatThrownBy(() -> CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0xE7, + 0x00, 0x00, 0x00, 0x00, + (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, 0x7F + })).hasMessageContaining("Array count exceeds payload"); + assertThatThrownBy(() -> CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0xF3, (byte) 0xF1, (byte) 0xD7, + (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF + })).hasMessageContaining("Uniform nested array count exceeds input"); + } + + @Test(groups = "unit") + public void rejectsTruncatedAndOverflowingValues() { + assertThatThrownBy(() -> CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0xDF, 0x01, 0x00, 0x00, 0x00 + })).hasMessageContaining("Unexpected end"); + assertThatThrownBy(() -> CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0xEC, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, 0x7F + })).hasMessageContaining("Scope exceeds input"); + assertThatThrownBy(() -> CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0xC3, (byte) 0xFF + })).hasMessageContaining("target exceeds input"); + } + + @Test(groups = "unit") + public void decodesCompressedStrings() { + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, 0x78, 0x08, 0x10, 0x32, (byte) 0xBA, (byte) 0xDC + })).isEqualTo("0123abcd"); + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, 0x7B, 0x05, 0x41, 0x10, 0x32, 0x04 + })).isEqualTo("ABCDE"); + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, 0x7A, 0x04, 0x21, 0x43 + })).isEqualTo("0123"); + } + + @Test(groups = "unit") + public void decodesEncodedGuidAndBase64Strings() { + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, 0x75, + 0x10, 0x32, 0x54, 0x76, (byte) 0x98, (byte) 0xBA, (byte) 0xDC, (byte) 0xFE, + 0x10, 0x32, 0x54, 0x76, (byte) 0x98, (byte) 0xBA, (byte) 0xDC, (byte) 0xFE + })).isEqualTo("01234567-89ab-cdef-0123-456789abcdef"); + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, 0x71, 0x02, 0x01, 'h', 'e', 'l', 'l', 'o' + })).isEqualTo("aGVsbG8="); + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, 0x71, 0x02, (byte) 0xFE, 'h', 'e', 'l', 'l', 'o' + })).isEqualTo("aGVsbG8"); + } + + @Test(groups = "unit") + public void decodesUniformNumericArrays() { + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0xF0, (byte) 0xD7, 0x03, 0x01, 0x02, (byte) 0xFF + })).isEqualTo(Arrays.asList(1L, 2L, 255L)); + assertThat(CosmosBinaryReader.decode(new byte[] { + (byte) 0x80, (byte) 0xF2, (byte) 0xF0, (byte) 0xD8, + 0x02, 0x02, 0x01, (byte) 0xFF, 0x02, (byte) 0xFE + })).isEqualTo(Arrays.asList(Arrays.asList(1L, -1L), Arrays.asList(2L, -2L))); + } + + @Test(groups = "unit") + public void countedObjectAllowsDuplicateKeys() { + byte[] document = new byte[] { + (byte) 0x80, (byte) 0xED, 0x06, 0x02, + (byte) 0x81, 'a', 0x01, + (byte) 0x81, 'a', 0x02 + }; + + assertThat(CosmosBinaryReader.decode(document)) + .isEqualTo(java.util.Collections.singletonMap("a", 2L)); + } + + @Test(groups = "unit") + public void decodesReferenceStrings() { + byte[] document = concat( + new byte[] { (byte) 0x80, (byte) 0xEA, 0x16, (byte) 0x85 }, + "first".getBytes(StandardCharsets.UTF_8), + new byte[] { (byte) 0x86 }, + "repeat".getBytes(StandardCharsets.UTF_8), + new byte[] { (byte) 0x86 }, + "second".getBytes(StandardCharsets.UTF_8), + new byte[] { (byte) 0xC3, 0x09 }); + + Object decoded = CosmosBinaryReader.decode(document); + + assertThat(decoded).isInstanceOf(Map.class); + Map values = (Map) decoded; + assertThat(values.get("first")).isEqualTo("repeat"); + assertThat(values.get("second")).isEqualTo("repeat"); + } + + @Test(groups = "unit") + public void rejectsReferenceToAnotherReference() { + byte[] document = new byte[] { + (byte) 0x80, (byte) 0xE1, (byte) 0xC3, 0x04, (byte) 0xC3, 0x02 + }; + + assertThatThrownBy(() -> CosmosBinaryReader.decode(document)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("another reference"); + } + + @Test(groups = "unit") + public void everyMarkerEitherDecodesOrFailsWithRuntimeException() { + for (int marker = 0; marker <= 0xFF; marker++) { + try { + CosmosBinaryReader.decode(new byte[] { (byte) 0x80, (byte) marker }); + } catch (RuntimeException expected) { + // Unsupported or incomplete one-byte values must fail without Error subclasses or hangs. + } + } + } + + @Test(groups = "unit") + public void rejectsDocumentTruncatedAtEveryBoundary() { + Map nested = new LinkedHashMap<>(); + nested.put("payload", CosmosBinary.fromBytes(sequence(300))); + nested.put("values", Arrays.asList(-1L, 0L, 256L, "multibyte-€")); + byte[] document = CosmosBinaryWriter.encode(Collections.singletonMap("nested", nested)); + + for (int length = 0; length < document.length; length++) { + int truncationLength = length; + byte[] truncated = Arrays.copyOf(document, truncationLength); + assertThatThrownBy(() -> CosmosBinaryReader.decode(truncated)) + .as("truncation at byte %s", truncationLength) + .isInstanceOf(RuntimeException.class); + } + } + + @Test(groups = "unit") + public void rejectsAmbiguousByteArrays() { + assertThatThrownBy(() -> CosmosBinaryWriter.encode(new byte[] { 1, 2, 3 })) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("CosmosBinary.fromBytes"); + } + + private static byte[] concat(byte[]... values) { + int length = 0; + for (byte[] value : values) { + length += value.length; + } + byte[] result = new byte[length]; + int offset = 0; + for (byte[] value : values) { + System.arraycopy(value, 0, result, offset, value.length); + offset += value.length; + } + return result; + } + + private static byte[] sequence(int length) { + byte[] value = new byte[length]; + for (int index = 0; index < value.length; index++) { + value[index] = (byte) index; + } + return value; + } + + static final class PayloadItem { + public String id; + public byte[] payload; + } +} diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 9cbcee519b89..2ef8e19e53b3 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.82.0-beta.1 (Unreleased) #### Features Added +* Added opt-in Cosmos Binary encoding for Direct-mode item create, read, replace, upsert, delete, query, change feed, transactional batch, and bulk operations. Enable with `AZURE_COSMOS_BINARY_ENCODING_ENABLED=true` or `-Dazure.cosmos.binaryEncodingEnabled=true`. Native binary item fields avoid JSON base64 expansion; batch and bulk requests use the service's RecordIO/HybridRow wire format. * Enabled Gateway V2 (thin-client) data-plane routing by default for `Cosmos(Async)Client` instances configured with `gatewayMode` and HTTP/2, gated by an HTTP/2 connectivity probe with automatic fallback to Gateway V1. - See [PR 49437](https://github.com/Azure/azure-sdk-for-java/pull/49437) * Added support for QueryPlan and Execute Stored Procedure requests to be routed to Gateway V2. - See [PR 47759](https://github.com/Azure/azure-sdk-for-java/pull/47759) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/BinaryEncodingHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/BinaryEncodingHelper.java new file mode 100644 index 000000000000..1f97abf13c67 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/BinaryEncodingHelper.java @@ -0,0 +1,63 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation; + +import com.azure.cosmos.ConnectionMode; + +public final class BinaryEncodingHelper { + private BinaryEncodingHelper() { + } + + static boolean canUseBinaryEncoding( + ConnectionMode connectionMode, + ResourceType resourceType, + OperationType operationType, + RequestOptions options) { + + return Configs.isBinaryEncodingEnabled() + && connectionMode == ConnectionMode.DIRECT + && resourceType == ResourceType.Document + && isSupportedPointOperation(operationType) + && !hasTriggers(options); + } + + public static boolean canUseBinaryBatch(ConnectionMode connectionMode) { + return Configs.isBinaryEncodingEnabled() && connectionMode == ConnectionMode.DIRECT; + } + + public static boolean canUseBinaryQueryResponse(RxDocumentServiceRequest request) { + if (!canUseBinaryDocumentResponse(request)) { + return false; + } + OperationType operationType = request.getOperationType(); + return operationType == OperationType.Query || operationType == OperationType.SqlQuery; + } + + public static boolean canUseBinaryChangeFeedResponse(RxDocumentServiceRequest request) { + return canUseBinaryDocumentResponse(request) + && request.getOperationType() == OperationType.ReadFeed + && request.isChangeFeedRequest(); + } + + private static boolean canUseBinaryDocumentResponse(RxDocumentServiceRequest request) { + return Configs.isBinaryEncodingEnabled() + && request != null + && !request.useThinClientMode + && request.getResourceType() == ResourceType.Document; + } + + private static boolean isSupportedPointOperation(OperationType operationType) { + return operationType == OperationType.Create + || operationType == OperationType.Upsert + || operationType == OperationType.Replace + || operationType == OperationType.Read + || operationType == OperationType.Delete; + } + + private static boolean hasTriggers(RequestOptions options) { + return options != null + && ((options.getPreTriggerInclude() != null && !options.getPreTriggerInclude().isEmpty()) + || (options.getPostTriggerInclude() != null && !options.getPostTriggerInclude().isEmpty())); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index 686efabe8188..42ca621df566 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -94,6 +94,9 @@ public class Configs { public static final String HTTP_PENDING_ACQUIRE_MAX_COUNT = "COSMOS.HTTP_PENDING_ACQUIRE_MAX_COUNT"; public static final String HTTP_PENDING_ACQUIRE_MAX_COUNT_VARIABLE = "COSMOS_HTTP_PENDING_ACQUIRE_MAX_COUNT"; + static final String BINARY_ENCODING_ENABLED = "azure.cosmos.binaryEncodingEnabled"; + static final String BINARY_ENCODING_ENABLED_VARIABLE = "AZURE_COSMOS_BINARY_ENCODING_ENABLED"; + public static final String ITEM_SERIALIZATION_INCLUSION_MODE = "COSMOS.ITEM_SERIALIZATION_INCLUSION_MODE"; public static final String ITEM_SERIALIZATION_INCLUSION_MODE_VARIABLE = "COSMOS_ITEM_SERIALIZATION_INCLUSION_MODE"; @@ -1724,6 +1727,12 @@ public static EnumSet getDefaultOtelSpanAttributeNamingSc return AttributeNamingScheme.parse(DEFAULT_OTEL_SPAN_ATTRIBUTE_NAMING_SCHEME); } + static boolean isBinaryEncodingEnabled() { + return Boolean.parseBoolean(System.getProperty( + BINARY_ENCODING_ENABLED, + firstNonNull(emptyToNull(System.getenv(BINARY_ENCODING_ENABLED_VARIABLE)), Boolean.FALSE.toString()))); + } + public static boolean isNonParseableDocumentLoggingEnabled() { String isNonParseableDocumentLoggingEnabledAsString = System.getProperty( IS_NON_PARSEABLE_DOCUMENT_LOGGING_ENABLED, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/InternalObjectNode.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/InternalObjectNode.java index 3c87a8f0993a..bf59faa89586 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/InternalObjectNode.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/InternalObjectNode.java @@ -4,6 +4,7 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosItemSerializer; +import com.azure.cosmos.implementation.json.CosmosBinaryJacksonCodec; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -114,6 +115,17 @@ public static ByteBuffer serializeJsonToByteBuffer( String trackingId, boolean isIdValidationEnabled) { + return serializeJsonToByteBuffer( + cosmosItem, itemSerializer, trackingId, isIdValidationEnabled, false); + } + + public static ByteBuffer serializeJsonToByteBuffer( + Object cosmosItem, + CosmosItemSerializer itemSerializer, + String trackingId, + boolean isIdValidationEnabled, + boolean binaryEncodingEnabled) { + checkNotNull(itemSerializer, "Argument 'itemSerializer' must not be null."); if (cosmosItem instanceof InternalObjectNode) { InternalObjectNode internalObjectNode = ((InternalObjectNode) cosmosItem); @@ -121,37 +133,49 @@ public static ByteBuffer serializeJsonToByteBuffer( if (trackingId != null) { onAfterSerialization = (node) -> node.put(Constants.Properties.TRACKING_ID, trackingId); } - return internalObjectNode.serializeJsonToByteBuffer(itemSerializer, onAfterSerialization, isIdValidationEnabled); + return internalObjectNode.serializeJsonToByteBuffer( + itemSerializer, onAfterSerialization, isIdValidationEnabled, binaryEncodingEnabled); } else if (cosmosItem instanceof Document) { Document doc = (Document) cosmosItem; Consumer> onAfterSerialization = null; if (trackingId != null) { onAfterSerialization = (node) -> node.put(Constants.Properties.TRACKING_ID, trackingId); } - return doc.serializeJsonToByteBuffer(itemSerializer, onAfterSerialization, isIdValidationEnabled); + return doc.serializeJsonToByteBuffer( + itemSerializer, onAfterSerialization, isIdValidationEnabled, binaryEncodingEnabled); } else if (cosmosItem instanceof ObjectNode) { ObjectNode objectNode = (ObjectNode)cosmosItem; Consumer> onAfterSerialization = null; if (trackingId != null) { onAfterSerialization = (node) -> node.put(Constants.Properties.TRACKING_ID, trackingId); } - return (new InternalObjectNode(objectNode).serializeJsonToByteBuffer(itemSerializer, onAfterSerialization, isIdValidationEnabled)); + return new InternalObjectNode(objectNode).serializeJsonToByteBuffer( + itemSerializer, onAfterSerialization, isIdValidationEnabled, binaryEncodingEnabled); } else if (cosmosItem instanceof byte[]) { if (trackingId != null) { InternalObjectNode internalObjectNode = new InternalObjectNode((byte[]) cosmosItem); return internalObjectNode.serializeJsonToByteBuffer( itemSerializer, (node) -> node.put(Constants.Properties.TRACKING_ID, trackingId), - isIdValidationEnabled); + isIdValidationEnabled, + binaryEncodingEnabled); + } + byte[] serializedItem = (byte[]) cosmosItem; + if (!binaryEncodingEnabled + || CosmosBinaryJacksonCodec.isBinaryFormat(serializedItem)) { + return ByteBuffer.wrap(serializedItem); } - return ByteBuffer.wrap((byte[]) cosmosItem); + InternalObjectNode internalObjectNode = new InternalObjectNode(serializedItem); + return internalObjectNode.serializeJsonToByteBuffer( + itemSerializer, null, isIdValidationEnabled, true); } else { Consumer> onAfterSerialization = null; if (trackingId != null) { onAfterSerialization = (node) -> node.put(Constants.Properties.TRACKING_ID, trackingId); } - return Utils.serializeJsonToByteBuffer(itemSerializer, cosmosItem, onAfterSerialization, isIdValidationEnabled); + return Utils.serializeJsonToByteBuffer( + itemSerializer, cosmosItem, onAfterSerialization, isIdValidationEnabled, binaryEncodingEnabled); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/JsonSerializable.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/JsonSerializable.java index 985fdcc494bf..32863ff33af3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/JsonSerializable.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/JsonSerializable.java @@ -631,9 +631,24 @@ private ObjectNode fromJson(ByteBuffer json) { } } - public ByteBuffer serializeJsonToByteBuffer(CosmosItemSerializer itemSerializer, Consumer> onAfterSerialization, boolean isIdValidationEnabled) { + public ByteBuffer serializeJsonToByteBuffer( + CosmosItemSerializer itemSerializer, + Consumer> onAfterSerialization, + boolean isIdValidationEnabled) { + + return this.serializeJsonToByteBuffer( + itemSerializer, onAfterSerialization, isIdValidationEnabled, false); + } + + public ByteBuffer serializeJsonToByteBuffer( + CosmosItemSerializer itemSerializer, + Consumer> onAfterSerialization, + boolean isIdValidationEnabled, + boolean binaryEncodingEnabled) { + this.populatePropertyBag(); - return Utils.serializeJsonToByteBuffer(itemSerializer, propertyBag, onAfterSerialization, isIdValidationEnabled); + return Utils.serializeJsonToByteBuffer( + itemSerializer, propertyBag, onAfterSerialization, isIdValidationEnabled, binaryEncodingEnabled); } private String toJson(Object object) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index c31d590e6d35..f4ab68f18d97 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -2085,6 +2085,15 @@ public void validateReadConsistencyStrategy(ReadConsistencyStrategy readConsiste } private Map getRequestHeaders(RequestOptions options, ResourceType resourceType, OperationType operationType) { + return getRequestHeaders(options, resourceType, operationType, null); + } + + private Map getRequestHeaders( + RequestOptions options, + ResourceType resourceType, + OperationType operationType, + Boolean binaryEncodingEnabled) { + Map headers = new HashMap<>(); // Apply client-level additional headers first (e.g., workload-id from CosmosClientBuilder.additionalHeaders()) @@ -2096,6 +2105,16 @@ private Map getRequestHeaders(RequestOptions options, ResourceTy headers.put(HttpConstants.HttpHeaders.ALLOW_TENTATIVE_WRITES, Boolean.TRUE.toString()); } + boolean useBinaryEncoding = binaryEncodingEnabled != null + ? binaryEncodingEnabled + : BinaryEncodingHelper.canUseBinaryEncoding( + this.connectionPolicy.getConnectionMode(), resourceType, operationType, options); + if (useBinaryEncoding) { + headers.put( + HttpConstants.HttpHeaders.CONTENT_SERIALIZATION_FORMAT, + ContentSerializationFormat.CosmosBinary.toString()); + } + if (consistencyLevel != null) { // adding the "x-ms-consistency-level" header with consistency level stricter than the // account's default consistency level in Compute Gateway will result in a 400 Bad Request @@ -2415,7 +2434,14 @@ private Mono requestHeaders = this.getRequestHeaders(options, ResourceType.Document, operationType); + Map requestHeaders = this.getRequestHeaders( + options, ResourceType.Document, operationType, binaryEncodingEnabled); RxDocumentServiceRequest request = RxDocumentServiceRequest.create( getEffectiveClientContext(clientContextOverride), @@ -2472,7 +2499,7 @@ private Mono getBatchDocumentRequest(DocumentClientRet checkNotNull(serverBatchRequest, "expected non null serverBatchRequest"); Instant serializationStartTimeUTC = Instant.now(); - ByteBuffer content = ByteBuffer.wrap(Utils.getUTF8Bytes(serverBatchRequest.getRequestBody())); + ByteBuffer content = ByteBuffer.wrap(serverBatchRequest.getRequestBody()); Instant serializationEndTimeUTC = Instant.now(); SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics( @@ -3588,8 +3615,10 @@ private Mono> replaceDocumentInternal( logger.debug("Replacing a Document. documentLink: [{}]", documentLink); final String path = Utils.joinPath(documentLink, null); - final Map requestHeaders = - getRequestHeaders(options, ResourceType.Document, OperationType.Replace); + boolean binaryEncodingEnabled = BinaryEncodingHelper.canUseBinaryEncoding( + this.connectionPolicy.getConnectionMode(), ResourceType.Document, OperationType.Replace, options); + final Map requestHeaders = getRequestHeaders( + options, ResourceType.Document, OperationType.Replace, binaryEncodingEnabled); Instant serializationStartTimeUTC = Instant.now(); Consumer> onAfterSerialization = null; if (options != null) { @@ -3606,7 +3635,8 @@ private Mono> replaceDocumentInternal( CosmosItemSerializer serializerForContent = itemAlreadySerialized ? CosmosItemSerializer.DEFAULT_SERIALIZER : options.getEffectiveItemSerializer(); - ByteBuffer content = document.serializeJsonToByteBuffer(serializerForContent, onAfterSerialization, false); + ByteBuffer content = document.serializeJsonToByteBuffer( + serializerForContent, onAfterSerialization, false, binaryEncodingEnabled); Instant serializationEndTime = Instant.now(); SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceResponse.java index f9866f565957..dfe98179896b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceResponse.java @@ -121,6 +121,10 @@ public JsonNode getResponseBody() { return this.storeResponse.getResponseBodyAsJson(); } + public byte[] getResponseBodyAsBytes() { + return this.storeResponse.getResponseBodyAsBytes(); + } + public RequestTimeline getGatewayHttpRequestTimeline() { return gatewayHttpRequestTimeline; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java index a387e72a64ce..b50d1a6f5e1a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java @@ -8,6 +8,7 @@ import com.azure.cosmos.models.CosmosAdditionalHeaderName; import com.azure.cosmos.CosmosItemSerializer; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.json.CosmosBinaryJacksonCodec; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.DedicatedGatewayRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; @@ -737,17 +738,27 @@ public static void validateIdValue(Object itemIdValue) { } } - @SuppressWarnings("unchecked") public static ByteBuffer serializeJsonToByteBuffer( CosmosItemSerializer serializer, Object object, Consumer> onAfterSerialization, boolean isIdValidationEnabled) { + return serializeJsonToByteBuffer( + serializer, object, onAfterSerialization, isIdValidationEnabled, false); + } + + @SuppressWarnings("unchecked") + public static ByteBuffer serializeJsonToByteBuffer( + CosmosItemSerializer serializer, + Object object, + Consumer> onAfterSerialization, + boolean isIdValidationEnabled, + boolean binaryEncodingEnabled) { + checkArgument(serializer != null || object instanceof Map, "Argument 'serializer' must not be null."); try { - ByteBufferOutputStream byteBufferOutputStream = new ByteBufferOutputStream(ONE_KB); Map jsonTreeMap = (object instanceof Map && serializer == null) ? (Map) object : ensureItemSerializerAccessor().serializeSafe(serializer, object); @@ -771,8 +782,13 @@ public static ByteBuffer serializeJsonToByteBuffer( jsonNode = mapper.convertValue(jsonTreeMap, JsonNode.class); } - mapper.writeValue(byteBufferOutputStream, jsonNode); - return byteBufferOutputStream.asByteBuffer(); + if (binaryEncodingEnabled) { + return ByteBuffer.wrap(CosmosBinaryJacksonCodec.encode(jsonNode)); + } + + ByteBufferOutputStream outputStream = new ByteBufferOutputStream(ONE_KB); + mapper.writeValue(outputStream, jsonNode); + return outputStream.asByteBuffer(); } catch (IOException e) { // TODO moderakh: on serialization/deserialization failure we should throw CosmosException here and elsewhere throw new IllegalArgumentException("Failed to serialize the object into json", e); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java index 001347c37f88..2c3cd5d1f446 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java @@ -8,6 +8,7 @@ import com.azure.cosmos.CosmosBridgeInternal; import com.azure.cosmos.CosmosItemSerializer; import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.BinaryEncodingHelper; import com.azure.cosmos.implementation.RequestOptions; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchResponse; @@ -53,10 +54,14 @@ public Mono executeAsync() { List operations = this.cosmosBatch.getOperations(); checkArgument(operations.size() > 0, "Number of operations should be more than 0."); + AsyncDocumentClient documentClient = CosmosBridgeInternal.getAsyncDocumentClient(container.getDatabase()); + boolean hybridRow = BinaryEncodingHelper.canUseBinaryBatch( + documentClient.getConnectionPolicy().getConnectionMode()); final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createBatchRequest( this.cosmosBatch.getPartitionKeyValue(), operations, - this.effectiveItemSerializer); + this.effectiveItemSerializer, + hybridRow); request.setAtomicBatch(true); request.setShouldContinueOnError(false); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java index d62817896ca5..9ff6b0efa425 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -6,15 +6,19 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.JsonSerializable; import com.azure.cosmos.implementation.RxDocumentServiceResponse; +import com.azure.cosmos.implementation.Utils; import com.azure.cosmos.models.CosmosBatchOperationResult; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosItemOperation; +import com.azure.cosmos.implementation.batch.hybridrow.HybridRowBatchCodec; +import com.azure.cosmos.implementation.json.CosmosBinaryJacksonCodec; import com.azure.cosmos.models.ModelBridgeInternal; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ArrayNode; import com.fasterxml.jackson.databind.node.ObjectNode; import io.netty.handler.codec.http.HttpResponseStatus; +import java.io.IOException; import java.time.Duration; import java.util.ArrayList; import java.util.Iterator; @@ -40,18 +44,21 @@ public static CosmosBatchResponse fromDocumentServiceResponse( CosmosBatchResponse response = null; final JsonNode responseContentAsJson = documentServiceResponse.getResponseBody(); - - if (responseContentAsJson != null) { + byte[] rawResponse = documentServiceResponse.getResponseBodyAsBytes(); + + if (rawResponse != null && rawResponse.length > 0 && (rawResponse[0] & 0xFF) == 0x81) { + try { + response = populateFromHybridRow( + documentServiceResponse, request, shouldPromoteOperationStatus, rawResponse); + } catch (RuntimeException decodingFailure) { + response = deserializationFailure(documentServiceResponse, decodingFailure); + } + } else if (responseContentAsJson != null) { response = BatchResponseParser.populateFromResponseContent(documentServiceResponse, request, shouldPromoteOperationStatus); if (response == null) { // Convert any payload read failures as InternalServerError - response = ModelBridgeInternal.createCosmosBatchResponse( - HttpResponseStatus.INTERNAL_SERVER_ERROR.code(), - HttpConstants.SubStatusCodes.UNKNOWN, - "ServerResponseDeserializationFailure", - documentServiceResponse.getResponseHeaders(), - documentServiceResponse.getCosmosDiagnostics()); + response = deserializationFailure(documentServiceResponse, null); } } @@ -94,6 +101,85 @@ public static CosmosBatchResponse fromDocumentServiceResponse( return response; } + private static CosmosBatchResponse deserializationFailure( + RxDocumentServiceResponse serviceResponse, RuntimeException cause) { + String message = cause == null + ? "ServerResponseDeserializationFailure" + : "ServerResponseDeserializationFailure: " + cause.getMessage(); + return ModelBridgeInternal.createCosmosBatchResponse( + HttpResponseStatus.INTERNAL_SERVER_ERROR.code(), + HttpConstants.SubStatusCodes.UNKNOWN, + message, + serviceResponse.getResponseHeaders(), + serviceResponse.getCosmosDiagnostics()); + } + + private static CosmosBatchResponse populateFromHybridRow( + RxDocumentServiceResponse serviceResponse, + ServerBatchRequest request, + boolean shouldPromoteOperationStatus, + byte[] payload) { + + List operations = request.getOperations(); + List wireResults = HybridRowBatchCodec.decodeResponse( + payload, operations.size()); + if (wireResults.size() != operations.size()) { + return null; + } + List results = new ArrayList<>(wireResults.size()); + for (int index = 0; index < wireResults.size(); index++) { + HybridRowBatchCodec.Result wireResult = wireResults.get(index); + results.add(ModelBridgeInternal.createCosmosBatchResult( + wireResult.getETag(), + wireResult.getRequestCharge(), + decodeResourceBody(wireResult.getResourceBody()), + wireResult.getStatusCode(), + Duration.ofMillis(wireResult.getRetryAfterMilliseconds()), + wireResult.getSubStatusCode(), + operations.get(index))); + } + int statusCode = promotedStatus(serviceResponse.getStatusCode(), results, shouldPromoteOperationStatus); + int subStatusCode = statusCode == serviceResponse.getStatusCode() + ? BatchExecUtils.getSubStatusCode(serviceResponse.getResponseHeaders()) + : results.stream().filter(result -> result.getStatusCode() == statusCode) + .map(CosmosBatchOperationResult::getSubStatusCode).findFirst().orElse(HttpConstants.SubStatusCodes.UNKNOWN); + CosmosBatchResponse response = ModelBridgeInternal.createCosmosBatchResponse( + statusCode, subStatusCode, null, serviceResponse.getResponseHeaders(), serviceResponse.getCosmosDiagnostics()); + ModelBridgeInternal.addCosmosBatchResultInResponse(response, results); + return response; + } + + private static ObjectNode decodeResourceBody(byte[] resourceBody) { + if (resourceBody == null || resourceBody.length == 0) { + return null; + } + JsonNode value; + try { + value = CosmosBinaryJacksonCodec.isBinaryFormat(resourceBody) + ? CosmosBinaryJacksonCodec.decode(resourceBody) + : Utils.getSimpleObjectMapper().readTree(resourceBody); + } catch (IOException error) { + throw new IllegalStateException("Unable to decode batch resource body", error); + } + if (!(value instanceof ObjectNode)) { + throw new IllegalStateException("Batch resource body is not an object"); + } + return (ObjectNode) value; + } + + private static int promotedStatus( + int responseStatus, List results, boolean shouldPromoteOperationStatus) { + if (responseStatus == HttpResponseStatus.MULTI_STATUS.code() && shouldPromoteOperationStatus) { + for (CosmosBatchOperationResult result : results) { + if (result.getStatusCode() != HttpResponseStatus.FAILED_DEPENDENCY.code() + && result.getStatusCode() >= 400) { + return result.getStatusCode(); + } + } + } + return responseStatus; + } + private static CosmosBatchResponse populateFromResponseContent( final RxDocumentServiceResponse documentServiceResponse, final ServerBatchRequest request, @@ -147,10 +233,7 @@ private static CosmosBatchResponse populateFromResponseContent( } /** - * Read batch operation result result. - * - * TODO(rakkuma): Similarly hybrid row result needs to be parsed. - * Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856 + * Read a JSON batch operation result. * * @param objectNode having response for a single operation. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BulkExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BulkExecutor.java index 031399ae95ca..ba6dcd70ba67 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BulkExecutor.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BulkExecutor.java @@ -12,6 +12,8 @@ import com.azure.cosmos.CosmosItemSerializer; import com.azure.cosmos.ThrottlingRetryOptions; import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.BinaryEncodingHelper; +import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.CosmosBulkExecutionOptionsImpl; import com.azure.cosmos.implementation.CosmosSchedulers; import com.azure.cosmos.implementation.HttpConstants; @@ -592,7 +594,12 @@ private Flux> executeOperations( String pkRange = thresholds.getPartitionKeyRangeId(); ServerOperationBatchRequest serverOperationBatchRequest = - BulkExecutorUtil.createBatchRequest(operations, pkRange, this.maxMicroBatchPayloadSizeInBytes, this.effectiveItemSerializer); + BulkExecutorUtil.createBatchRequest( + operations, + pkRange, + this.maxMicroBatchPayloadSizeInBytes, + this.effectiveItemSerializer, + BinaryEncodingHelper.canUseBinaryBatch(this.docClientWrapper.getConnectionPolicy().getConnectionMode())); if (serverOperationBatchRequest.getBatchPendingOperations().size() > 0) { serverOperationBatchRequest.getBatchPendingOperations().forEach(groupSink::next); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BulkExecutorUtil.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BulkExecutorUtil.java index 22746b76a68d..018ca32af786 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BulkExecutorUtil.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BulkExecutorUtil.java @@ -43,14 +43,16 @@ static ServerOperationBatchRequest createBatchRequest( List operations, String partitionKeyRangeId, int maxMicroBatchPayloadSizeInBytes, - CosmosItemSerializer clientItemSerializer) { + CosmosItemSerializer clientItemSerializer, + boolean hybridRow) { return PartitionKeyRangeServerBatchRequest.createBatchRequest( partitionKeyRangeId, operations, maxMicroBatchPayloadSizeInBytes, Math.min(operations.size(), BatchRequestResponseConstants.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST), - clientItemSerializer); + clientItemSerializer, + hybridRow); } static void setRetryPolicyForBulk( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/HybridRowBatchMapper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/HybridRowBatchMapper.java new file mode 100644 index 000000000000..dfb08de5da0b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/HybridRowBatchMapper.java @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.CosmosItemSerializer; +import com.azure.cosmos.implementation.JsonSerializable; +import com.azure.cosmos.implementation.RequestOptions; +import com.azure.cosmos.implementation.batch.hybridrow.HybridRowBatchCodec; +import com.azure.cosmos.implementation.json.CosmosBinaryJacksonCodec; +import com.azure.cosmos.models.CosmosItemOperation; +import com.azure.cosmos.models.CosmosItemOperationType; +import com.azure.cosmos.models.IndexingDirective; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ObjectNode; + +import java.nio.charset.StandardCharsets; + +final class HybridRowBatchMapper { + private HybridRowBatchMapper() { + } + + static byte[] encodeOperation(CosmosItemOperation operation, CosmosItemSerializer serializer) { + if (!(operation instanceof CosmosItemOperationBase)) { + throw new IllegalArgumentException("Unsupported Cosmos batch operation: " + operation.getClass()); + } + CosmosItemOperationBase internalOperation = (CosmosItemOperationBase) operation; + JsonSerializable serialized = internalOperation.getSerializedOperation(serializer); + ObjectNode fields = serialized.getPropertyBag(); + WireOptions options = options(internalOperation); + return HybridRowBatchCodec.encodeOperation( + operation.getOperationType(), + text(fields, BatchRequestResponseConstants.FIELD_PARTITION_KEY), + operation.getId(), + resourceBody(operation.getOperationType(), fields), + options.indexingDirective, + options.ifMatch, + options.ifNoneMatch, + options.minimalReturn); + } + + private static byte[] resourceBody(CosmosItemOperationType operationType, ObjectNode fields) { + JsonNode value = fields.get(BatchRequestResponseConstants.FIELD_RESOURCE_BODY); + if (value == null || value.isNull()) { + return null; + } + if (operationType == CosmosItemOperationType.PATCH) { + return value.toString().getBytes(StandardCharsets.UTF_8); + } + return CosmosBinaryJacksonCodec.encode(value); + } + + private static WireOptions options(CosmosItemOperationBase operation) { + RequestOptions options = null; + if (operation instanceof ItemBatchOperation) { + options = ((ItemBatchOperation) operation).getRequestOptions(); + } else if (operation instanceof ItemBulkOperation) { + options = ((ItemBulkOperation) operation).getRequestOptions(); + } + if (options == null) { + return WireOptions.NONE; + } + IndexingDirective indexing = options.getIndexingDirective(); + boolean minimalReturn = options.isContentResponseOnWriteEnabled() != null + && !options.isContentResponseOnWriteEnabled(); + return new WireOptions( + indexing == null ? null : indexing.toString(), + options.getIfMatchETag(), + options.getIfNoneMatchETag(), + minimalReturn); + } + + private static final class WireOptions { + private static final WireOptions NONE = new WireOptions(null, null, null, false); + private final String indexingDirective; + private final String ifMatch; + private final String ifNoneMatch; + private final boolean minimalReturn; + + private WireOptions(String indexingDirective, String ifMatch, String ifNoneMatch, boolean minimalReturn) { + this.indexingDirective = indexingDirective; + this.ifMatch = ifMatch; + this.ifNoneMatch = ifNoneMatch; + this.minimalReturn = minimalReturn; + } + } + + private static String text(ObjectNode fields, String name) { + JsonNode value = fields.get(name); + return value == null || value.isNull() ? null : value.textValue(); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/PartitionKeyRangeServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/PartitionKeyRangeServerBatchRequest.java index 5e9224ddd65c..b114a8815326 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/PartitionKeyRangeServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/PartitionKeyRangeServerBatchRequest.java @@ -53,6 +53,18 @@ static ServerOperationBatchRequest createBatchRequest( final int maxOperationCount, final CosmosItemSerializer clientItemSerializer) { + return createBatchRequest( + partitionKeyRangeId, operations, maxBodyLength, maxOperationCount, clientItemSerializer, false); + } + + static ServerOperationBatchRequest createBatchRequest( + final String partitionKeyRangeId, + final List operations, + final int maxBodyLength, + final int maxOperationCount, + final CosmosItemSerializer clientItemSerializer, + final boolean hybridRow) { + final PartitionKeyRangeServerBatchRequest request = new PartitionKeyRangeServerBatchRequest( partitionKeyRangeId, maxBodyLength, @@ -61,7 +73,8 @@ static ServerOperationBatchRequest createBatchRequest( request.setAtomicBatch(false); request.setShouldContinueOnError(true); - List pendingOperations = request.createBodyOfBatchRequest(operations, clientItemSerializer); + List pendingOperations = request.createBodyOfBatchRequest( + operations, clientItemSerializer, hybridRow); return new ServerOperationBatchRequest(request, pendingOperations); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java index 6748be3cb080..9847bb43b96b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java @@ -7,9 +7,12 @@ import com.azure.cosmos.implementation.JsonSerializable; import com.azure.cosmos.implementation.Utils; import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; +import com.azure.cosmos.implementation.batch.hybridrow.HybridRowBatchCodec; import com.azure.cosmos.models.CosmosItemOperation; import com.fasterxml.jackson.databind.node.ArrayNode; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.List; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -24,6 +27,8 @@ public abstract class ServerBatchRequest { private final int maxOperationCount; private String requestBody; + private byte[] hybridRowRequestBody; + private boolean hybridRow; private List operations; private boolean isAtomicBatch = false; private boolean shouldContinueOnError = false; @@ -41,10 +46,9 @@ public abstract class ServerBatchRequest { /** * Adds as many operations as possible from the given list of operations. - * TODO(rakkuma): Similarly for hybrid row, request needs to be parsed to create a request body in any form. - * Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856 * - * Operations are added in order while ensuring the request body never exceeds {@link #maxBodyLength}. + * Operations are added in order while ensuring the encoded request body never exceeds + * {@link #maxBodyLength}. * * @param operations operations to be added; read-only. * @@ -52,22 +56,33 @@ public abstract class ServerBatchRequest { */ final List createBodyOfBatchRequest( final List operations, - final CosmosItemSerializer effectiveItemSerializer) { + final CosmosItemSerializer effectiveItemSerializer, + final boolean hybridRow) { checkNotNull(operations, "expected non-null operations"); - int totalSerializedLength = 0; + int totalSerializedLength = hybridRow ? 10 : 0; int totalOperationCount = 0; + List hybridRowOperations = hybridRow ? new ArrayList<>() : null; - final ArrayNode arrayNode = Utils.getSimpleObjectMapper().createArrayNode(); + final ArrayNode arrayNode = hybridRow ? null : Utils.getSimpleObjectMapper().createArrayNode(); for(CosmosItemOperation operation : operations) { JsonSerializable operationJsonSerializable; int operationSerializedLength; + byte[] hybridRowOperation = null; if (operation instanceof CosmosItemOperationBase) { - operationJsonSerializable = ((CosmosItemOperationBase) operation).getSerializedOperation(effectiveItemSerializer); - operationSerializedLength = ((CosmosItemOperationBase) operation).getSerializedLength(effectiveItemSerializer); + if (hybridRow) { + operationJsonSerializable = null; + hybridRowOperation = HybridRowBatchMapper.encodeOperation(operation, effectiveItemSerializer); + operationSerializedLength = hybridRowOperation.length + 13; + } else { + operationJsonSerializable = + ((CosmosItemOperationBase) operation).getSerializedOperation(effectiveItemSerializer); + operationSerializedLength = ((CosmosItemOperationBase) operation) + .getSerializedLength(effectiveItemSerializer); + } } else { throw new UnsupportedOperationException("Unknown CosmosItemOperation."); } @@ -81,21 +96,32 @@ final List createBodyOfBatchRequest( totalSerializedLength += operationSerializedLength; totalOperationCount++; - arrayNode.add(operationJsonSerializable.getPropertyBag()); + if (hybridRow) { + hybridRowOperations.add(hybridRowOperation); + } else { + arrayNode.add(operationJsonSerializable.getPropertyBag()); + } } - // TODO(rakkuma): This should change to byte array later as optimisation. + // TODO(rakkuma): The JSON path should change to byte array later as optimisation. // Issue: https://github.com/Azure/azure-sdk-for-java/issues/16112 - this.requestBody = arrayNode.toString(); + this.requestBody = hybridRow ? null : arrayNode.toString(); this.operations = operations.subList(0, totalOperationCount); + this.hybridRow = hybridRow; + if (hybridRow) { + this.hybridRowRequestBody = HybridRowBatchCodec.encodeRecordIo(hybridRowOperations); + } return operations.subList(totalOperationCount, operations.size()); } - public final String getRequestBody() { - checkState(this.requestBody != null, "expected non-null body"); - - return this.requestBody; + public final byte[] getRequestBody() { + if (hybridRow) { + checkState(this.hybridRowRequestBody != null, "expected non-null HybridRow body"); + return this.hybridRowRequestBody.clone(); + } + checkState(this.requestBody != null, "expected non-null JSON body"); + return this.requestBody.getBytes(StandardCharsets.UTF_8); } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java index 9765ca9a6721..0897df6ad017 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java @@ -39,11 +39,20 @@ static SinglePartitionKeyServerBatchRequest createBatchRequest( final List operations, final CosmosItemSerializer effectiveItemSerializer) { + return createBatchRequest(partitionKey, operations, effectiveItemSerializer, false); + } + + static SinglePartitionKeyServerBatchRequest createBatchRequest( + final PartitionKey partitionKey, + final List operations, + final CosmosItemSerializer effectiveItemSerializer, + final boolean hybridRow) { + checkNotNull(partitionKey, "expected non-null partitionKey"); checkNotNull(operations, "expected non-null operations"); final SinglePartitionKeyServerBatchRequest request = new SinglePartitionKeyServerBatchRequest(partitionKey); - request.createBodyOfBatchRequest(operations, effectiveItemSerializer); + request.createBodyOfBatchRequest(operations, effectiveItemSerializer, hybridRow); return request; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/BatchOperationCodec.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/BatchOperationCodec.java new file mode 100644 index 000000000000..197d3e98245f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/BatchOperationCodec.java @@ -0,0 +1,124 @@ +// 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.ByteBuf; +import io.netty.buffer.Unpooled; + +import java.util.Objects; + +/** Encodes one operation using the fixed Cosmos BatchOperation HybridRow V1 schema. */ +final class BatchOperationCodec { + private BatchOperationCodec() { + } + + static byte[] encode(Operation operation) { + Objects.requireNonNull(operation, "operation"); + int capacity = 64 + length(operation.partitionKey) + length(operation.id) + + length(operation.resourceBody) + operation.options.estimatedLength(); + ByteBuf output = Unpooled.buffer(capacity); + HybridRowWireWriter writer = new HybridRowWireWriter(output); + writer.writeByte(HybridRowBatchSchema.VERSION); + writer.writeInt32(HybridRowBatchSchema.OPERATION_SCHEMA_ID); + writer.writeByte(operation.presenceMask()); + writer.writeInt32(toWireOperation(operation.operationType)); + writer.writeInt32(HybridRowBatchSchema.DOCUMENT_RESOURCE_TYPE); + writer.writeVariable(operation.partitionKey); + writer.writeVariable(operation.id); + writer.writeVariable(operation.resourceBody); + writer.writeSparseString(HybridRowBatchSchema.OperationField.INDEXING_DIRECTIVE, + operation.options.indexingDirective); + writer.writeSparseString(HybridRowBatchSchema.OperationField.IF_MATCH, operation.options.ifMatch); + writer.writeSparseString(HybridRowBatchSchema.OperationField.IF_NONE_MATCH, operation.options.ifNoneMatch); + writer.writeSparseBoolean(HybridRowBatchSchema.OperationField.MINIMAL_RETURN, + operation.options.minimalReturn); + byte[] result = new byte[output.readableBytes()]; + output.readBytes(result); + return result; + } + + private static int toWireOperation(CosmosItemOperationType operationType) { + switch (operationType) { + case CREATE: return 0; + case PATCH: return 1; + case READ: return 2; + case DELETE: return 4; + case REPLACE: return 5; + case UPSERT: return 20; + default: throw new IllegalArgumentException("Unsupported batch operation: " + operationType); + } + } + + private static int length(String value) { + return value == null ? 0 : value.length() * 3; + } + + private static int length(byte[] value) { + return value == null ? 0 : value.length; + } + + static final class Operation { + private final CosmosItemOperationType operationType; + private String partitionKey; + private String id; + private byte[] resourceBody; + private Options options = Options.NONE; + + Operation(CosmosItemOperationType operationType) { + this.operationType = Objects.requireNonNull(operationType, "operationType"); + } + + Operation partitionKey(String value) { + this.partitionKey = value; + return this; + } + + Operation id(String value) { + this.id = value; + return this; + } + + Operation resourceBody(byte[] value) { + this.resourceBody = value == null ? null : value.clone(); + return this; + } + + Operation options(Options value) { + this.options = Objects.requireNonNull(value, "options"); + return this; + } + + private int presenceMask() { + return HybridRowBatchSchema.OperationField.OPERATION_TYPE.presenceBit() + | HybridRowBatchSchema.OperationField.RESOURCE_TYPE.presenceBit() + | present(partitionKey, HybridRowBatchSchema.OperationField.PARTITION_KEY) + | present(id, HybridRowBatchSchema.OperationField.ID) + | present(resourceBody, HybridRowBatchSchema.OperationField.RESOURCE_BODY); + } + + private static int present(Object value, HybridRowBatchSchema.OperationField field) { + return value == null ? 0 : field.presenceBit(); + } + } + + static final class Options { + static final Options NONE = new Options(null, null, null, false); + private final String indexingDirective; + private final String ifMatch; + private final String ifNoneMatch; + private final boolean minimalReturn; + + Options(String indexingDirective, String ifMatch, String ifNoneMatch, boolean minimalReturn) { + this.indexingDirective = indexingDirective; + this.ifMatch = ifMatch; + this.ifNoneMatch = ifNoneMatch; + this.minimalReturn = minimalReturn; + } + + private int estimatedLength() { + return length(indexingDirective) + length(ifMatch) + length(ifNoneMatch) + 16; + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/BatchResultCodec.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/BatchResultCodec.java new file mode 100644 index 000000000000..6ad51a5f272e --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/BatchResultCodec.java @@ -0,0 +1,94 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch.hybridrow; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.netty.handler.codec.CorruptedFrameException; + +/** Decodes one result using the fixed Cosmos BatchResult HybridRow V1 schema. */ +final class BatchResultCodec { + private BatchResultCodec() { + } + + static Result decode(byte[] row) { + ByteBuf input = Unpooled.wrappedBuffer(row); + HybridRowWireReader reader = new HybridRowWireReader(input); + reader.expectByte(HybridRowBatchSchema.VERSION, "result version"); + reader.expectInt32(HybridRowBatchSchema.RESULT_SCHEMA_ID, "result schema"); + int presence = reader.readUnsignedByte(); + int required = HybridRowBatchSchema.ResultField.STATUS_CODE.presenceBit() + | HybridRowBatchSchema.ResultField.SUB_STATUS_CODE.presenceBit(); + if ((presence & required) != required) { + throw new CorruptedFrameException("Batch result is missing required fields"); + } + Result result = new Result(reader.readInt32(), reader.readInt32()); + if ((presence & HybridRowBatchSchema.ResultField.ETAG.presenceBit()) != 0) { + result.eTag = reader.readVariableString(); + } + if ((presence & HybridRowBatchSchema.ResultField.RESOURCE_BODY.presenceBit()) != 0) { + result.resourceBody = reader.readVariableBytes(); + } + boolean retryAfterRead = false; + boolean requestChargeRead = false; + while (reader.readableBytes() != 0) { + int typeCode = reader.readUnsignedByte(); + int pathToken = readPathToken(reader); + HybridRowBatchSchema.ResultField field = HybridRowBatchSchema.ResultField.sparseFromPathToken(pathToken); + HybridRowBatchSchema.SparseType type = HybridRowBatchSchema.SparseType.fromCode(typeCode); + if (field == HybridRowBatchSchema.ResultField.RETRY_AFTER_MILLISECONDS + && type == HybridRowBatchSchema.SparseType.UINT32) { + if (retryAfterRead) { + throw new CorruptedFrameException("Duplicate BatchResult retry-after field"); + } + result.retryAfterMilliseconds = reader.readUnsignedInt32(); + retryAfterRead = true; + } else if (field == HybridRowBatchSchema.ResultField.REQUEST_CHARGE + && type == HybridRowBatchSchema.SparseType.FLOAT64) { + if (requestChargeRead) { + throw new CorruptedFrameException("Duplicate BatchResult request-charge field"); + } + result.requestCharge = reader.readFloat64(); + requestChargeRead = true; + } else { + throw new CorruptedFrameException( + "Unsupported BatchResult sparse field token/type: " + pathToken + "/" + typeCode); + } + } + return result; + } + + private static int readPathToken(HybridRowWireReader reader) { + int first = reader.readUnsignedByte(); + if ((first & 0x80) == 0) { + return first; + } + int second = reader.readUnsignedByte(); + if ((second & 0x80) != 0) { + throw new CorruptedFrameException("BatchResult path token is too large"); + } + return (first & 0x7F) | (second << 7); + } + + static final class Result { + private final int statusCode; + private final int subStatusCode; + private String eTag; + private byte[] resourceBody; + private long retryAfterMilliseconds; + private double requestCharge; + + private Result(int statusCode, int subStatusCode) { + this.statusCode = statusCode; + this.subStatusCode = subStatusCode; + } + + int getStatusCode() { return statusCode; } + int getSubStatusCode() { return subStatusCode; } + String getETag() { return eTag; } + byte[] getResourceBody() { return resourceBody == null ? null : resourceBody.clone(); } + long getRetryAfterMilliseconds() { return retryAfterMilliseconds; } + double getRequestCharge() { return requestCharge; } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchCodec.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchCodec.java new file mode 100644 index 000000000000..3c3cb006abe8 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchCodec.java @@ -0,0 +1,62 @@ +// 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 java.util.ArrayList; +import java.util.List; + +/** Internal facade for the fixed Cosmos HybridRow batch wire format. */ +public final class HybridRowBatchCodec { + private HybridRowBatchCodec() { + } + + public static byte[] encodeOperation( + CosmosItemOperationType operationType, + String partitionKey, + String id, + byte[] resourceBody, + String indexingDirective, + String ifMatch, + String ifNoneMatch, + boolean minimalReturn) { + + BatchOperationCodec.Options options = new BatchOperationCodec.Options( + indexingDirective, ifMatch, ifNoneMatch, minimalReturn); + return BatchOperationCodec.encode(new BatchOperationCodec.Operation(operationType) + .partitionKey(partitionKey) + .id(id) + .resourceBody(resourceBody) + .options(options)); + } + + public static byte[] encodeRecordIo(List operations) { + return RecordIoCodec.encode(operations); + } + + public static List decodeResponse(byte[] payload, int maxResults) { + List rows = RecordIoCodec.decode(payload, maxResults); + List results = new ArrayList<>(rows.size()); + for (byte[] row : rows) { + results.add(new Result(BatchResultCodec.decode(row))); + } + return results; + } + + public static final class Result { + private final BatchResultCodec.Result result; + + private Result(BatchResultCodec.Result result) { + this.result = result; + } + + public int getStatusCode() { return result.getStatusCode(); } + public int getSubStatusCode() { return result.getSubStatusCode(); } + public String getETag() { return result.getETag(); } + public byte[] getResourceBody() { return result.getResourceBody(); } + public long getRetryAfterMilliseconds() { return result.getRetryAfterMilliseconds(); } + public double getRequestCharge() { return result.getRequestCharge(); } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchSchema.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchSchema.java new file mode 100644 index 000000000000..f599db388017 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchSchema.java @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch.hybridrow; + +/** Wire constants for the fixed Cosmos batch HybridRow schemas. */ +final class HybridRowBatchSchema { + static final int VERSION = 0x81; + static final int DOCUMENT_RESOURCE_TYPE = 2; + static final int OPERATION_SCHEMA_ID = 2_145_473_648; + static final int RESULT_SCHEMA_ID = 2_145_473_649; + static final int RECORD_IO_SEGMENT_SCHEMA_ID = 2_147_473_648; + static final int RECORD_IO_RECORD_SCHEMA_ID = 2_147_473_649; + static final int RECORD_IO_SEGMENT_LENGTH = 10; + + private HybridRowBatchSchema() { + } + + enum OperationField { + OPERATION_TYPE(0, Storage.FIXED), + RESOURCE_TYPE(1, Storage.FIXED), + PARTITION_KEY(2, Storage.VARIABLE), + EFFECTIVE_PARTITION_KEY(3, Storage.VARIABLE), + ID(4, Storage.VARIABLE), + BINARY_ID(5, Storage.VARIABLE), + RESOURCE_BODY(6, Storage.VARIABLE), + INDEXING_DIRECTIVE(8, Storage.SPARSE), + IF_MATCH(9, Storage.SPARSE), + IF_NONE_MATCH(10, Storage.SPARSE), + TTL_SECONDS(11, Storage.SPARSE), + MINIMAL_RETURN(12, Storage.SPARSE); + + private final int index; + private final Storage storage; + + OperationField(int index, Storage storage) { + this.index = index; + this.storage = storage; + } + + int presenceBit() { + if (storage == Storage.SPARSE) { + throw new IllegalStateException(name() + " is sparse"); + } + return 1 << index; + } + + int pathToken() { + if (storage != Storage.SPARSE) { + throw new IllegalStateException(name() + " is not sparse"); + } + return index; + } + } + + enum ResultField { + STATUS_CODE(0, Storage.FIXED), + SUB_STATUS_CODE(1, Storage.FIXED), + ETAG(2, Storage.VARIABLE), + RESOURCE_BODY(3, Storage.VARIABLE), + // Sparse tokens include the schema's root token at index zero. + RETRY_AFTER_MILLISECONDS(5, Storage.SPARSE), + REQUEST_CHARGE(6, Storage.SPARSE); + + private final int index; + private final Storage storage; + + ResultField(int index, Storage storage) { + this.index = index; + this.storage = storage; + } + + int presenceBit() { + if (storage == Storage.SPARSE) { + throw new IllegalStateException(name() + " is sparse"); + } + return 1 << index; + } + + static ResultField sparseFromPathToken(int token) { + for (ResultField field : values()) { + if (field.storage == Storage.SPARSE && field.index == token) { + return field; + } + } + return null; + } + } + + enum SparseType { + BOOLEAN_TRUE(0x03), + UINT32(0x0B), + FLOAT64(0x10), + UTF8(0x14); + + private final int code; + + SparseType(int code) { + this.code = code; + } + + int code() { + return code; + } + + static SparseType fromCode(int code) { + for (SparseType type : values()) { + if (type.code == code) { + return type; + } + } + return null; + } + } + + private enum Storage { + FIXED, + VARIABLE, + SPARSE + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowWireReader.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowWireReader.java new file mode 100644 index 000000000000..a243371b2328 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowWireReader.java @@ -0,0 +1,118 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch.hybridrow; + +import io.netty.buffer.ByteBuf; +import io.netty.handler.codec.CorruptedFrameException; + +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.CodingErrorAction; +import java.nio.charset.StandardCharsets; +import java.util.Objects; + +final class HybridRowWireReader { + private final ByteBuf buffer; + + HybridRowWireReader(ByteBuf buffer) { + this.buffer = Objects.requireNonNull(buffer, "buffer"); + } + + int readableBytes() { + return buffer.readableBytes(); + } + + int readUnsignedByte() { + require(Byte.BYTES); + return buffer.readUnsignedByte(); + } + + int readInt32() { + require(Integer.BYTES); + return buffer.readIntLE(); + } + + int readNonNegativeInt32(String field) { + int value = readInt32(); + if (value < 0) { + throw corrupt(field + " is negative"); + } + return value; + } + + long readUnsignedInt32() { + require(Integer.BYTES); + return buffer.readUnsignedIntLE(); + } + + double readFloat64() { + require(Double.BYTES); + return buffer.readDoubleLE(); + } + + String readVariableString() { + byte[] value = readVariableBytes(); + try { + return StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(ByteBuffer.wrap(value)) + .toString(); + } catch (CharacterCodingException error) { + throw new CorruptedFrameException("Invalid UTF-8 in HybridRow string", error); + } + } + + byte[] readVariableBytes() { + int length = readVarUInt(); + require(length); + byte[] value = new byte[length]; + buffer.readBytes(value); + return value; + } + + ByteBuf readRetainedSlice(int length) { + require(length); + return buffer.readRetainedSlice(length); + } + + void expectByte(int expected, String field) { + int actual = readUnsignedByte(); + if (actual != expected) { + throw corrupt("Invalid " + field + ": " + actual); + } + } + + void expectInt32(int expected, String field) { + int actual = readInt32(); + if (actual != expected) { + throw corrupt("Invalid " + field + ": " + actual); + } + } + + private int readVarUInt() { + int value = 0; + for (int index = 0; index < 5; index++) { + int next = readUnsignedByte(); + if (index == 4 && (next & 0xF0) != 0) { + throw corrupt("Variable length exceeds int range"); + } + value |= (next & 0x7F) << (index * 7); + if ((next & 0x80) == 0) { + return value; + } + } + throw corrupt("Invalid variable length"); + } + + private void require(int length) { + if (length < 0 || buffer.readableBytes() < length) { + throw corrupt("Truncated HybridRow payload"); + } + } + + private static CorruptedFrameException corrupt(String message) { + return new CorruptedFrameException(message); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowWireWriter.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowWireWriter.java new file mode 100644 index 000000000000..b3d4792d8218 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowWireWriter.java @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch.hybridrow; + +import io.netty.buffer.ByteBuf; + +import java.nio.charset.StandardCharsets; +import java.util.Objects; + +final class HybridRowWireWriter { + private final ByteBuf buffer; + + HybridRowWireWriter(ByteBuf buffer) { + this.buffer = Objects.requireNonNull(buffer, "buffer"); + } + + void writeByte(int value) { + buffer.writeByte(value); + } + + void writeInt32(int value) { + buffer.writeIntLE(value); + } + + void writeFloat64(double value) { + buffer.writeDoubleLE(value); + } + + void writeVariable(String value) { + writeVariable(value == null ? null : value.getBytes(StandardCharsets.UTF_8)); + } + + void writeVariable(byte[] value) { + if (value == null) { + return; + } + writeVarUInt(value.length); + buffer.writeBytes(value); + } + + void writeSparseString(HybridRowBatchSchema.OperationField field, String value) { + if (value == null) { + return; + } + writeByte(HybridRowBatchSchema.SparseType.UTF8.code()); + writeVarUInt(field.pathToken()); + writeVariable(value); + } + + void writeSparseBoolean(HybridRowBatchSchema.OperationField field, boolean value) { + if (!value) { + return; + } + writeByte(HybridRowBatchSchema.SparseType.BOOLEAN_TRUE.code()); + writeVarUInt(field.pathToken()); + } + + void writeVarUInt(int value) { + if (value < 0) { + throw new IllegalArgumentException("value must be nonnegative"); + } + int remaining = value; + do { + int next = remaining & 0x7F; + remaining >>>= 7; + writeByte(remaining == 0 ? next : next | 0x80); + } while (remaining != 0); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/RecordIoCodec.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/RecordIoCodec.java new file mode 100644 index 000000000000..6bdf1dd804c3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/RecordIoCodec.java @@ -0,0 +1,85 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch.hybridrow; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.netty.handler.codec.CorruptedFrameException; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.zip.CRC32; + +/** Encodes and decodes the RecordIO envelope used by Cosmos batch requests. */ +final class RecordIoCodec { + private RecordIoCodec() { + } + + static byte[] encode(List records) { + Objects.requireNonNull(records, "records"); + int capacity = HybridRowBatchSchema.RECORD_IO_SEGMENT_LENGTH; + for (byte[] record : records) { + capacity = Math.addExact(capacity, Math.addExact(13, Objects.requireNonNull(record, "record").length)); + } + ByteBuf output = Unpooled.buffer(capacity, capacity); + HybridRowWireWriter writer = new HybridRowWireWriter(output); + writer.writeByte(HybridRowBatchSchema.VERSION); + writer.writeInt32(HybridRowBatchSchema.RECORD_IO_SEGMENT_SCHEMA_ID); + writer.writeByte(1); + writer.writeInt32(HybridRowBatchSchema.RECORD_IO_SEGMENT_LENGTH); + for (byte[] record : records) { + CRC32 crc = new CRC32(); + crc.update(record); + writer.writeByte(HybridRowBatchSchema.VERSION); + writer.writeInt32(HybridRowBatchSchema.RECORD_IO_RECORD_SCHEMA_ID); + writer.writeInt32(record.length); + writer.writeInt32((int) crc.getValue()); + output.writeBytes(record); + } + byte[] result = new byte[output.readableBytes()]; + output.readBytes(result); + return result; + } + + static List decode(byte[] payload, int maxRecords) { + if (maxRecords < 0) { + throw new IllegalArgumentException("maxRecords must be nonnegative"); + } + if (payload == null || payload.length == 0) { + return Collections.emptyList(); + } + ByteBuf input = Unpooled.wrappedBuffer(payload); + HybridRowWireReader reader = new HybridRowWireReader(input); + reader.expectByte(HybridRowBatchSchema.VERSION, "segment version"); + reader.expectInt32(HybridRowBatchSchema.RECORD_IO_SEGMENT_SCHEMA_ID, "segment schema"); + reader.expectByte(1, "segment presence"); + reader.expectInt32(HybridRowBatchSchema.RECORD_IO_SEGMENT_LENGTH, "segment length"); + List records = new ArrayList<>(); + while (reader.readableBytes() != 0) { + if (records.size() >= maxRecords) { + throw new CorruptedFrameException("RecordIO record count exceeds request operation count"); + } + reader.expectByte(HybridRowBatchSchema.VERSION, "record version"); + reader.expectInt32(HybridRowBatchSchema.RECORD_IO_RECORD_SCHEMA_ID, "record schema"); + int length = reader.readNonNegativeInt32("record length"); + long expectedCrc = reader.readUnsignedInt32(); + ByteBuf record = reader.readRetainedSlice(length); + try { + byte[] bytes = new byte[length]; + record.getBytes(record.readerIndex(), bytes); + CRC32 crc = new CRC32(); + crc.update(bytes); + if (crc.getValue() != expectedCrc) { + throw new CorruptedFrameException("RecordIO CRC mismatch"); + } + records.add(bytes); + } finally { + record.release(); + } + } + return records; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java index ae81cd22c1c3..50df7b0ec59d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java @@ -6,6 +6,7 @@ import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.json.CosmosBinaryJacksonCodec; import com.fasterxml.jackson.databind.JsonNode; import io.netty.buffer.ByteBufInputStream; import io.netty.util.internal.StringUtil; @@ -27,13 +28,17 @@ public class JsonNodeStorePayload implements StorePayload { private static final CharsetDecoder fallbackCharsetDecoder = getFallbackCharsetDecoder(); private final int responsePayloadSize; private final JsonNode jsonValue; + private final byte[] rawPayload; public JsonNodeStorePayload(ByteBufInputStream bufferStream, int readableBytes, Map responseHeaders) { if (readableBytes > 0) { this.responsePayloadSize = readableBytes; - this.jsonValue = parseJson(bufferStream, readableBytes, () -> responseHeaders); + byte[] payload = readPayload(bufferStream, readableBytes); + this.rawPayload = isHybridRow(payload) ? payload : null; + this.jsonValue = parsePayload(payload, () -> responseHeaders); } else { this.responsePayloadSize = 0; + this.rawPayload = null; this.jsonValue = null; } } @@ -50,22 +55,54 @@ public JsonNodeStorePayload( if (readableBytes > 0) { this.responsePayloadSize = readableBytes; - this.jsonValue = parseJson(bufferStream, readableBytes, () -> buildHeaderMap(headerNames, headerValues)); + byte[] payload = readPayload(bufferStream, readableBytes); + this.rawPayload = isHybridRow(payload) ? payload : null; + this.jsonValue = parsePayload(payload, () -> buildHeaderMap(headerNames, headerValues)); } else { this.responsePayloadSize = 0; + this.rawPayload = null; this.jsonValue = null; } } - private static JsonNode parseJson( - ByteBufInputStream bufferStream, - int readableBytes, + private static byte[] readPayload(ByteBufInputStream bufferStream, int readableBytes) { + byte[] bytes = new byte[readableBytes]; + try { + int offset = 0; + while (offset < bytes.length) { + int read = bufferStream.read(bytes, offset, bytes.length - offset); + if (read < 0) { + throw new IOException("Unexpected end of store payload"); + } + offset += read; + } + return bytes; + } catch (IOException error) { + throw new IllegalStateException("Unable to read store payload", error); + } + } + + private static JsonNode parsePayload( + byte[] bytes, Supplier> headersSupplier) { - byte[] bytes = new byte[readableBytes]; try { - bufferStream.read(bytes); + if (CosmosBinaryJacksonCodec.isBinaryFormat(bytes)) { + return CosmosBinaryJacksonCodec.decode(bytes); + } + if (isHybridRow(bytes)) { + return null; + } return Utils.getSimpleObjectMapper().readTree(bytes); + } catch (RuntimeException error) { + if (!CosmosBinaryJacksonCodec.isBinaryFormat(bytes)) { + throw error; + } + throw Utils.createCosmosException( + HttpConstants.StatusCodes.BADREQUEST, + HttpConstants.SubStatusCodes.FAILED_TO_PARSE_SERVER_RESPONSE, + new IllegalStateException("Unable to parse Cosmos Binary response.", error), + headersSupplier.get()); } catch (IOException e) { Map responseHeaders = headersSupplier.get(); if (fallbackCharsetDecoder != null) { @@ -92,6 +129,10 @@ private static JsonNode parseJson( } } + private static boolean isHybridRow(byte[] bytes) { + return bytes.length > 0 && (bytes[0] & 0xFF) == 0x81; + } + private static Map buildHeaderMap(String[] headerNames, String[] headerValues) { Map map = new HashMap<>(HttpUtils.mapCapacityForSize(headerNames.length)); for (int i = 0; i < headerNames.length; i++) { @@ -140,6 +181,10 @@ public JsonNode getPayload() { return jsonValue; } + byte[] getRawPayload() { + return rawPayload == null ? null : rawPayload.clone(); + } + private static CharsetDecoder getFallbackCharsetDecoder() { if (StringUtil.isNullOrEmpty(Configs.getCharsetDecoderErrorActionOnMalformedInput()) && StringUtil.isNullOrEmpty(Configs.getCharsetDecoderErrorActionOnUnmappedCharacter())) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java index afce8c917428..f2cb4b533456 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java @@ -217,6 +217,10 @@ public int getResponseBodyLength() { return this.responsePayload.getResponsePayloadSize(); } + public byte[] getResponseBodyAsBytes() { + return this.responsePayload == null ? null : this.responsePayload.getRawPayload(); + } + public long getLSN() { String lsnString = this.getHeaderValue(WFConstants.BackendHeaders.LSN); if (StringUtils.isNotEmpty(lsnString)) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java index f9eb191fc20c..29fd2ab10cc7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java @@ -84,6 +84,22 @@ public byte id() { } } + public enum RntbdSupportedSerializationFormats { + + JsonText((byte) 0x01), + CosmosBinary((byte) 0x02); + + private final byte id; + + RntbdSupportedSerializationFormats(final byte id) { + this.id = id; + } + + public byte id() { + return this.id; + } + } + @SuppressWarnings("UnstableApiUsage") enum RntbdContextHeader implements RntbdHeader { @@ -619,6 +635,7 @@ public enum RntbdRequestHeader implements RntbdHeader { SDKSupportedCapabilities((short) 0x00A2, RntbdTokenType.ULong, false), ChangeFeedWireFormatVersion((short) 0x00B2, RntbdTokenType.String, false), PriorityLevel((short) 0x00BF, RntbdTokenType.Byte, false), + SupportedSerializationFormats((short) 0x00C4, RntbdTokenType.Byte, false), GlobalDatabaseAccountName((short) 0x00CE, RntbdTokenType.String, false), PopulateQueryAdvice((short) 0x00DA, RntbdTokenType.Byte, false), ThroughputBucket((short)0x00DB, RntbdTokenType.Byte, false), diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java index d00f397790c9..591da204c10e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java @@ -6,6 +6,7 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ConsistencyLevel; import com.azure.cosmos.ReadConsistencyStrategy; +import com.azure.cosmos.implementation.BinaryEncodingHelper; import com.azure.cosmos.implementation.ContentSerializationFormat; import com.azure.cosmos.implementation.EnumerationDirection; import com.azure.cosmos.implementation.FanoutOperationState; @@ -47,6 +48,7 @@ import static com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdConstants.RntbdReadFeedKeyType; import static com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdConstants.RntbdRemoteStorageType; import static com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdConstants.RntbdRequestHeader; +import static com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdConstants.RntbdSupportedSerializationFormats; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @JsonFilter("RntbdToken") @@ -99,6 +101,8 @@ private static ImplementationBridgeHelpers.PriorityLevelHelper.PriorityLevelAcce this.addCollectionRemoteStorageSecurityIdentifier(headers); this.addConsistencyLevelHeader(headers); this.addContentSerializationFormat(headers); + this.addBinaryChangeFeedResponseFormat(request); + this.addSupportedSerializationFormats(request); this.addContinuationToken(request, headers); this.addDateHeader(headers); this.addDisableRUPerMinuteUsage(headers); @@ -299,6 +303,10 @@ private RntbdToken getContentSerializationFormat() { return this.get(RntbdRequestHeader.ContentSerializationFormat); } + private RntbdToken getSupportedSerializationFormats() { + return this.get(RntbdRequestHeader.SupportedSerializationFormats); + } + private RntbdToken getContinuationToken() { return this.get(RntbdRequestHeader.ContinuationToken); } @@ -785,6 +793,20 @@ private void addContentSerializationFormat(final Map headers) { } } + private void addBinaryChangeFeedResponseFormat(final RxDocumentServiceRequest request) { + if (BinaryEncodingHelper.canUseBinaryChangeFeedResponse(request)) { + this.getContentSerializationFormat().setValue(RntbdContentSerializationFormat.CosmosBinary.id()); + } + } + + private void addSupportedSerializationFormats(final RxDocumentServiceRequest request) { + if (BinaryEncodingHelper.canUseBinaryQueryResponse(request)) { + this.getSupportedSerializationFormats().setValue((byte) ( + RntbdSupportedSerializationFormats.JsonText.id() + | RntbdSupportedSerializationFormats.CosmosBinary.id())); + } + } + private void addContinuationToken(final RxDocumentServiceRequest request, final Map headers) { String value = request.getContinuation(); if (StringUtils.isEmpty(value)) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinary.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinary.java new file mode 100644 index 000000000000..c8865a893f12 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinary.java @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.json; + +import java.util.Arrays; +import java.util.Objects; + +/** An explicit native Cosmos binary scalar. */ +final class CosmosBinary { + private final byte[] value; + + private CosmosBinary(byte[] value) { + this.value = value; + } + + static CosmosBinary fromBytes(byte[] value) { + Objects.requireNonNull(value, "value"); + return new CosmosBinary(value.clone()); + } + + byte[] toByteArray() { + return value.clone(); + } + + int size() { + return value.length; + } + + byte[] unsafeValue() { + return value; + } + + @Override + public boolean equals(Object other) { + return other instanceof CosmosBinary && Arrays.equals(value, ((CosmosBinary) other).value); + } + + @Override + public int hashCode() { + return Arrays.hashCode(value); + } + + @Override + public String toString() { + return "CosmosBinary[" + value.length + " bytes]"; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryJacksonCodec.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryJacksonCodec.java new file mode 100644 index 000000000000..fd362cb8bde5 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryJacksonCodec.java @@ -0,0 +1,131 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.json; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.BinaryNode; +import com.fasterxml.jackson.databind.node.BooleanNode; +import com.fasterxml.jackson.databind.node.DoubleNode; +import com.fasterxml.jackson.databind.node.JsonNodeFactory; +import com.fasterxml.jackson.databind.node.BigIntegerNode; +import com.fasterxml.jackson.databind.node.LongNode; +import com.fasterxml.jackson.databind.node.NullNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.fasterxml.jackson.databind.node.TextNode; +import java.io.IOException; +import java.math.BigInteger; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** Bridges Jackson's item tree to Cosmos Binary without converting {@link BinaryNode} to base64. */ +public final class CosmosBinaryJacksonCodec { + private static final int MAX_NESTING_DEPTH = 128; + + private CosmosBinaryJacksonCodec() {} + + public static boolean isBinaryFormat(byte[] document) { + return document != null && document.length > 0 && (document[0] & 0xFF) == 0x80; + } + + public static byte[] encode(JsonNode value) { + if (value == null) { + throw new NullPointerException("value"); + } + return CosmosBinaryWriter.encode(toCosmosValue(value, 0)); + } + + public static JsonNode decode(byte[] document) { + if (!isBinaryFormat(document)) { + throw new IllegalArgumentException("Not a Cosmos Binary document"); + } + return toJsonNode(CosmosBinaryReader.decode(document), 0); + } + + private static Object toCosmosValue(JsonNode value, int depth) { + if (depth > MAX_NESTING_DEPTH) { + throw new IllegalArgumentException("Maximum nesting depth exceeded"); + } + if (value.isNull()) { + return null; + } + if (value.isTextual()) { + return value.textValue(); + } + if (value.isBoolean()) { + return value.booleanValue(); + } + if (value.isIntegralNumber()) { + return value.canConvertToLong() ? value.longValue() : value.bigIntegerValue(); + } + if (value.isFloatingPointNumber()) { + return value.doubleValue(); + } + if (value.isBinary()) { + try { + return CosmosBinary.fromBytes(value.binaryValue()); + } catch (IOException error) { + throw new IllegalArgumentException("Unable to read Jackson binary value", error); + } + } + if (value.isArray()) { + List values = new ArrayList<>(value.size()); + value.forEach(element -> values.add(toCosmosValue(element, depth + 1))); + return values; + } + if (value.isObject()) { + Map fields = new LinkedHashMap<>(); + Iterator> iterator = value.fields(); + while (iterator.hasNext()) { + Map.Entry field = iterator.next(); + fields.put(field.getKey(), toCosmosValue(field.getValue(), depth + 1)); + } + return fields; + } + throw new IllegalArgumentException("Unsupported Jackson node: " + value.getNodeType()); + } + + private static JsonNode toJsonNode(Object value, int depth) { + if (depth > MAX_NESTING_DEPTH) { + throw new IllegalArgumentException("Maximum nesting depth exceeded"); + } + JsonNodeFactory factory = JsonNodeFactory.instance; + if (value == null) { + return NullNode.instance; + } + if (value instanceof String) { + return TextNode.valueOf((String) value); + } + if (value instanceof Boolean) { + return BooleanNode.valueOf((Boolean) value); + } + if (value instanceof Long) { + return LongNode.valueOf((Long) value); + } + if (value instanceof BigInteger) { + return BigIntegerNode.valueOf((BigInteger) value); + } + if (value instanceof Double) { + return DoubleNode.valueOf((Double) value); + } + if (value instanceof CosmosBinary) { + return BinaryNode.valueOf(((CosmosBinary) value).toByteArray()); + } + if (value instanceof List) { + ArrayNode array = factory.arrayNode(((List) value).size()); + ((List) value).forEach(element -> array.add(toJsonNode(element, depth + 1))); + return array; + } + if (value instanceof Map) { + ObjectNode object = factory.objectNode(); + ((Map) value).forEach( + (name, element) -> object.set((String) name, toJsonNode(element, depth + 1))); + return object; + } + throw new IllegalArgumentException("Unsupported decoded value: " + value.getClass().getName()); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryReader.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryReader.java new file mode 100644 index 000000000000..251d8d0f0b3b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryReader.java @@ -0,0 +1,486 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.json; + +import java.math.BigInteger; +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.CodingErrorAction; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** Decodes the core Cosmos Binary JSON value types. */ +final class CosmosBinaryReader { + private static final int MAX_NESTING_DEPTH = 128; + private static final String[] SYSTEM_STRINGS = { + "$s", "$t", "$v", "_attachments", "_etag", "_rid", "_self", "_ts", + "attachments/", "coordinates", "geometry", "GeometryCollection", "id", "url", + "Value", "label", "LineString", "link", "MultiLineString", "MultiPoint", + "MultiPolygon", "name", "Name", "Type", "Point", "Polygon", "properties", + "type", "value", "Feature", "FeatureCollection", "_id" + }; + + private final byte[] bytes; + private int position; + + private CosmosBinaryReader(byte[] bytes) { + this.bytes = bytes; + } + + static Object decode(byte[] document) { + if (document == null) { + throw new NullPointerException("document"); + } + CosmosBinaryReader reader = new CosmosBinaryReader(document); + reader.require(reader.readUnsignedByte() == 0x80, "Not Cosmos Binary format"); + Object value = reader.readValue(0); + reader.require(reader.position == document.length, "Trailing content"); + return value; + } + + static byte[] readBinaryProperty(byte[] document, String propertyName) { + Object decoded = decode(document); + if (!(decoded instanceof Map)) { + throw new IllegalArgumentException("Cosmos document root is not an object"); + } + Object value = ((Map) decoded).get(propertyName); + if (!(value instanceof CosmosBinary)) { + throw new IllegalArgumentException("Property is not native binary: " + propertyName); + } + return ((CosmosBinary) value).toByteArray(); + } + + static String[] systemStrings() { + return SYSTEM_STRINGS.clone(); + } + + private Object readValue(int depth) { + require(depth <= MAX_NESTING_DEPTH, "Maximum nesting depth exceeded"); + int marker = readUnsignedByte(); + if (marker < 0x20) { + return (long) marker; + } + if (isStringMarker(marker)) { + return readString(marker); + } + switch (marker) { + case 0xC7: return readUnsignedLong(); + case 0xC8: return (long) readUnsignedByte(); + case 0xC9: return (long) readSignedInteger(2); + case 0xCA: return (long) readSignedInteger(4); + case 0xCB: return readSignedLong(); + case 0xCC: return Double.longBitsToDouble(readLongBits()); + case 0xCD: return (double) Float.intBitsToFloat((int) readUnsignedInteger(4)); + case 0xCE: return Double.longBitsToDouble(readLongBits()); + case 0xCF: return readFloat16(); + case 0xD0: return null; + case 0xD1: return false; + case 0xD2: return true; + case 0xD3: return readGuid(); + case 0xD7: return (long) readUnsignedByte(); + case 0xD8: return (long) (byte) readUnsignedByte(); + case 0xD9: return (long) readSignedInteger(2); + case 0xDA: return (long) readSignedInteger(4); + case 0xDB: return readSignedLong(); + case 0xDC: return readUnsignedInteger(4); + case 0xDD: return readBinary(1); + case 0xDE: return readBinary(2); + case 0xDF: return readBinary(4); + case 0xE0: return new ArrayList<>(); + case 0xE1: return readSingleArray(depth + 1); + case 0xE2: return readArray(readLength(1), -1, depth + 1); + case 0xE3: return readArray(readLength(2), -1, depth + 1); + case 0xE4: return readArray(readLength(4), -1, depth + 1); + case 0xE5: return readArray(readLength(1), readLength(1), depth + 1); + case 0xE6: return readArray(readLength(2), readLength(2), depth + 1); + case 0xE7: return readArray(readLength(4), readLength(4), depth + 1); + case 0xE8: return new LinkedHashMap<>(); + case 0xE9: return readSingleObject(depth + 1); + case 0xEA: return readObject(readLength(1), -1, depth + 1); + case 0xEB: return readObject(readLength(2), -1, depth + 1); + case 0xEC: return readObject(readLength(4), -1, depth + 1); + case 0xED: return readObject(readLength(1), readLength(1), depth + 1); + case 0xEE: return readObject(readLength(2), readLength(2), depth + 1); + case 0xEF: return readObject(readLength(4), readLength(4), depth + 1); + case 0xF0: return readUniformNumberArray(1); + case 0xF1: return readUniformNumberArray(2); + case 0xF2: return readUniformNumberArrayArray(1); + case 0xF3: return readUniformNumberArrayArray(2); + default: + throw unsupported(marker); + } + } + + private List readSingleArray(int depth) { + List values = new ArrayList<>(1); + values.add(readValue(depth)); + return values; + } + + private List readArray(int payloadLength, int expectedCount, int depth) { + int end = scopeEnd(payloadLength); + require(expectedCount < 0 || expectedCount <= payloadLength, "Array count exceeds payload"); + List values = new ArrayList<>(expectedCount < 0 ? 4 : expectedCount); + while (position < end) { + values.add(readValue(depth)); + } + require(position == end, "Array length mismatch"); + require(expectedCount < 0 || values.size() == expectedCount, "Array count mismatch"); + return values; + } + + private Map readSingleObject(int depth) { + Map values = new LinkedHashMap<>(1); + values.put(readFieldName(), readValue(depth)); + return values; + } + + private Map readObject(int payloadLength, int expectedCount, int depth) { + int end = scopeEnd(payloadLength); + require(expectedCount < 0 || expectedCount <= payloadLength / 2, "Object count exceeds payload"); + Map values = new LinkedHashMap<>(expectedCount < 0 ? 4 : expectedCount); + int actualCount = 0; + while (position < end) { + values.put(readFieldName(), readValue(depth)); + actualCount++; + } + require(position == end, "Object length mismatch"); + require(expectedCount < 0 || actualCount == expectedCount, "Object count mismatch"); + return values; + } + + private List readUniformNumberArray(int countWidth) { + int itemMarker = readUnsignedByte(); + int count = readLength(countWidth); + validateUniformNumberCount(itemMarker, count); + return readUniformNumbers(itemMarker, count); + } + + private List readUniformNumberArrayArray(int countWidth) { + int nestedMarker = readUnsignedByte(); + require( + nestedMarker == (countWidth == 1 ? 0xF0 : 0xF1), + "Uniform nested array marker does not match count width"); + int itemMarker = readUnsignedByte(); + int innerCount = readLength(countWidth); + int outerCount = readLength(countWidth); + int itemWidth = uniformNumberWidth(itemMarker); + long totalCount = (long) innerCount * outerCount; + require(totalCount <= (bytes.length - position) / itemWidth, "Uniform nested array count exceeds input"); + require(totalCount <= Integer.MAX_VALUE, "Uniform nested array is too large"); + List result = new ArrayList<>(outerCount); + for (int index = 0; index < outerCount; index++) { + result.add(readUniformNumbers(itemMarker, innerCount)); + } + return result; + } + + private void validateUniformNumberCount(int itemMarker, int count) { + int itemWidth = uniformNumberWidth(itemMarker); + require(count <= (bytes.length - position) / itemWidth, "Uniform array count exceeds input"); + } + + private int uniformNumberWidth(int itemMarker) { + switch (itemMarker) { + case 0xD7: + case 0xD8: return 1; + case 0xD9: return 2; + case 0xDA: + case 0xCD: return 4; + case 0xDB: + case 0xCE: return 8; + default: throw unsupported(itemMarker); + } + } + + private List readUniformNumbers(int itemMarker, int count) { + List values = new ArrayList<>(count); + for (int index = 0; index < count; index++) { + switch (itemMarker) { + case 0xD7: values.add((long) readUnsignedByte()); break; + case 0xD8: values.add((long) (byte) readUnsignedByte()); break; + case 0xD9: values.add((long) readSignedInteger(2)); break; + case 0xDA: values.add((long) readSignedInteger(4)); break; + case 0xDB: values.add(readSignedLong()); break; + case 0xCD: values.add((double) Float.intBitsToFloat((int) readUnsignedInteger(4))); break; + case 0xCE: values.add(Double.longBitsToDouble(readLongBits())); break; + default: throw unsupported(itemMarker); + } + } + return values; + } + + private String readFieldName() { + int marker = readUnsignedByte(); + require(isStringMarker(marker), String.format("Object key marker is 0x%02X", marker)); + return readString(marker); + } + + private boolean isStringMarker(int marker) { + return (marker >= 0x20 && marker < 0x68) + || (marker >= 0x71 && marker <= 0xC6); + } + + private String readString(int marker) { + if (marker >= 0x20 && marker < 0x40) { + return SYSTEM_STRINGS[marker - 0x20]; + } + if (marker >= 0x40 && marker < 0x68) { + throw unsupported(marker); // User dictionary strings need external dictionary state. + } + if (marker >= 0x71 && marker <= 0x7F) { + return readEncodedString(marker); + } + if (marker >= 0xC3 && marker <= 0xC6) { + return readReferenceString(marker - 0xC2); + } + int length; + if (marker >= 0x80 && marker < 0xC0) { + length = marker - 0x80; + } else if (marker == 0xC0) { + length = readLength(1); + } else if (marker == 0xC1) { + length = readLength(2); + } else if (marker == 0xC2) { + length = readLength(4); + } else { + throw unsupported(marker); + } + byte[] utf8 = take(length); + try { + return StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(ByteBuffer.wrap(utf8)) + .toString(); + } catch (CharacterCodingException error) { + throw new IllegalArgumentException("Invalid UTF-8 string at byte " + position, error); + } + } + + private String readEncodedString(int marker) { + switch (marker) { + case 0x71: return readBase64String(1, false); + case 0x72: return readBase64String(2, false); + case 0x73: return readBase64String(1, true); + case 0x74: return readBase64String(2, true); + case 0x75: return readGuidString(false, false); + case 0x76: return readGuidString(true, false); + case 0x77: return readGuidString(false, true); + case 0x78: return readFourBitString("0123456789abcdef"); + case 0x79: return readFourBitString("0123456789ABCDEF"); + case 0x7A: return readFourBitString(" 0123456789:-.TZ"); + case 0x7B: return readPackedString(4, true); + case 0x7C: return readPackedString(5, true); + case 0x7D: return readPackedString(6, true); + case 0x7E: return readPackedString(7, false); + case 0x7F: return readPackedString(7, false, 2); + default: throw unsupported(marker); + } + } + + private String readBase64String(int lengthWidth, boolean urlSafe) { + int groups = readLength(lengthWidth); + int paddingMarker = readUnsignedByte(); + int padding = paddingMarker > 2 ? (~paddingMarker) & 0xFF : paddingMarker; + int valueLength = groups * 4 - (paddingMarker > 2 ? padding : 0); + int byteCount = (groups * 4 - padding) * 3 / 4; + byte[] raw = take(byteCount); + String encoded = java.util.Base64.getEncoder().encodeToString(raw); + encoded = encoded.substring(0, valueLength); + return urlSafe ? encoded.replace('+', '-').replace('/', '_') : encoded; + } + + private String readGuidString(boolean upperCase, boolean quoted) { + byte[] value = take(16); + StringBuilder result = new StringBuilder(quoted ? 38 : 36); + if (quoted) { + result.append('"'); + } + for (int index = 0; index < value.length; index++) { + if (index == 4 || index == 6 || index == 8 || index == 10) { + result.append('-'); + } + int packed = value[index] & 0xFF; + result.append(hexDigit(packed & 0x0F, upperCase)); + result.append(hexDigit(packed >>> 4, upperCase)); + } + if (quoted) { + result.append('"'); + } + return result.toString(); + } + + private String readFourBitString(String alphabet) { + int characterCount = readLength(1); + byte[] packed = take(compressedLength(characterCount, 4)); + StringBuilder result = new StringBuilder(characterCount); + for (int index = 0; index < characterCount; index++) { + int value = packed[index / 2] & 0xFF; + int alphabetIndex = index % 2 == 0 ? value & 0x0F : value >>> 4; + result.append(alphabet.charAt(alphabetIndex)); + } + return result.toString(); + } + + private String readPackedString(int bits, boolean hasBaseCharacter) { + return readPackedString(bits, hasBaseCharacter, 1); + } + + private String readPackedString(int bits, boolean hasBaseCharacter, int lengthWidth) { + int characterCount = readLength(lengthWidth); + int baseCharacter = hasBaseCharacter ? readUnsignedByte() : 0; + byte[] packed = take(compressedLength(characterCount, bits)); + StringBuilder result = new StringBuilder(characterCount); + long bitBuffer = 0; + int bufferedBits = 0; + int sourceIndex = 0; + int mask = (1 << bits) - 1; + for (int index = 0; index < characterCount; index++) { + while (bufferedBits < bits) { + bitBuffer |= (long) (packed[sourceIndex++] & 0xFF) << bufferedBits; + bufferedBits += 8; + } + result.append((char) ((bitBuffer & mask) + baseCharacter)); + bitBuffer >>>= bits; + bufferedBits -= bits; + } + return result.toString(); + } + + private static int compressedLength(int characterCount, int bits) { + long length = ((long) characterCount * bits + 7) / 8; + if (length > Integer.MAX_VALUE) { + throw new IllegalArgumentException("Compressed string is too large"); + } + return (int) length; + } + + private static char hexDigit(int value, boolean upperCase) { + return (char) (value < 10 ? '0' + value : (upperCase ? 'A' : 'a') + value - 10); + } + + private String readReferenceString(int width) { + int targetOffset = readLength(width); + require(targetOffset < bytes.length, "Reference string target exceeds input"); + int returnPosition = position; + position = targetOffset; + int targetMarker = readUnsignedByte(); + require(isStringMarker(targetMarker), "Reference string target is not a string"); + require(targetMarker < 0xC3 || targetMarker > 0xC6, "Reference string target is another reference"); + String value = readString(targetMarker); + position = returnPosition; + return value; + } + + private CosmosBinary readBinary(int width) { + return CosmosBinary.fromBytes(take(readLength(width))); + } + + private int scopeEnd(int payloadLength) { + require(payloadLength <= bytes.length - position, "Scope exceeds input"); + return position + payloadLength; + } + + private int readUnsignedByte() { + require(position < bytes.length, "Unexpected end of input"); + return bytes[position++] & 0xFF; + } + + private int readLength(int width) { + long value = readUnsignedInteger(width); + if (value > Integer.MAX_VALUE) { + throw new IllegalArgumentException("Value is too large at byte " + position); + } + return (int) value; + } + + private long readUnsignedInteger(int width) { + require(position + width <= bytes.length, "Unexpected end of integer"); + long value = 0; + for (int index = 0; index < width; index++) { + value |= (long) (bytes[position++] & 0xFF) << (index * 8); + } + return value; + } + + private int readSignedInteger(int width) { + long value = readUnsignedInteger(width); + int shift = (Long.BYTES - width) * 8; + return (int) (value << shift >> shift); + } + + private long readSignedLong() { + return readLongBits(); + } + + private Number readUnsignedLong() { + byte[] littleEndian = take(8); + byte[] bigEndian = new byte[8]; + for (int index = 0; index < 8; index++) { + bigEndian[index] = littleEndian[7 - index]; + } + BigInteger value = new BigInteger(1, bigEndian); + return value.bitLength() < 64 ? value.longValue() : value; + } + + private double readFloat16() { + int bits = (int) readUnsignedInteger(2); + int sign = (bits >>> 15) & 1; + int exponent = (bits >>> 10) & 0x1F; + int fraction = bits & 0x3FF; + double value; + if (exponent == 0) { + value = Math.scalb((double) fraction, -24); + } else if (exponent == 0x1F) { + value = fraction == 0 ? Double.POSITIVE_INFINITY : Double.NaN; + } else { + value = Math.scalb((double) (0x400 + fraction), exponent - 25); + } + return sign == 0 ? value : -value; + } + + private String readGuid() { + byte[] value = take(16); + return String.format( + "%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x", + value[3] & 0xFF, value[2] & 0xFF, value[1] & 0xFF, value[0] & 0xFF, + value[5] & 0xFF, value[4] & 0xFF, + value[7] & 0xFF, value[6] & 0xFF, + value[8] & 0xFF, value[9] & 0xFF, + value[10] & 0xFF, value[11] & 0xFF, value[12] & 0xFF, + value[13] & 0xFF, value[14] & 0xFF, value[15] & 0xFF); + } + + private long readLongBits() { + require(position + 8 <= bytes.length, "Unexpected end of 64-bit value"); + long value = 0; + for (int index = 0; index < 8; index++) { + value |= (long) (bytes[position++] & 0xFF) << (index * 8); + } + return value; + } + + private byte[] take(int length) { + require(length <= bytes.length - position, "Unexpected end of value"); + byte[] value = Arrays.copyOfRange(bytes, position, position + length); + position += length; + return value; + } + + private IllegalArgumentException unsupported(int marker) { + return new IllegalArgumentException( + String.format("Unsupported Cosmos Binary marker 0x%02X at byte %d", marker, position - 1)); + } + + private void require(boolean condition, String message) { + if (!condition) { + throw new IllegalArgumentException(message + " at byte " + position); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryWriter.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryWriter.java new file mode 100644 index 000000000000..032f0d2f8cc1 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryWriter.java @@ -0,0 +1,302 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.json; + +import java.math.BigInteger; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** Encodes JSON-compatible values plus {@link CosmosBinary} into Cosmos Binary JSON. */ +final class CosmosBinaryWriter { + private static final int FORMAT_BINARY = 0x80; + private static final int STRING_LENGTH_BASE = 0x80; + private static final int STRING_LENGTH_1 = 0xC0; + private static final int STRING_LENGTH_2 = 0xC1; + private static final int STRING_LENGTH_4 = 0xC2; + private static final int NUMBER_UINT64 = 0xC7; + private static final int NUMBER_UINT8 = 0xC8; + private static final int NUMBER_INT16 = 0xC9; + private static final int NUMBER_INT32 = 0xCA; + private static final int NUMBER_INT64 = 0xCB; + private static final int NUMBER_DOUBLE = 0xCC; + private static final int NULL = 0xD0; + private static final int FALSE = 0xD1; + private static final int TRUE = 0xD2; + private static final int BINARY_LENGTH_1 = 0xDD; + private static final int BINARY_LENGTH_2 = 0xDE; + private static final int BINARY_LENGTH_4 = 0xDF; + private static final int ARRAY_EMPTY = 0xE0; + private static final int ARRAY_SINGLE = 0xE1; + private static final int ARRAY_LENGTH_1 = 0xE2; + private static final int ARRAY_LENGTH_2 = 0xE3; + private static final int ARRAY_LENGTH_4 = 0xE4; + private static final int OBJECT_EMPTY = 0xE8; + private static final int OBJECT_SINGLE = 0xE9; + private static final int OBJECT_LENGTH_1 = 0xEA; + private static final int OBJECT_LENGTH_2 = 0xEB; + private static final int OBJECT_LENGTH_4 = 0xEC; + private static final int MAX_SCOPE_PREFIX = 5; + private static final int MAX_NESTING_DEPTH = 128; + + private static final Map SYSTEM_STRING_IDS = systemStringIds(); + + private final Buffer buffer; + + private CosmosBinaryWriter(int initialCapacity) { + this.buffer = new Buffer(initialCapacity); + this.buffer.writeByte(FORMAT_BINARY); + } + + static byte[] encode(Object value) { + CosmosBinaryWriter writer = new CosmosBinaryWriter(256); + writer.writeValue(value, 0); + return writer.buffer.toByteArray(); + } + + private void writeValue(Object value, int depth) { + if (depth > MAX_NESTING_DEPTH) { + throw new IllegalArgumentException("Maximum nesting depth exceeded"); + } + if (value == null) { + buffer.writeByte(NULL); + } else if (value instanceof String) { + writeString((String) value, false); + } else if (value instanceof Boolean) { + buffer.writeByte((Boolean) value ? TRUE : FALSE); + } else if (value instanceof Byte || value instanceof Short + || value instanceof Integer || value instanceof Long) { + writeInteger(((Number) value).longValue()); + } else if (value instanceof BigInteger) { + writeBigInteger((BigInteger) value); + } else if (value instanceof Float || value instanceof Double) { + buffer.writeByte(NUMBER_DOUBLE); + buffer.writeLittleEndian(Double.doubleToRawLongBits(((Number) value).doubleValue()), 8); + } else if (value instanceof CosmosBinary) { + writeBinary((CosmosBinary) value); + } else if (value instanceof byte[]) { + throw new IllegalArgumentException( + "byte[] is ambiguous; wrap native bytes with CosmosBinary.fromBytes(...)"); + } else if (value instanceof Map) { + writeObject((Map) value, depth + 1); + } else if (value instanceof List) { + writeArray((List) value, depth + 1); + } else { + throw new IllegalArgumentException("Unsupported Cosmos value: " + value.getClass().getName()); + } + } + + private void writeObject(Map fields, int depth) { + if (fields.isEmpty()) { + buffer.writeByte(OBJECT_EMPTY); + return; + } + if (fields.size() == 1) { + buffer.writeByte(OBJECT_SINGLE); + writeFields(fields, depth); + return; + } + int scopeStart = buffer.reserve(MAX_SCOPE_PREFIX); + int payloadStart = buffer.position(); + writeFields(fields, depth); + finishScope(scopeStart, payloadStart, OBJECT_LENGTH_1, OBJECT_LENGTH_2, OBJECT_LENGTH_4); + } + + private void writeFields(Map fields, int depth) { + for (Map.Entry field : fields.entrySet()) { + if (!(field.getKey() instanceof String)) { + throw new IllegalArgumentException("Cosmos object keys must be strings"); + } + writeString((String) field.getKey(), true); + writeValue(field.getValue(), depth); + } + } + + private void writeArray(List values, int depth) { + if (values.isEmpty()) { + buffer.writeByte(ARRAY_EMPTY); + return; + } + if (values.size() == 1) { + buffer.writeByte(ARRAY_SINGLE); + writeValue(values.get(0), depth); + return; + } + int scopeStart = buffer.reserve(MAX_SCOPE_PREFIX); + int payloadStart = buffer.position(); + for (Object value : values) { + writeValue(value, depth); + } + finishScope(scopeStart, payloadStart, ARRAY_LENGTH_1, ARRAY_LENGTH_2, ARRAY_LENGTH_4); + } + + private void finishScope(int scopeStart, int payloadStart, int marker1, int marker2, int marker4) { + int payloadLength = buffer.position() - payloadStart; + int prefixLength; + int marker; + if (payloadLength <= 0xFF) { + prefixLength = 2; + marker = marker1; + } else if (payloadLength <= 0xFFFF) { + prefixLength = 3; + marker = marker2; + } else { + prefixLength = 5; + marker = marker4; + } + buffer.compact(payloadStart, payloadLength, scopeStart + prefixLength); + buffer.setPosition(scopeStart); + buffer.writeByte(marker); + buffer.writeLittleEndian(payloadLength, prefixLength - 1); + buffer.setPosition(scopeStart + prefixLength + payloadLength); + } + + private void writeString(String value, boolean fieldName) { + if (fieldName) { + Integer systemId = SYSTEM_STRING_IDS.get(value); + if (systemId != null) { + buffer.writeByte(0x20 + systemId); + return; + } + } + byte[] utf8 = value.getBytes(StandardCharsets.UTF_8); + if (utf8.length <= 63) { + buffer.writeByte(STRING_LENGTH_BASE + utf8.length); + } else if (utf8.length <= 0xFF) { + buffer.writeByte(STRING_LENGTH_1); + buffer.writeLittleEndian(utf8.length, 1); + } else if (utf8.length <= 0xFFFF) { + buffer.writeByte(STRING_LENGTH_2); + buffer.writeLittleEndian(utf8.length, 2); + } else { + buffer.writeByte(STRING_LENGTH_4); + buffer.writeLittleEndian(utf8.length, 4); + } + buffer.writeBytes(utf8); + } + + private void writeInteger(long value) { + if (value >= 0 && value < 32) { + buffer.writeByte((int) value); + } else if (value >= 0 && value <= 0xFF) { + buffer.writeByte(NUMBER_UINT8); + buffer.writeByte((int) value); + } else if (value >= Short.MIN_VALUE && value <= Short.MAX_VALUE) { + buffer.writeByte(NUMBER_INT16); + buffer.writeLittleEndian(value, 2); + } else if (value >= Integer.MIN_VALUE && value <= Integer.MAX_VALUE) { + buffer.writeByte(NUMBER_INT32); + buffer.writeLittleEndian(value, 4); + } else { + buffer.writeByte(NUMBER_INT64); + buffer.writeLittleEndian(value, 8); + } + } + + private void writeBigInteger(BigInteger value) { + if (value.signum() < 0 || value.bitLength() > 64) { + throw new IllegalArgumentException("Integer is outside unsigned 64-bit range: " + value); + } + if (value.bitLength() < 64) { + writeInteger(value.longValue()); + return; + } + buffer.writeByte(NUMBER_UINT64); + byte[] bigEndian = value.toByteArray(); + for (int index = 0; index < 8; index++) { + buffer.writeByte(bigEndian[bigEndian.length - 1 - index]); + } + } + + private void writeBinary(CosmosBinary value) { + byte[] bytes = value.unsafeValue(); + writeLengthPrefix(bytes.length, BINARY_LENGTH_1, BINARY_LENGTH_2, BINARY_LENGTH_4); + buffer.writeBytes(bytes); + } + + private void writeLengthPrefix(int length, int marker1, int marker2, int marker4) { + if (length <= 0xFF) { + buffer.writeByte(marker1); + buffer.writeLittleEndian(length, 1); + } else if (length <= 0xFFFF) { + buffer.writeByte(marker2); + buffer.writeLittleEndian(length, 2); + } else { + buffer.writeByte(marker4); + buffer.writeLittleEndian(length, 4); + } + } + + private static Map systemStringIds() { + String[] strings = CosmosBinaryReader.systemStrings(); + Map ids = new HashMap<>(); + for (int index = 0; index < strings.length; index++) { + ids.put(strings[index], index); + } + return ids; + } + + private static final class Buffer { + private byte[] bytes; + private int position; + + private Buffer(int capacity) { + this.bytes = new byte[capacity]; + } + + private int position() { + return position; + } + + private void setPosition(int position) { + this.position = position; + } + + private int reserve(int length) { + ensureCapacity(length); + int start = position; + position += length; + return start; + } + + private void writeByte(int value) { + ensureCapacity(1); + bytes[position++] = (byte) value; + } + + private void writeBytes(byte[] value) { + ensureCapacity(value.length); + System.arraycopy(value, 0, bytes, position, value.length); + position += value.length; + } + + private void writeLittleEndian(long value, int width) { + ensureCapacity(width); + for (int index = 0; index < width; index++) { + bytes[position++] = (byte) (value >>> (index * 8)); + } + } + + private void compact(int source, int length, int destination) { + System.arraycopy(bytes, source, bytes, destination, length); + } + + private byte[] toByteArray() { + return java.util.Arrays.copyOf(bytes, position); + } + + private void ensureCapacity(int additionalLength) { + int required = position + additionalLength; + if (required <= bytes.length) { + return; + } + int capacity = bytes.length; + while (capacity < required) { + capacity = Math.max(capacity << 1, required); + } + bytes = java.util.Arrays.copyOf(bytes, capacity); + } + } +}