Optimize byte buffer/output stream

Author: alextwoods

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/ExposedByteArrayOutputStream.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.protocols.json;
+
+import java.io.ByteArrayOutputStream;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+
+/**
+ * A thin subclass of {@link ByteArrayOutputStream} that exposes the internal buffer and count
+ * without copying. This allows {@link SdkJsonGenerator} to create a {@code ContentStreamProvider}
+ * that wraps the buffer directly via {@code ByteArrayInputStream(buf, 0, count)}, avoiding the
+ * contiguous copy that {@link ByteArrayOutputStream#toByteArray()} performs.
+ *
+ * <p>The write path is identical to {@code ByteArrayOutputStream} — no overhead is added.
+ * Only the final "get the bytes" step is optimized.
+ *
+ * <p>This class is not thread-safe.
+ */
+@SdkInternalApi
+final class ExposedByteArrayOutputStream extends ByteArrayOutputStream {
+
+    ExposedByteArrayOutputStream(int size) {
+        super(size);
+    }
+
+    /**
+     * Returns the internal buffer. The valid data is in {@code buf[0..count-1]}.
+     * The returned array may be larger than {@link #size()}; callers must use
+     * {@link #size()} to determine the valid range.
+     *
+     * <p><b>Warning:</b> The returned array is the live internal buffer. Do not modify it,
+     * and do not write to this stream after capturing the reference — the buffer may be
+     * replaced by a larger one on the next write if growth is needed.
+     */
+    byte[] buf() {
+        return buf;
+    }
+}

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/SdkJsonGenerator.java

@@ -15,14 +15,15 @@
 
 package software.amazon.awssdk.protocols.json;
 
-import java.io.ByteArrayOutputStream;
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
 import java.time.Instant;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.thirdparty.jackson.core.JsonFactory;
 import software.amazon.awssdk.thirdparty.jackson.core.JsonGenerator;
 import software.amazon.awssdk.utils.BinaryUtils;
@@ -39,7 +40,7 @@ public class SdkJsonGenerator implements StructuredJsonGenerator {
      * prevent frequent resizings but small enough to avoid wasted allocations for small requests.
      */
     private static final int DEFAULT_BUFFER_SIZE = 1024;
-    private final ByteArrayOutputStream baos = new ByteArrayOutputStream(DEFAULT_BUFFER_SIZE);
+    private final ExposedByteArrayOutputStream baos = new ExposedByteArrayOutputStream(DEFAULT_BUFFER_SIZE);
     private final JsonGenerator generator;
     private final String contentType;
 
@@ -277,6 +278,28 @@ public byte[] getBytes() {
         return baos.toByteArray();
     }
 
+    /**
+     * Returns the size of the generated content in bytes without copying.
+     */
+    public int contentSize() {
+        close();
+        return baos.size();
+    }
+
+    /**
+     * Returns a {@link ContentStreamProvider} that wraps the internal buffer directly,
+     * avoiding the contiguous copy that {@link #getBytes()} performs via
+     * {@code ByteArrayOutputStream.toByteArray()}. Each call to
+     * {@link ContentStreamProvider#newStream()} creates a fresh {@code ByteArrayInputStream}
+     * over the same buffer for retry safety.
+     */
+    public ContentStreamProvider contentStreamProvider() {
+        close();
+        byte[] buf = baos.buf();
+        int count = baos.size();
+        return () -> new ByteArrayInputStream(buf, 0, count);
+    }
+
     @Override
     public String getContentType() {
         return contentType;

core/protocols/aws-json-protocol/src/main/java/software/amazon/awssdk/protocols/json/internal/marshall/JsonProtocolMarshaller.java

@@ -43,6 +43,7 @@
 import software.amazon.awssdk.core.traits.RequiredTrait;
 import software.amazon.awssdk.core.traits.TimestampFormatTrait;
 import software.amazon.awssdk.core.traits.TraitType;
+import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.protocols.core.InstantToString;
 import software.amazon.awssdk.protocols.core.OperationInfo;
@@ -52,6 +53,7 @@
 import software.amazon.awssdk.protocols.json.AwsJsonProtocol;
 import software.amazon.awssdk.protocols.json.AwsJsonProtocolMetadata;
 import software.amazon.awssdk.protocols.json.BaseAwsJsonProtocolFactory;
+import software.amazon.awssdk.protocols.json.SdkJsonGenerator;
 import software.amazon.awssdk.protocols.json.StructuredJsonGenerator;
 import software.amazon.awssdk.protocols.json.internal.ProtocolFact;
 
@@ -288,12 +290,24 @@ private SdkHttpFullRequest finishMarshalling() {
                 jsonGenerator.writeEndObject();
             }
 
-            byte[] content = jsonGenerator.getBytes();
+            if (jsonGenerator instanceof SdkJsonGenerator) {
+                // Optimized path: stream directly from chunked buffers, avoiding a single
+                // contiguous byte[] allocation that can cause G1GC humongous allocations.
+                SdkJsonGenerator sdkGenerator = (SdkJsonGenerator) jsonGenerator;
+                ContentStreamProvider contentProvider = sdkGenerator.contentStreamProvider();
+                request.contentStreamProvider(contentProvider);
+                int contentSize = sdkGenerator.contentSize();
+                if (contentSize > 0) {
+                    request.putHeader(CONTENT_LENGTH, Integer.toString(contentSize));
+                }
+            } else {
+                byte[] content = jsonGenerator.getBytes();
 
-            if (content != null) {
-                request.contentStreamProvider(() -> new ByteArrayInputStream(content));
-                if (content.length > 0) {
-                    request.putHeader(CONTENT_LENGTH, Integer.toString(content.length));
+                if (content != null) {
+                    request.contentStreamProvider(() -> new ByteArrayInputStream(content));
+                    if (content.length > 0) {
+                        request.putHeader(CONTENT_LENGTH, Integer.toString(content.length));
+                    }
                 }
             }
         }

core/protocols/aws-json-protocol/src/test/java/software/amazon/awssdk/protocols/json/ExposedByteArrayOutputStreamTest.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.protocols.json;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.Arrays;
+import org.junit.jupiter.api.Test;
+
+class ExposedByteArrayOutputStreamTest {
+
+    @Test
+    void emptyStream_hasZeroSize() {
+        ExposedByteArrayOutputStream stream = new ExposedByteArrayOutputStream(64);
+        assertThat(stream.size()).isEqualTo(0);
+        assertThat(stream.toByteArray()).isEmpty();
+    }
+
+    @Test
+    void buf_returnsInternalBuffer() {
+        ExposedByteArrayOutputStream stream = new ExposedByteArrayOutputStream(64);
+        byte[] data = {1, 2, 3, 4, 5};
+        stream.write(data, 0, data.length);
+
+        byte[] buf = stream.buf();
+        // buf is the live internal buffer — it may be larger than size()
+        assertThat(buf.length).isGreaterThanOrEqualTo(stream.size());
+        // The valid data in buf[0..size()-1] matches what was written
+        assertThat(Arrays.copyOf(buf, stream.size())).isEqualTo(data);
+    }
+
+    @Test
+    void buf_reflectsWrittenData_afterGrowth() {
+        // Start with a tiny buffer to force growth
+        ExposedByteArrayOutputStream stream = new ExposedByteArrayOutputStream(4);
+        byte[] data = new byte[100];
+        for (int i = 0; i < data.length; i++) {
+            data[i] = (byte) i;
+        }
+        stream.write(data, 0, data.length);
+
+        byte[] buf = stream.buf();
+        assertThat(stream.size()).isEqualTo(100);
+        assertThat(Arrays.copyOf(buf, stream.size())).isEqualTo(data);
+    }
+
+    @Test
+    void toByteArray_returnsCopy_notSameReference() {
+        ExposedByteArrayOutputStream stream = new ExposedByteArrayOutputStream(64);
+        stream.write(new byte[]{1, 2, 3}, 0, 3);
+
+        byte[] copy = stream.toByteArray();
+        byte[] buf = stream.buf();
+        // toByteArray returns a copy, buf returns the live buffer
+        assertThat(copy).isNotSameAs(buf);
+        assertThat(copy).isEqualTo(Arrays.copyOf(buf, stream.size()));
+    }
+
+    @Test
+    void singleByteWrite_worksCorrectly() {
+        ExposedByteArrayOutputStream stream = new ExposedByteArrayOutputStream(64);
+        stream.write(0x42);
+        assertThat(stream.size()).isEqualTo(1);
+        assertThat(stream.buf()[0]).isEqualTo((byte) 0x42);
+    }
+}

core/protocols/aws-json-protocol/src/test/java/software/amazon/awssdk/protocols/json/SdkJsonGeneratorTest.java

@@ -21,11 +21,13 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.time.Instant;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.protocols.jsoncore.JsonNode;
 import software.amazon.awssdk.thirdparty.jackson.core.JsonFactory;
 import software.amazon.awssdk.thirdparty.jackson.core.StreamReadFeature;
@@ -178,4 +180,131 @@ private JsonNode toJsonNode() throws IOException {
         return JsonNode.parser().parse(new ByteArrayInputStream(jsonGenerator.getBytes()));
     }
 
+    @Test
+    public void contentSize_matchesGetBytesLength() {
+        SdkJsonGenerator gen = newSdkJsonGenerator();
+        gen.writeStartObject();
+        gen.writeFieldName("key").writeValue("value");
+        gen.writeFieldName("num").writeValue(42);
+        gen.writeEndObject();
+
+        byte[] bytes = gen.getBytes();
+
+        SdkJsonGenerator gen2 = newSdkJsonGenerator();
+        gen2.writeStartObject();
+        gen2.writeFieldName("key").writeValue("value");
+        gen2.writeFieldName("num").writeValue(42);
+        gen2.writeEndObject();
+
+        assertEquals(bytes.length, gen2.contentSize());
+    }
+
+    @Test
+    public void contentStreamProvider_producesSameBytesAsGetBytes() throws IOException {
+        SdkJsonGenerator gen = newSdkJsonGenerator();
+        gen.writeStartObject();
+        gen.writeFieldName("hello").writeValue("world");
+        gen.writeFieldName("count").writeValue(123);
+        gen.writeEndObject();
+
+        byte[] expected = gen.getBytes();
+
+        SdkJsonGenerator gen2 = newSdkJsonGenerator();
+        gen2.writeStartObject();
+        gen2.writeFieldName("hello").writeValue("world");
+        gen2.writeFieldName("count").writeValue(123);
+        gen2.writeEndObject();
+
+        ContentStreamProvider provider = gen2.contentStreamProvider();
+        byte[] actual = readAllBytes(provider.newStream());
+
+        assertTrue(java.util.Arrays.equals(expected, actual),
+            "contentStreamProvider should produce identical bytes to getBytes");
+    }
+
+    @Test
+    public void contentStreamProvider_isResettable() throws IOException {
+        SdkJsonGenerator gen = newSdkJsonGenerator();
+        gen.writeStartObject();
+        gen.writeFieldName("data").writeValue("test");
+        gen.writeEndObject();
+
+        ContentStreamProvider provider = gen.contentStreamProvider();
+        byte[] first = readAllBytes(provider.newStream());
+        byte[] second = readAllBytes(provider.newStream());
+
+        assertTrue(java.util.Arrays.equals(first, second),
+            "Multiple calls to newStream() should produce identical content");
+        assertTrue(first.length > 0, "Content should not be empty");
+    }
+
+    @Test
+    public void emptyGenerator_contentSizeIsZero() throws IOException {
+        SdkJsonGenerator gen = newSdkJsonGenerator();
+        assertEquals(0, gen.contentSize());
+
+        ContentStreamProvider provider = gen.contentStreamProvider();
+        assertTrue(provider != null, "Provider should not be null even for empty content");
+        byte[] content = readAllBytes(provider.newStream());
+        assertEquals(0, content.length, "Empty generator should produce empty stream");
+    }
+
+    @Test
+    public void largePayload_contentStreamProviderStreamsCorrectData() throws IOException {
+        // Generate JSON exceeding 64 KB to verify contentStreamProvider works for large payloads
+        SdkJsonGenerator gen = newSdkJsonGenerator();
+        gen.writeStartObject();
+        gen.writeFieldName("items");
+        gen.writeStartArray();
+        for (int i = 0; i < 2000; i++) {
+            gen.writeStartObject();
+            gen.writeFieldName("index").writeValue(i);
+            gen.writeFieldName("description").writeValue(
+                "This is a moderately long string value for item number " + i +
+                " that helps push the total payload size beyond the 64KB chunk boundary.");
+            gen.writeEndObject();
+        }
+        gen.writeEndArray();
+        gen.writeEndObject();
+
+        byte[] expected = gen.getBytes();
+        assertTrue(expected.length > 64 * 1024, "Payload should exceed 64 KB");
+
+        SdkJsonGenerator gen2 = newSdkJsonGenerator();
+        gen2.writeStartObject();
+        gen2.writeFieldName("items");
+        gen2.writeStartArray();
+        for (int i = 0; i < 2000; i++) {
+            gen2.writeStartObject();
+            gen2.writeFieldName("index").writeValue(i);
+            gen2.writeFieldName("description").writeValue(
+                "This is a moderately long string value for item number " + i +
+                " that helps push the total payload size beyond the 64KB chunk boundary.");
+            gen2.writeEndObject();
+        }
+        gen2.writeEndArray();
+        gen2.writeEndObject();
+
+        assertEquals(expected.length, gen2.contentSize());
+        byte[] actual = readAllBytes(gen2.contentStreamProvider().newStream());
+        assertTrue(java.util.Arrays.equals(expected, actual),
+            "Large payload should stream correctly via contentStreamProvider");
+    }
+
+    private SdkJsonGenerator newSdkJsonGenerator() {
+        return new SdkJsonGenerator(JsonFactory.builder()
+                                               .enable(StreamReadFeature.INCLUDE_SOURCE_IN_LOCATION)
+                                               .build(), "application/json");
+    }
+
+    private static byte[] readAllBytes(InputStream is) throws IOException {
+        java.io.ByteArrayOutputStream bos = new java.io.ByteArrayOutputStream();
+        byte[] buf = new byte[1024];
+        int n;
+        while ((n = is.read(buf)) != -1) {
+            bos.write(buf, 0, n);
+        }
+        return bos.toByteArray();
+    }
+
 }