Optimize GC usage in JSON marshalling (#6935)

* Add SdkByteArrayOutputStream, zero-copy streaming, and writeBinaryValue optimizations

Optimization 3: SdkByteArrayOutputStream - a hybrid ByteArrayOutputStream
that caps buffer growth at 128 KB and overflows into fixed 64 KB chunks,
eliminating G1GC humongous object allocations for large payloads.

Optimization 4: Zero-copy content streaming via contentStreamProvider()
and contentSize() on StructuredJsonGenerator/SdkJsonGenerator, avoiding
the toByteArray() copy in finishMarshalling(). SdkJsonGenerator now uses
SdkByteArrayOutputStream as its internal buffer.

Also adds writeBinaryValue(byte[]) to StructuredJsonGenerator and
SdkJsonGenerator, and updates the SDK_BYTES switch case to use
asByteArrayUnsafe() to bypass the ByteBuffer allocation and
copyBytesFrom copy.

Includes the large SDK_BYTES equivalence test (200 KB payload) that
exercises the SdkByteArrayOutputStream overflow path.

These optimizations build on the megamorphic dispatch and registry
caching changes in alexwoo/json_marshall_opt.

* Minor cleanups

* Add partial read unit tests

Author: alextwoods

.changes/next-release/bugfix-AWSSDKforJavav2-1db72bd.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Optimized GC usage (specifically G1GC humongous allocations) in JSON marshalling"
+}

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

@@ -0,0 +1,226 @@
+/*
+ * 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.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.SequenceInputStream;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import software.amazon.awssdk.annotations.NotThreadSafe;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.http.ContentStreamProvider;
+
+/**
+ * A {@link ByteArrayOutputStream} subclass that behaves identically to the JDK implementation for
+ * small payloads, but caps internal buffer growth to avoid large object allocations for
+ * large payloads.
+ *
+ * <p>
+ * Writes flow into the inherited {@code ByteArrayOutputStream} buffer
+ * normally. When a write would cause the buffer to grow beyond {@link #MAX_BUFFER_SIZE}, the
+ * current buffer contents are frozen and subsequent writes go into fixed-size chunks
+ * ({@link #CHUNK_SIZE} bytes each). No single allocation ever exceeds {@code MAX_BUFFER_SIZE}.
+ *
+ * <p>
+ * The stream is in "chunked mode" when {@code chunks != null}. This is derived from state
+ * rather than tracked by a separate boolean, eliminating a class of state-sync bugs.
+ *
+ */
+@NotThreadSafe
+@SdkInternalApi
+final class SdkByteArrayOutputStream extends ByteArrayOutputStream {
+    // 128 KB, chosen to be well below 1 MB "humongous threshold" for most heap sizes
+    static final int MAX_BUFFER_SIZE = 128 * 1024;
+    static final int CHUNK_SIZE = 64 * 1024;
+
+    private List<byte[]> chunks;
+    private int chunkOffset;
+    private int chunkedBytes;
+
+    SdkByteArrayOutputStream(int initialCapacity) {
+        super(initialCapacity);
+    }
+
+    @Override
+    public void write(int b) {
+        if (chunks != null) {
+            ensureChunkCapacity(1);
+            currentChunk()[chunkOffset++] = (byte) b;
+            chunkedBytes++;
+        } else if (count + 1 > MAX_BUFFER_SIZE) {
+            startChunking();
+            write(b);
+        } else {
+            super.write(b);
+        }
+    }
+
+    @Override
+    public void write(byte[] b, int off, int len) {
+        if (chunks != null) {
+            writeToChunks(b, off, len);
+        } else if (count + len > MAX_BUFFER_SIZE) {
+            // Write what fits into the base buffer, then chunk the rest
+            int fits = MAX_BUFFER_SIZE - count;
+            if (fits > 0) {
+                super.write(b, off, fits);
+            }
+            startChunking();
+            writeToChunks(b, off + fits, len - fits);
+        } else {
+            super.write(b, off, len);
+        }
+    }
+
+    /**
+     * Returns the total number of bytes written (base buffer + chunks).
+     */
+    @Override
+    public int size() {
+        return count + chunkedBytes;
+    }
+
+    /**
+     * Returns all written data as a single contiguous byte array. Exists for backward
+     * compatibility via {@link #toByteArray()} but should not be used on the hot path.
+     */
+    @Override
+    public byte[] toByteArray() {
+        if (chunks == null) {
+            return super.toByteArray();
+        }
+        int total = size();
+        byte[] result = new byte[total];
+        // Copy base buffer
+        System.arraycopy(buf, 0, result, 0, count);
+        // Copy chunks
+        int destOff = count;
+        for (int i = 0; i < chunks.size(); i++) {
+            int len = (i < chunks.size() - 1) ? chunks.get(i).length : chunkOffset;
+            System.arraycopy(chunks.get(i), 0, result, destOff, len);
+            destOff += len;
+        }
+        return result;
+    }
+
+    /**
+     * Resets this stream so that all currently accumulated output is discarded, including any
+     * chunks. After calling this method, the stream can be reused as if freshly constructed.
+     */
+    @Override
+    public void reset() {
+        super.reset();
+        chunks = null;
+        chunkOffset = 0;
+        chunkedBytes = 0;
+    }
+
+    /**
+     * Writes the complete contents of this stream to the specified output stream, including
+     * any chunks.
+     */
+    @Override
+    public void writeTo(OutputStream out) throws IOException {
+        if (chunks == null) {
+            super.writeTo(out);
+            return;
+        }
+        // Write base buffer
+        out.write(buf, 0, count);
+        // Write chunks
+        for (int i = 0; i < chunks.size(); i++) {
+            int len = (i < chunks.size() - 1) ? chunks.get(i).length : chunkOffset;
+            out.write(chunks.get(i), 0, len);
+        }
+    }
+
+    /**
+     * Returns a {@link ContentStreamProvider} that streams directly from the internal buffers
+     * without creating a contiguous copy. For small payloads this wraps the single base buffer;
+     * for large payloads it chains the base buffer and chunks via {@link SequenceInputStream}.
+     */
+    ContentStreamProvider contentStreamProvider() {
+        if (chunks == null) {
+            // Small payload: single buffer, wrap directly.
+            // Safe to capture buf because callers (SdkJsonGenerator.contentStreamProvider()) close the
+            // generator before calling this, guaranteeing no further writes will mutate buf.
+            byte[] b = buf;
+            int c = count;
+            return () -> new ByteArrayInputStream(b, 0, c);
+        }
+
+        // Large payload: chain base buffer + chunks
+        byte[] baseBuf = buf;
+        int baseCount = count;
+        List<byte[]> capturedChunks = chunks;
+        int lastChunkLen = chunkOffset;
+
+        return () -> {
+            List<InputStream> streams = new ArrayList<>(1 + capturedChunks.size());
+            streams.add(new ByteArrayInputStream(baseBuf, 0, baseCount));
+            for (int i = 0; i < capturedChunks.size(); i++) {
+                int len = (i < capturedChunks.size() - 1) ? capturedChunks.get(i).length : lastChunkLen;
+                streams.add(new ByteArrayInputStream(capturedChunks.get(i), 0, len));
+            }
+            return new SequenceInputStream(Collections.enumeration(streams));
+        };
+    }
+
+    /**
+     * Returns the content size without copying.
+     */
+    int contentSize() {
+        return size();
+    }
+
+    private void startChunking() {
+        chunks = new ArrayList<>();
+        chunks.add(new byte[CHUNK_SIZE]);
+        chunkOffset = 0;
+        chunkedBytes = 0;
+    }
+
+    private void writeToChunks(byte[] b, int off, int len) {
+        int remaining = len;
+        int srcOff = off;
+        while (remaining > 0) {
+            ensureChunkCapacity(1);
+            int space = currentChunk().length - chunkOffset;
+            int toCopy = Math.min(remaining, space);
+            System.arraycopy(b, srcOff, currentChunk(), chunkOffset, toCopy);
+            chunkOffset += toCopy;
+            chunkedBytes += toCopy;
+            srcOff += toCopy;
+            remaining -= toCopy;
+        }
+    }
+
+    private byte[] currentChunk() {
+        return chunks.get(chunks.size() - 1);
+    }
+
+    private void ensureChunkCapacity(int needed) {
+        if (chunkOffset + needed > currentChunk().length) {
+            chunks.add(new byte[CHUNK_SIZE]);
+            chunkOffset = 0;
+        }
+    }
+}

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

@@ -15,14 +15,14 @@
 
 package software.amazon.awssdk.protocols.json;
 
-import java.io.ByteArrayOutputStream;
 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 +39,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 SdkByteArrayOutputStream baos = new SdkByteArrayOutputStream(DEFAULT_BUFFER_SIZE);
     private final JsonGenerator generator;
     private final String contentType;
 
@@ -206,6 +206,16 @@ public StructuredJsonGenerator writeValue(ByteBuffer bytes) {
         return this;
     }
 
+    @Override
+    public StructuredJsonGenerator writeBinaryValue(byte[] bytes) {
+        try {
+            generator.writeBinary(bytes);
+        } catch (IOException e) {
+            throw new JsonGenerationException(e);
+        }
+        return this;
+    }
+
     @Override
     //TODO: This date formatting is coupled to AWS's format. Should generalize it
     public StructuredJsonGenerator writeValue(Instant instant) {
@@ -277,6 +287,24 @@ public byte[] getBytes() {
         return baos.toByteArray();
     }
 
+    /**
+     * Returns the size of the generated content in bytes without copying.
+     */
+    public int contentSize() {
+        close();
+        return baos.contentSize();
+    }
+
+    /**
+     * Returns a {@link ContentStreamProvider} that streams directly from the internal buffers
+     * without creating a contiguous copy. For small payloads this wraps the single base buffer;
+     * for large payloads it chains the base buffer and overflow chunks.
+     */
+    public ContentStreamProvider contentStreamProvider() {
+        close();
+        return baos.contentStreamProvider();
+    }
+
     @Override
     public String getContentType() {
         return contentType;

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

@@ -15,11 +15,13 @@
 
 package software.amazon.awssdk.protocols.json;
 
+import java.io.ByteArrayInputStream;
 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.http.ContentStreamProvider;
 
 /**
  * Interface for generating a JSON
@@ -102,6 +104,11 @@ public StructuredJsonGenerator writeValue(ByteBuffer bytes) {
             return this;
         }
 
+        @Override
+        public StructuredJsonGenerator writeBinaryValue(byte[] bytes) {
+            return this;
+        }
+
         @Override
         public StructuredJsonGenerator writeValue(Instant instant) {
             return this;
@@ -169,6 +176,15 @@ default StructuredJsonGenerator writeValue(byte val) {
 
     StructuredJsonGenerator writeValue(ByteBuffer bytes);
 
+    /**
+     * Writes binary data directly from a byte array, avoiding the overhead of wrapping in a
+     * {@link ByteBuffer}. The default implementation wraps the array and delegates to
+     * {@link #writeValue(ByteBuffer)}.
+     */
+    default StructuredJsonGenerator writeBinaryValue(byte[] bytes) {
+        return writeValue(ByteBuffer.wrap(bytes));
+    }
+
     StructuredJsonGenerator writeValue(Instant instant);
 
     StructuredJsonGenerator writeNumber(String number);
@@ -184,4 +200,28 @@ default StructuredJsonGenerator writeValue(byte val) {
      */
     @Deprecated
     String getContentType();
+
+    /**
+     * Returns the size of the generated content in bytes without copying. The default
+     * implementation falls back to {@link #getBytes()}.length.
+     */
+    default int contentSize() {
+        byte[] bytes = getBytes();
+        return bytes == null ? 0 : bytes.length;
+    }
+
+    /**
+     * Returns a {@link ContentStreamProvider} that streams the generated content. The default
+     * implementation wraps the result of {@link #getBytes()} in a {@code ByteArrayInputStream}.
+     * Implementations may override this to stream directly from internal buffers without copying.
+     *
+     * @return a content stream provider, or {@code null} if {@link #getBytes()} returns null
+     */
+    default ContentStreamProvider contentStreamProvider() {
+        byte[] bytes = getBytes();
+        if (bytes == null) {
+            return null;
+        }
+        return () -> new ByteArrayInputStream(bytes);
+    }
 }