Skip to content

Commit b516da9

Browse files
alextwoodsalextwoods
committed
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.
1 parent a966125 commit b516da9

7 files changed

Lines changed: 803 additions & 9 deletions

File tree

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

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.protocols.json;
17+
18+
import java.io.ByteArrayInputStream;
19+
import java.io.ByteArrayOutputStream;
20+
import java.io.IOException;
21+
import java.io.InputStream;
22+
import java.io.OutputStream;
23+
import java.io.SequenceInputStream;
24+
import java.util.ArrayList;
25+
import java.util.Collections;
26+
import java.util.List;
27+
import software.amazon.awssdk.annotations.NotThreadSafe;
28+
import software.amazon.awssdk.annotations.SdkInternalApi;
29+
import software.amazon.awssdk.http.ContentStreamProvider;
30+
31+
/**
32+
* A {@link ByteArrayOutputStream} subclass that behaves identically to the JDK implementation for
33+
* small payloads, but caps internal buffer growth to avoid large object allocations for
34+
* large payloads.
35+
*
36+
* <p>
37+
* Writes flow into the inherited {@code ByteArrayOutputStream} buffer
38+
* normally. When a write would cause the buffer to grow beyond {@link #MAX_BUFFER_SIZE}, the
39+
* current buffer contents are frozen into the first "chunk" and subsequent writes go into
40+
* fixed-size overflow chunks ({@link #CHUNK_SIZE} bytes each). No single allocation ever exceeds
41+
* {@code MAX_BUFFER_SIZE}.
42+
*
43+
*/
44+
@NotThreadSafe
45+
@SdkInternalApi
46+
final class SdkByteArrayOutputStream extends ByteArrayOutputStream {
47+
// 128 KB, choosen to be well below 1 MB "humongous threshold" for most heap sizes
48+
static final int MAX_BUFFER_SIZE = 128 * 1024;
49+
static final int CHUNK_SIZE = 64 * 1024;
50+
51+
private List<byte[]> overflowChunks;
52+
private int overflowChunkOffset;
53+
private int overflowTotalBytes;
54+
private boolean overflowing;
55+
56+
SdkByteArrayOutputStream(int initialCapacity) {
57+
super(initialCapacity);
58+
}
59+
60+
@Override
61+
public void write(int b) {
62+
if (overflowing) {
63+
ensureOverflowCapacity(1);
64+
currentOverflowChunk()[overflowChunkOffset++] = (byte) b;
65+
overflowTotalBytes++;
66+
} else if (count + 1 > MAX_BUFFER_SIZE) {
67+
startOverflow();
68+
write(b);
69+
} else {
70+
super.write(b);
71+
}
72+
}
73+
74+
@Override
75+
public void write(byte[] b, int off, int len) {
76+
if (overflowing) {
77+
writeToOverflow(b, off, len);
78+
} else if (count + len > MAX_BUFFER_SIZE) {
79+
// Write what fits into the base buffer, then overflow the rest
80+
int fits = MAX_BUFFER_SIZE - count;
81+
if (fits > 0) {
82+
super.write(b, off, fits);
83+
}
84+
startOverflow();
85+
writeToOverflow(b, off + fits, len - fits);
86+
} else {
87+
super.write(b, off, len);
88+
}
89+
}
90+
91+
/**
92+
* Returns the total number of bytes written (base buffer + overflow).
93+
*/
94+
@Override
95+
public int size() {
96+
return count + overflowTotalBytes;
97+
}
98+
99+
/**
100+
* Returns all written data as a single contiguous byte array. Exists for backward
101+
* compatibility via {@link #toByteArray()} but should not be used on the hot path.
102+
*/
103+
@Override
104+
public byte[] toByteArray() {
105+
if (!overflowing) {
106+
return super.toByteArray();
107+
}
108+
int total = size();
109+
byte[] result = new byte[total];
110+
// Copy base buffer
111+
System.arraycopy(buf, 0, result, 0, count);
112+
// Copy overflow chunks
113+
int destOff = count;
114+
for (int i = 0; i < overflowChunks.size(); i++) {
115+
int len = (i < overflowChunks.size() - 1) ? overflowChunks.get(i).length : overflowChunkOffset;
116+
System.arraycopy(overflowChunks.get(i), 0, result, destOff, len);
117+
destOff += len;
118+
}
119+
return result;
120+
}
121+
122+
/**
123+
* Resets this stream so that all currently accumulated output is discarded, including any
124+
* overflow chunks. After calling this method, the stream can be reused as if freshly constructed.
125+
*/
126+
@Override
127+
public void reset() {
128+
super.reset();
129+
overflowing = false;
130+
overflowChunks = null;
131+
overflowChunkOffset = 0;
132+
overflowTotalBytes = 0;
133+
}
134+
135+
/**
136+
* Writes the complete contents of this stream to the specified output stream, including
137+
* any overflow chunks.
138+
*/
139+
@Override
140+
public void writeTo(OutputStream out) throws IOException {
141+
if (!overflowing) {
142+
super.writeTo(out);
143+
return;
144+
}
145+
// Write base buffer
146+
out.write(buf, 0, count);
147+
// Write overflow chunks
148+
for (int i = 0; i < overflowChunks.size(); i++) {
149+
int len = (i < overflowChunks.size() - 1) ? overflowChunks.get(i).length : overflowChunkOffset;
150+
out.write(overflowChunks.get(i), 0, len);
151+
}
152+
}
153+
154+
/**
155+
* Returns a {@link ContentStreamProvider} that streams directly from the internal buffers
156+
* without creating a contiguous copy. For small payloads this wraps the single base buffer;
157+
* for large payloads it chains the base buffer and overflow chunks via
158+
* {@link SequenceInputStream}.
159+
*/
160+
ContentStreamProvider contentStreamProvider() {
161+
if (!overflowing) {
162+
// Small payload: single buffer, wrap directly
163+
byte[] b = buf;
164+
int c = count;
165+
return () -> new ByteArrayInputStream(b, 0, c);
166+
}
167+
168+
// Large payload: chain base buffer + overflow chunks
169+
byte[] baseBuf = buf;
170+
int baseCount = count;
171+
List<byte[]> chunks = overflowChunks;
172+
int lastChunkLen = overflowChunkOffset;
173+
174+
return () -> {
175+
List<InputStream> streams = new ArrayList<>(1 + chunks.size());
176+
streams.add(new ByteArrayInputStream(baseBuf, 0, baseCount));
177+
for (int i = 0; i < chunks.size(); i++) {
178+
int len = (i < chunks.size() - 1) ? chunks.get(i).length : lastChunkLen;
179+
streams.add(new ByteArrayInputStream(chunks.get(i), 0, len));
180+
}
181+
return new SequenceInputStream(Collections.enumeration(streams));
182+
};
183+
}
184+
185+
/**
186+
* Returns the content size without copying.
187+
*/
188+
int contentSize() {
189+
return size();
190+
}
191+
192+
private void startOverflow() {
193+
overflowing = true;
194+
overflowChunks = new ArrayList<>();
195+
overflowChunks.add(new byte[CHUNK_SIZE]);
196+
overflowChunkOffset = 0;
197+
overflowTotalBytes = 0;
198+
}
199+
200+
private void writeToOverflow(byte[] b, int off, int len) {
201+
int remaining = len;
202+
int srcOff = off;
203+
while (remaining > 0) {
204+
ensureOverflowCapacity(1);
205+
int space = currentOverflowChunk().length - overflowChunkOffset;
206+
int toCopy = Math.min(remaining, space);
207+
System.arraycopy(b, srcOff, currentOverflowChunk(), overflowChunkOffset, toCopy);
208+
overflowChunkOffset += toCopy;
209+
overflowTotalBytes += toCopy;
210+
srcOff += toCopy;
211+
remaining -= toCopy;
212+
}
213+
}
214+
215+
private byte[] currentOverflowChunk() {
216+
return overflowChunks.get(overflowChunks.size() - 1);
217+
}
218+
219+
private void ensureOverflowCapacity(int needed) {
220+
if (overflowChunkOffset + needed > currentOverflowChunk().length) {
221+
overflowChunks.add(new byte[CHUNK_SIZE]);
222+
overflowChunkOffset = 0;
223+
}
224+
}
225+
}

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

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,14 @@
1515

1616
package software.amazon.awssdk.protocols.json;
1717

18-
import java.io.ByteArrayOutputStream;
1918
import java.io.IOException;
2019
import java.math.BigDecimal;
2120
import java.math.BigInteger;
2221
import java.nio.ByteBuffer;
2322
import java.time.Instant;
2423
import software.amazon.awssdk.annotations.SdkProtectedApi;
2524
import software.amazon.awssdk.core.exception.SdkClientException;
25+
import software.amazon.awssdk.http.ContentStreamProvider;
2626
import software.amazon.awssdk.thirdparty.jackson.core.JsonFactory;
2727
import software.amazon.awssdk.thirdparty.jackson.core.JsonGenerator;
2828
import software.amazon.awssdk.utils.BinaryUtils;
@@ -39,7 +39,7 @@ public class SdkJsonGenerator implements StructuredJsonGenerator {
3939
* prevent frequent resizings but small enough to avoid wasted allocations for small requests.
4040
*/
4141
private static final int DEFAULT_BUFFER_SIZE = 1024;
42-
private final ByteArrayOutputStream baos = new ByteArrayOutputStream(DEFAULT_BUFFER_SIZE);
42+
private final SdkByteArrayOutputStream baos = new SdkByteArrayOutputStream(DEFAULT_BUFFER_SIZE);
4343
private final JsonGenerator generator;
4444
private final String contentType;
4545

@@ -206,6 +206,16 @@ public StructuredJsonGenerator writeValue(ByteBuffer bytes) {
206206
return this;
207207
}
208208

209+
@Override
210+
public StructuredJsonGenerator writeBinaryValue(byte[] bytes) {
211+
try {
212+
generator.writeBinary(bytes);
213+
} catch (IOException e) {
214+
throw new JsonGenerationException(e);
215+
}
216+
return this;
217+
}
218+
209219
@Override
210220
//TODO: This date formatting is coupled to AWS's format. Should generalize it
211221
public StructuredJsonGenerator writeValue(Instant instant) {
@@ -277,6 +287,24 @@ public byte[] getBytes() {
277287
return baos.toByteArray();
278288
}
279289

290+
/**
291+
* Returns the size of the generated content in bytes without copying.
292+
*/
293+
public int contentSize() {
294+
close();
295+
return baos.contentSize();
296+
}
297+
298+
/**
299+
* Returns a {@link ContentStreamProvider} that streams directly from the internal buffers
300+
* without creating a contiguous copy. For small payloads this wraps the single base buffer;
301+
* for large payloads it chains the base buffer and overflow chunks.
302+
*/
303+
public ContentStreamProvider contentStreamProvider() {
304+
close();
305+
return baos.contentStreamProvider();
306+
}
307+
280308
@Override
281309
public String getContentType() {
282310
return contentType;

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,13 @@
1515

1616
package software.amazon.awssdk.protocols.json;
1717

18+
import java.io.ByteArrayInputStream;
1819
import java.math.BigDecimal;
1920
import java.math.BigInteger;
2021
import java.nio.ByteBuffer;
2122
import java.time.Instant;
2223
import software.amazon.awssdk.annotations.SdkProtectedApi;
24+
import software.amazon.awssdk.http.ContentStreamProvider;
2325

2426
/**
2527
* Interface for generating a JSON
@@ -102,6 +104,11 @@ public StructuredJsonGenerator writeValue(ByteBuffer bytes) {
102104
return this;
103105
}
104106

107+
@Override
108+
public StructuredJsonGenerator writeBinaryValue(byte[] bytes) {
109+
return this;
110+
}
111+
105112
@Override
106113
public StructuredJsonGenerator writeValue(Instant instant) {
107114
return this;
@@ -169,6 +176,15 @@ default StructuredJsonGenerator writeValue(byte val) {
169176

170177
StructuredJsonGenerator writeValue(ByteBuffer bytes);
171178

179+
/**
180+
* Writes binary data directly from a byte array, avoiding the overhead of wrapping in a
181+
* {@link ByteBuffer}. The default implementation wraps the array and delegates to
182+
* {@link #writeValue(ByteBuffer)}.
183+
*/
184+
default StructuredJsonGenerator writeBinaryValue(byte[] bytes) {
185+
return writeValue(ByteBuffer.wrap(bytes));
186+
}
187+
172188
StructuredJsonGenerator writeValue(Instant instant);
173189

174190
StructuredJsonGenerator writeNumber(String number);
@@ -184,4 +200,28 @@ default StructuredJsonGenerator writeValue(byte val) {
184200
*/
185201
@Deprecated
186202
String getContentType();
203+
204+
/**
205+
* Returns the size of the generated content in bytes without copying. The default
206+
* implementation falls back to {@link #getBytes()}.length.
207+
*/
208+
default int contentSize() {
209+
byte[] bytes = getBytes();
210+
return bytes == null ? 0 : bytes.length;
211+
}
212+
213+
/**
214+
* Returns a {@link ContentStreamProvider} that streams the generated content. The default
215+
* implementation wraps the result of {@link #getBytes()} in a {@code ByteArrayInputStream}.
216+
* Implementations may override this to stream directly from internal buffers without copying.
217+
*
218+
* @return a content stream provider, or {@code null} if {@link #getBytes()} returns null
219+
*/
220+
default ContentStreamProvider contentStreamProvider() {
221+
byte[] bytes = getBytes();
222+
if (bytes == null) {
223+
return null;
224+
}
225+
return () -> new ByteArrayInputStream(bytes);
226+
}
187227
}

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
import software.amazon.awssdk.core.traits.RequiredTrait;
4444
import software.amazon.awssdk.core.traits.TimestampFormatTrait;
4545
import software.amazon.awssdk.core.traits.TraitType;
46+
import software.amazon.awssdk.http.ContentStreamProvider;
4647
import software.amazon.awssdk.http.SdkHttpFullRequest;
4748
import software.amazon.awssdk.protocols.core.InstantToString;
4849
import software.amazon.awssdk.protocols.core.OperationInfo;
@@ -290,12 +291,12 @@ private SdkHttpFullRequest finishMarshalling() {
290291
jsonGenerator.writeEndObject();
291292
}
292293

293-
byte[] content = jsonGenerator.getBytes();
294-
295-
if (content != null) {
296-
request.contentStreamProvider(() -> new ByteArrayInputStream(content));
297-
if (content.length > 0) {
298-
request.putHeader(CONTENT_LENGTH, Integer.toString(content.length));
294+
ContentStreamProvider contentProvider = jsonGenerator.contentStreamProvider();
295+
if (contentProvider != null) {
296+
request.contentStreamProvider(contentProvider);
297+
int contentSize = jsonGenerator.contentSize();
298+
if (contentSize > 0) {
299+
request.putHeader(CONTENT_LENGTH, Integer.toString(contentSize));
299300
}
300301
}
301302
}
@@ -389,7 +390,7 @@ private void marshallPayloadField(SdkField<?> field, Object val) {
389390
break;
390391
case SDK_BYTES:
391392
gen.writeFieldName(fieldName);
392-
gen.writeValue(((SdkBytes) val).asByteBuffer());
393+
gen.writeBinaryValue(((SdkBytes) val).asByteArrayUnsafe());
393394
break;
394395
case SDK_POJO:
395396
SimpleTypeJsonMarshaller.SDK_POJO.marshall((SdkPojo) val, marshallerContext,

0 commit comments

Comments
 (0)