From b89e7b40531b8e7e2c8810eb4cece44d742f1d38 Mon Sep 17 00:00:00 2001 From: Julien Ellie Date: Wed, 15 Jul 2026 15:05:51 -0700 Subject: [PATCH 1/4] Add Cosmos binary JSON codec Implement the Cosmos Binary value format, Jackson tree bridging, native binary scalars, compressed values, strict UTF-8 decoding, bounded nesting, and hostile-input validation. --- .../json/CosmosBinaryCodecTest.java | 300 +++++++++++ .../implementation/json/CosmosBinary.java | 48 ++ .../json/CosmosBinaryJacksonCodec.java | 131 +++++ .../json/CosmosBinaryReader.java | 486 ++++++++++++++++++ .../json/CosmosBinaryWriter.java | 302 +++++++++++ 5 files changed, 1267 insertions(+) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/json/CosmosBinaryCodecTest.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinary.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryJacksonCodec.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryReader.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/json/CosmosBinaryWriter.java 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/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); + } + } +} From be6eac0a3a115b6a7cb73f3fc029411ae4fd4579 Mon Sep 17 00:00:00 2001 From: Julien Ellie Date: Wed, 15 Jul 2026 15:07:08 -0700 Subject: [PATCH 2/4] Add HybridRow batch wire codec Implement the fixed Cosmos BatchOperation and BatchResult schemas, RecordIO framing, CRC validation, exact .NET golden vectors, response-count bounds, and a single internal facade. --- .../hybridrow/HybridRowBatchCodecTest.java | 110 ++++++++++++++++ .../batch/hybridrow/BatchOperationCodec.java | 124 ++++++++++++++++++ .../batch/hybridrow/BatchResultCodec.java | 94 +++++++++++++ .../batch/hybridrow/HybridRowBatchCodec.java | 62 +++++++++ .../batch/hybridrow/HybridRowBatchSchema.java | 121 +++++++++++++++++ .../batch/hybridrow/HybridRowWireReader.java | 118 +++++++++++++++++ .../batch/hybridrow/HybridRowWireWriter.java | 70 ++++++++++ .../batch/hybridrow/RecordIoCodec.java | 85 ++++++++++++ 8 files changed, 784 insertions(+) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchCodecTest.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/BatchOperationCodec.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/BatchResultCodec.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchCodec.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowBatchSchema.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowWireReader.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/HybridRowWireWriter.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/hybridrow/RecordIoCodec.java 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/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; + } +} From ebe27fbeead3261960122f8e0c2bb600badb3451 Mon Sep 17 00:00:00 2001 From: Julien Ellie Date: Wed, 15 Jul 2026 15:08:54 -0700 Subject: [PATCH 3/4] Enable opt-in Cosmos binary encoding Add the Direct-mode feature gate and integrate Cosmos Binary with point operations, transactional batches, and bulk. Preserve HybridRow responses, decode structured batch results, size requests from encoded records, and retain JSON behavior for Gateway mode. --- .../BinaryEncodingHelperTest.java | 113 ++++++++++++++++++ .../InternalObjectNodeTest.java | 15 +++ ...titionKeyRangeServerBatchRequestTests.java | 49 ++++++++ .../JsonNodeStorePayloadTest.java | 77 ++++++++++++ .../rntbd/RntbdBinaryEncodingHeaderTest.java | 112 +++++++++++++++++ .../implementation/BinaryEncodingHelper.java | 63 ++++++++++ .../azure/cosmos/implementation/Configs.java | 9 ++ .../implementation/InternalObjectNode.java | 36 +++++- .../implementation/JsonSerializable.java | 19 ++- .../implementation/RxDocumentClientImpl.java | 42 ++++++- .../RxDocumentServiceResponse.java | 4 + .../azure/cosmos/implementation/Utils.java | 24 +++- .../implementation/batch/BatchExecutor.java | 7 +- .../batch/BatchResponseParser.java | 107 +++++++++++++++-- .../implementation/batch/BulkExecutor.java | 9 +- .../batch/BulkExecutorUtil.java | 6 +- .../batch/HybridRowBatchMapper.java | 92 ++++++++++++++ .../PartitionKeyRangeServerBatchRequest.java | 15 ++- .../batch/ServerBatchRequest.java | 56 ++++++--- .../SinglePartitionKeyServerBatchRequest.java | 11 +- .../JsonNodeStorePayload.java | 59 +++++++-- .../directconnectivity/StoreResponse.java | 4 + .../rntbd/RntbdConstants.java | 17 +++ .../rntbd/RntbdRequestHeaders.java | 22 ++++ 24 files changed, 910 insertions(+), 58 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/BinaryEncodingHelperTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdBinaryEncodingHeaderTest.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/BinaryEncodingHelper.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/HybridRowBatchMapper.java 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/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/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/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)) { From 4b4b24a791576c7fac00ce44a955d50f9800bc8c Mon Sep 17 00:00:00 2001 From: Julien Ellie Date: Wed, 15 Jul 2026 15:09:54 -0700 Subject: [PATCH 4/4] Document Cosmos binary encoding Describe the hidden Direct-mode opt-in, supported operation set, and native binary payload benefit. --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) 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)