From 7030ce8dc4e85689a02348d43c3846ff3924890d Mon Sep 17 00:00:00 2001 From: Adwait Kumar Singh Date: Sun, 11 Oct 2026 09:27:55 +0530 Subject: [PATCH 1/2] Add output buffer pooling --- .../scheme/sigv4/AwsChunkedDataStream.java | 14 +- .../client/auth/scheme/sigv4/SigV4Signer.java | 2 +- .../auth/scheme/sigv4/SigningResources.java | 27 +- .../sigv4/AwsChunkedDataStreamTest.java | 19 + .../sigv4/SigningResourcesPropertyTest.java | 10 + .../aws/client/awsjson/AwsJsonProtocol.java | 4 +- .../QueryFormSerializerBenchmark.java | 4 +- .../awsquery/AwsQueryClientProtocol.java | 16 +- .../awsquery/Ec2QueryClientProtocol.java | 16 +- .../client/awsquery/QueryFormSerializer.java | 62 ++-- .../awsquery/QueryFormSerializerTest.java | 24 +- .../restjson/RestJsonClientProtocol.java | 2 + .../client/restxml/RestXmlClientProtocol.java | 2 + .../java/benchmarks/e2e/BenchmarkClient.java | 2 +- .../smithy/java/benchmarks/e2e/Bodies.java | 59 +++- .../java/benchmarks/e2e/BodiesTest.java | 33 ++ benchmarks/serde-benchmarks/README.md | 5 + .../AwsJson1_0CodecSerializeBenchmark.java | 69 ++++ .../serde/AwsJson1_0SerializeBenchmark.java | 6 +- .../serde/AwsQuerySerializeBenchmark.java | 6 +- .../serde/RestJson1SerializeBenchmark.java | 6 +- .../serde/RestXmlSerializeBenchmark.java | 6 +- .../RpcV2CborCodecSerializeBenchmark.java | 58 +++ .../serde/RpcV2CborSerializeBenchmark.java | 6 +- .../java/benchmarks/serde/SerializeState.java | 13 + .../java/client/core/ClientPipeline.java | 32 +- .../java/client/core/ClientProtocol.java | 14 + .../java/client/core/ClientPipelineTest.java | 186 ++++++++++ .../java/client/http/HttpClientProtocol.java | 20 ++ .../java/client/http/compression/Gzip.java | 9 +- .../http/plugins/HttpChecksumPlugin.java | 8 +- .../java/client/http/mock/MockPlugin.java | 9 +- .../java/client/http/mock/MockPluginTest.java | 85 +++++ .../java/client/rpcv2/RpcV2CborProtocol.java | 5 +- .../client/rpcv2json/RpcV2JsonProtocol.java | 2 + .../rpcv2/AbstractRpcV2ClientProtocol.java | 2 +- codecs/cbor-codec/build.gradle.kts | 2 +- .../java/cbor/CborStringWriteBenchmark.java | 3 +- .../smithy/java/cbor/CborSerdeProvider.java | 5 + .../smithy/java/cbor/CborSerializer.java | 65 ++-- .../amazon/smithy/java/cbor/CborSettings.java | 22 ++ .../java/cbor/DefaultCborSerdeProvider.java | 20 +- .../smithy/java/cbor/Rpcv2CborCodec.java | 20 ++ .../smithy/java/cbor/CborCodecTest.java | 77 +++- .../smithy/java/cbor/CborSerializerTest.java | 3 +- .../java/codecs/commons/OutputBuffers.java | 159 +++++++++ .../java/codecs/commons/StripedPool.java | 36 +- .../codecs/commons/OutputBuffersTest.java | 159 +++++++++ .../java/codecs/commons/StripedPoolTest.java | 104 ++++++ .../java/json/CodegenJsonSerdeProvider.java | 25 ++ .../amazon/smithy/java/json/JsonCodec.java | 27 ++ .../smithy/java/json/JsonSerdeProvider.java | 10 + .../amazon/smithy/java/json/JsonSettings.java | 27 ++ .../java/json/smithy/JsonCodegenWriter.java | 42 ++- .../json/smithy/SmithyGeneratedJsonSerde.java | 34 ++ .../json/smithy/SmithyJsonSerdeProvider.java | 19 + .../json/smithy/SmithyJsonSerializer.java | 60 +++- .../smithy/java/json/JsonSerializerTest.java | 53 +++ .../smithy/java/xml/LazyXmlSerializer.java | 23 +- .../smithy/java/xml/SmithyXmlSerializer.java | 57 ++- .../amazon/smithy/java/xml/XmlCodec.java | 40 ++- .../amazon/smithy/java/xml/XmlCodecTest.java | 60 ++++ .../amazon/smithy/java/core/serde/Codec.java | 16 + .../java/core/serde/MemberSubsetCodec.java | 11 + .../http/binding/HttpBindingSerializer.java | 35 +- .../http/binding/HttpBindingWriterTest.java | 32 ++ .../java/http/client/DefaultHttpClient.java | 13 +- .../http/client/h2/H2StreamRequestBody.java | 6 +- .../http/client/DefaultHttpClientTest.java | 112 ++++++ .../java/io/datastream/ByteArrayRecycler.java | 21 ++ .../smithy/java/io/datastream/DataStream.java | 52 ++- .../java/io/datastream/LeasedDataStream.java | 174 +++++++++ .../datastream/MultiByteBufferDataStream.java | 8 + .../java/io/datastream/PooledDataStream.java | 66 ++++ .../java/io/datastream/WrappedDataStream.java | 8 +- .../datastream/WrappedPooledDataStream.java | 50 +++ .../io/datastream/LeasedDataStreamTest.java | 331 ++++++++++++++++++ ...StreamClientTestsProtocolTestProvider.java | 3 +- ...HttpClientRequestProtocolTestProvider.java | 3 +- 79 files changed, 2718 insertions(+), 218 deletions(-) create mode 100644 benchmarks/e2e-benchmarks/src/test/java/software/amazon/smithy/java/benchmarks/e2e/BodiesTest.java create mode 100644 benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsJson1_0CodecSerializeBenchmark.java create mode 100644 benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RpcV2CborCodecSerializeBenchmark.java create mode 100644 codecs/codec-commons/src/main/java/software/amazon/smithy/java/codecs/commons/OutputBuffers.java create mode 100644 codecs/codec-commons/src/test/java/software/amazon/smithy/java/codecs/commons/OutputBuffersTest.java create mode 100644 codecs/codec-commons/src/test/java/software/amazon/smithy/java/codecs/commons/StripedPoolTest.java create mode 100644 io/src/main/java/software/amazon/smithy/java/io/datastream/ByteArrayRecycler.java create mode 100644 io/src/main/java/software/amazon/smithy/java/io/datastream/LeasedDataStream.java create mode 100644 io/src/main/java/software/amazon/smithy/java/io/datastream/PooledDataStream.java create mode 100644 io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedPooledDataStream.java create mode 100644 io/src/test/java/software/amazon/smithy/java/io/datastream/LeasedDataStreamTest.java diff --git a/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStream.java b/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStream.java index 4f1dc6d297..4a5b223dd4 100644 --- a/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStream.java +++ b/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStream.java @@ -111,19 +111,11 @@ public InputStream asInputStream() { @Override public void writeTo(OutputStream out) throws IOException { - var buf = inner.asByteBuffer(); - long size = buf.remaining(); var crc = new CRC32(); - crc.update(buf.duplicate()); - out.write(hexBytes(size)); + inner.forEachBuffer(crc::update); + out.write(hexBytes(decodedLength)); out.write(CRLF); - if (buf.hasArray()) { - out.write(buf.array(), buf.arrayOffset() + buf.position(), buf.remaining()); - } else { - byte[] tmp = new byte[buf.remaining()]; - buf.duplicate().get(tmp); - out.write(tmp); - } + inner.writeTo(out); out.write(CRLF); out.write(TERMINATOR); out.write(TRAILER_NAME); diff --git a/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Signer.java b/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Signer.java index 6fc21fa429..9d3d900aee 100644 --- a/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Signer.java +++ b/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Signer.java @@ -110,7 +110,7 @@ private void computePayloadHash(DataStream dataStream) { "Cannot SigV4-sign a body whose length is unknown without UNSIGNED-PAYLOAD " + "or chunked signing, neither of which is implemented yet."); } - r.hash(dataStream.asByteBuffer(), r.hashBytes); + r.hash(dataStream, r.hashBytes); SigningResources.writeHex(r.hashBytes, r.payloadHashHex, 0); } diff --git a/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResources.java b/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResources.java index 125706c16f..fc8541f56e 100644 --- a/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResources.java +++ b/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResources.java @@ -21,6 +21,7 @@ import javax.crypto.spec.SecretKeySpec; import software.amazon.smithy.java.codecs.commons.CompactStringAccess; import software.amazon.smithy.java.codecs.commons.StripedPool; +import software.amazon.smithy.java.io.datastream.DataStream; /** * Canonical requests use ASCII except for header values, which use UTF-8. @@ -269,7 +270,19 @@ byte[] hash(byte[] data, int offset, int length, byte[] output) { /** Hashes the remaining bytes without changing the buffer position. */ byte[] hash(ByteBuffer data, byte[] output) { try { - updateFromBuffer(data.duplicate()); + updateFromView(data); + sha256Digest.digest(output, 0, output.length); + return output; + } catch (DigestException | RuntimeException e) { + sha256Digest.reset(); + throw e instanceof RuntimeException re ? re : new RuntimeException(e); + } + } + + /** Hashes each buffer view in turn so a pooled body is never copied into one contiguous buffer. */ + byte[] hash(DataStream body, byte[] output) { + try { + body.forEachBuffer(this::updateFromView); sha256Digest.digest(output, 0, output.length); return output; } catch (DigestException | RuntimeException e) { @@ -280,21 +293,21 @@ byte[] hash(ByteBuffer data, byte[] output) { /** * Copy non-array buffers through scratch to avoid retained JDK buffers that each digest clone would copy. - * This method consumes the buffer. + * Leaves the view's position unchanged. */ - private void updateFromBuffer(ByteBuffer data) { + private void updateFromView(ByteBuffer data) { if (data.hasArray()) { sha256Digest.update(data.array(), data.arrayOffset() + data.position(), data.remaining()); - data.position(data.limit()); return; } byte[] tmp = scratch; if (tmp == null) { tmp = scratch = new byte[SCRATCH_SIZE]; } - while (data.hasRemaining()) { - int n = Math.min(data.remaining(), tmp.length); - data.get(tmp, 0, n); + ByteBuffer copy = data.duplicate(); + while (copy.hasRemaining()) { + int n = Math.min(copy.remaining(), tmp.length); + copy.get(tmp, 0, n); sha256Digest.update(tmp, 0, n); } } diff --git a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStreamTest.java b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStreamTest.java index 1dcc11b982..f453b36f50 100644 --- a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStreamTest.java +++ b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStreamTest.java @@ -24,6 +24,7 @@ import java.util.zip.CRC32; import org.junit.jupiter.api.Test; import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; class AwsChunkedDataStreamTest { @@ -194,4 +195,22 @@ public void onComplete() { // Frames: header + body + footer + complete = 3 onNext + 1 onComplete. assertThat(calls.get(), equalTo(List.of("next", "next", "next", "complete"))); } + + @Test + public void writeToDoesNotAliasAPooledBody() throws Exception { + byte[] payload = "hello chunked world".getBytes(StandardCharsets.UTF_8); + var recycled = new AtomicReference(); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, "text/plain", recycled::set); + + var encoded = new ByteArrayOutputStream(); + new AwsChunkedDataStream(pooled).writeTo(encoded); + + var expected = new ByteArrayOutputStream(); + new AwsChunkedDataStream(DataStream.ofBytes(payload)).writeTo(expected); + assertArrayEquals(expected.toByteArray(), encoded.toByteArray()); + + assertThat(pooled.escaped(), is(false)); + pooled.release(); + assertThat(recycled.get(), is(payload)); + } } diff --git a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResourcesPropertyTest.java b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResourcesPropertyTest.java index fa108fd1a9..f058125b5c 100644 --- a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResourcesPropertyTest.java +++ b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResourcesPropertyTest.java @@ -20,6 +20,7 @@ import java.util.Locale; import java.util.Random; import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.io.datastream.DataStream; class SigningResourcesPropertyTest { @@ -145,6 +146,15 @@ static void assertHash(SigningResources r, byte[] data, int pad, int kind) throw r.hash(buffer, viaBuffer); assertThat(viaBuffer).as("hash via buffer kind %d size %d", kind, data.length).isEqualTo(expected); assertThat(buffer.position()).as("buffer position must be untouched").isEqualTo(position); + + int split = data.length / 2; + var first = buffer.duplicate().limit(position + split); + var second = buffer.duplicate().position(position + split); + var stream = DataStream.ofByteBuffers(List.of(first, second), data.length, "application/octet-stream"); + byte[] viaStream = new byte[32]; + r.hash(stream, viaStream); + assertThat(viaStream).as("hash via stream kind %d size %d", kind, data.length).isEqualTo(expected); + assertThat(first.position()).as("stream view position must be untouched").isEqualTo(position); } @Test diff --git a/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java b/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java index 80de7b1251..142cd2a34e 100644 --- a/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java +++ b/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java @@ -13,6 +13,7 @@ import software.amazon.smithy.java.client.http.AmznErrorHeaderExtractor; import software.amazon.smithy.java.client.http.HttpClientProtocol; import software.amazon.smithy.java.client.http.HttpErrorDeserializer; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.schema.ApiOperation; import software.amazon.smithy.java.core.schema.SerializableStruct; @@ -54,6 +55,7 @@ public AwsJsonProtocol( this.codec = JsonCodec.builder() .defaultNamespace(service.getNamespace()) .useTimestampFormat(true) + .outputBuffers(OutputBuffers.create()) .build(); this.errorDeserializer = HttpErrorDeserializer.builder() @@ -94,7 +96,7 @@ public HttpRequest builder.addHeader(HeaderName.X_AMZ_TARGET, target) .addHeader(HeaderName.CONTENT_TYPE, contentType()); } - return builder.setBody(DataStream.ofByteBuffer(codec.serialize(input), contentType())); + return builder.setBody(codec.serializeToDataStream(input, contentType())); } @Override diff --git a/aws/client/aws-client-awsquery/src/jmh/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerBenchmark.java b/aws/client/aws-client-awsquery/src/jmh/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerBenchmark.java index 4f921cde22..bffe4d0390 100644 --- a/aws/client/aws-client-awsquery/src/jmh/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerBenchmark.java +++ b/aws/client/aws-client-awsquery/src/jmh/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerBenchmark.java @@ -11,6 +11,7 @@ import org.openjdk.jmh.annotations.Scope; import org.openjdk.jmh.annotations.Setup; import org.openjdk.jmh.annotations.State; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.schema.PreludeSchemas; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.model.shapes.ShapeId; @@ -54,7 +55,8 @@ public ByteBuffer serializeString() { QueryFormSerializer serializer = QueryFormSerializer.acquire( QueryFormSerializer.QueryVariant.AWS_QUERY, "Benchmark", - "2020-01-01"); + "2020-01-01", + OutputBuffers.shared()); serializer.writeString(member, value); return serializer.finish(); } diff --git a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java index bf0d6bced4..67c3792c29 100644 --- a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java +++ b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java @@ -16,6 +16,7 @@ import software.amazon.smithy.java.client.core.ProtocolSettings; import software.amazon.smithy.java.client.http.HttpClientProtocol; import software.amazon.smithy.java.client.http.HttpErrorDeserializer; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.error.CallException; import software.amazon.smithy.java.core.error.ModeledException; @@ -43,7 +44,10 @@ public final class AwsQueryClientProtocol extends HttpClientProtocol { private final ShapeId service; private final String version; private final HttpErrorDeserializer errorDeserializer; - private final XmlCodec codec = XmlCodec.builder().strictRootElement(false).build(); + private final XmlCodec codec = XmlCodec.builder() + .strictRootElement(false) + .outputBuffers(OutputBuffers.create()) + .build(); public AwsQueryClientProtocol(ShapeId service, String version) { super(AwsQueryTrait.ID); @@ -70,22 +74,22 @@ public HttpRequest SmithyUri endpoint ) { String operationName = operation.schema().id().getName(); - QueryFormSerializer serializer = QueryFormSerializer.acquire( - QueryFormSerializer.QueryVariant.AWS_QUERY, + QueryFormSerializer serializer = QueryFormSerializer.acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, operationName, - version); + version, + codec.outputBuffers()); if (!operation.inputSchema().hasTrait(TraitKey.UNIT_TYPE_TRAIT)) { input.serializeMembers(serializer); } - ByteBuffer body = serializer.finish(); + DataStream body = serializer.finishLeased(CONTENT_TYPE); return HttpRequest.create() .setMethod("POST") .setUri(endpoint) .setHeader(HeaderName.CONTENT_TYPE, CONTENT_TYPE) - .setBody(DataStream.ofByteBuffer(body, CONTENT_TYPE)); + .setBody(body); } @Override diff --git a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java index 9bd49f69ef..1a89910f42 100644 --- a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java +++ b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java @@ -15,6 +15,7 @@ import software.amazon.smithy.java.client.core.ProtocolSettings; import software.amazon.smithy.java.client.http.HttpClientProtocol; import software.amazon.smithy.java.client.http.HttpErrorDeserializer; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.error.CallException; import software.amazon.smithy.java.core.error.ModeledException; @@ -43,7 +44,10 @@ public final class Ec2QueryClientProtocol extends HttpClientProtocol { private final String version; private final HttpErrorDeserializer errorDeserializer; - private final XmlCodec codec = XmlCodec.builder().strictRootElement(false).build(); + private final XmlCodec codec = XmlCodec.builder() + .strictRootElement(false) + .outputBuffers(OutputBuffers.create()) + .build(); public Ec2QueryClientProtocol(ShapeId service, String version) { super(Ec2QueryTrait.ID); @@ -70,22 +74,22 @@ public HttpRequest SmithyUri endpoint ) { String operationName = operation.schema().id().getName(); - QueryFormSerializer serializer = QueryFormSerializer.acquire( - QueryFormSerializer.QueryVariant.EC2_QUERY, + QueryFormSerializer serializer = QueryFormSerializer.acquire(QueryFormSerializer.QueryVariant.EC2_QUERY, operationName, - version); + version, + codec.outputBuffers()); if (!operation.inputSchema().hasTrait(TraitKey.UNIT_TYPE_TRAIT)) { input.serializeMembers(serializer); } - ByteBuffer body = serializer.finish(); + DataStream body = serializer.finishLeased(CONTENT_TYPE); return HttpRequest.create() .setMethod("POST") .setUri(endpoint) .setHeader(HeaderName.CONTENT_TYPE, CONTENT_TYPE) - .setBody(DataStream.ofByteBuffer(body, CONTENT_TYPE)); + .setBody(body); } @Override diff --git a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializer.java b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializer.java index 8c05601086..373ed81663 100644 --- a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializer.java +++ b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializer.java @@ -14,6 +14,7 @@ import java.util.function.BiConsumer; import software.amazon.smithy.java.codecs.commons.CompactStringAccess; import software.amazon.smithy.java.codecs.commons.NumberCodec; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.codecs.commons.StripedPool; import software.amazon.smithy.java.codecs.commons.TimestampCodec; import software.amazon.smithy.java.core.schema.Schema; @@ -25,6 +26,7 @@ import software.amazon.smithy.java.core.serde.TimestampFormatter; import software.amazon.smithy.java.core.serde.document.Document; import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.model.traits.TimestampFormatTrait; /** @@ -72,30 +74,29 @@ enum QueryVariant { } private static final int DEFAULT_BUF_SIZE = 1024; - private static final int MAX_CACHEABLE_BUF = DEFAULT_BUF_SIZE * 4; private static final int MAX_BYTES_PER_CHAR = 9; private static final int MAX_BYTES_PER_LATIN1_BYTE = 6; - record AcquireContext(QueryVariant variant, String action, String version) {} + record AcquireContext(QueryVariant variant, String action, String version, OutputBuffers pool) {} private static final StripedPool POOL = new StripedPool<>() { @Override protected QueryFormSerializer create(AcquireContext ctx) { - return new QueryFormSerializer(); + return new QueryFormSerializer(ctx.pool()); } @Override - protected boolean canPool(QueryFormSerializer s) { - return true; + protected void cleanup(QueryFormSerializer s) { + s.releaseBuffer(); + // Clear the client pool reference because idle serializers remain in a JVM-wide pool. + s.pool = null; } @Override - protected void prepareForPool(QueryFormSerializer s) { - if (s.buf.length > MAX_CACHEABLE_BUF) { - s.buf = new byte[DEFAULT_BUF_SIZE]; - } + protected boolean canPool(QueryFormSerializer s) { + return true; } @Override @@ -103,12 +104,16 @@ protected boolean reset(QueryFormSerializer s, AcquireContext ctx) { s.prefixLen = 0; s.prefixDepth = 0; s.pos = 0; + s.pool = ctx.pool(); + s.buf = s.pool.acquire(s.sizeHint); return true; } }; private byte[] buf; private int pos; + private int sizeHint = DEFAULT_BUF_SIZE; + private OutputBuffers pool; private QueryVariant variant; private byte[] prefixBuf = new byte[128]; @@ -120,13 +125,21 @@ protected boolean reset(QueryFormSerializer s, AcquireContext ctx) { private final QueryMapSerializer mapSerializer = new QueryMapSerializer(); private final MapValueSerializer mapValueSerializer = new MapValueSerializer(); - private QueryFormSerializer() { - this.buf = new byte[DEFAULT_BUF_SIZE]; + private QueryFormSerializer(OutputBuffers pool) { + this.pool = pool; + this.buf = pool.acquire(DEFAULT_BUF_SIZE); + } + + private void releaseBuffer() { + if (buf != null) { + pool.release(buf); + buf = null; + } } @SuppressWarnings("deprecation") - static QueryFormSerializer acquire(QueryVariant variant, String action, String version) { - QueryFormSerializer s = POOL.acquire(new AcquireContext(variant, action, version)); + static QueryFormSerializer acquire(QueryVariant variant, String action, String version, OutputBuffers pool) { + QueryFormSerializer s = POOL.acquire(new AcquireContext(variant, action, version, pool)); s.variant = variant; int headerLen = ACTION_PREFIX.length + action.length() + VERSION_PREFIX.length + version.length(); @@ -143,22 +156,27 @@ static QueryFormSerializer acquire(QueryVariant variant, String action, String v } ByteBuffer finish() { - ByteBuffer result; - if (buf.length > MAX_CACHEABLE_BUF) { - byte[] resultBuf = buf; - buf = new byte[DEFAULT_BUF_SIZE]; - result = ByteBuffer.wrap(resultBuf, 0, pos).slice(); - } else { - result = ByteBuffer.wrap(Arrays.copyOf(buf, pos)); - } + sizeHint = OutputBuffers.nextSizeHint(sizeHint, pos); + ByteBuffer result = ByteBuffer.wrap(Arrays.copyOf(buf, pos)); POOL.release(this); return result; } + DataStream finishLeased(String contentType) { + byte[] detached = buf; + int length = pos; + var recycler = pool.recycler(); + sizeHint = OutputBuffers.nextSizeHint(sizeHint, length); + buf = null; + pos = 0; + POOL.release(this); + return DataStream.ofLeasedBytes(detached, 0, length, contentType, recycler); + } + private void ensureCapacity(int needed) { int required = pos + needed; if (required > buf.length) { - buf = Arrays.copyOf(buf, Math.max(required, buf.length + (buf.length >> 1))); + buf = pool.grow(buf, pos, required); } } diff --git a/aws/client/aws-client-awsquery/src/test/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerTest.java b/aws/client/aws-client-awsquery/src/test/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerTest.java index 24292901a4..b6b1be7431 100644 --- a/aws/client/aws-client-awsquery/src/test/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerTest.java +++ b/aws/client/aws-client-awsquery/src/test/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerTest.java @@ -24,6 +24,7 @@ import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; import software.amazon.smithy.java.codecs.commons.NumberCodec; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.schema.PreludeSchemas; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableStruct; @@ -524,10 +525,10 @@ void maxBigIntegerLengthBoundsDecimalEncoding() { @Test void finishDoesNotAliasThePooledBuffer() { ByteBuffer first = QueryFormSerializer - .acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A1", "V1") + .acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A1", "V1", OutputBuffers.shared()) .finish(); ByteBuffer second = QueryFormSerializer - .acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A2", "V2") + .acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A2", "V2", OutputBuffers.shared()) .finish(); assertThat(first.array(), not(sameInstance(second.array()))); } @@ -535,23 +536,22 @@ void finishDoesNotAliasThePooledBuffer() { @Test void finishedBodyIsUnaffectedByTheNextSerializer() { ByteBuffer first = QueryFormSerializer - .acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A1", "V1") + .acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A1", "V1", OutputBuffers.shared()) .finish(); QueryFormSerializer - .acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A2", "V2") + .acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A2", "V2", OutputBuffers.shared()) .finish(); assertThat(StandardCharsets.UTF_8.decode(first).toString(), equalTo("Action=A1&Version=V1")); } @Test - void finishTransfersOversizedBuffer() { + void finishReturnsAnExactCopyAndReusesThePooledBuffer() { String value = "a".repeat(8192); ByteBuffer body = serializeStringBuffer(value); - // The returned view has an exact limit, but keeps the oversized backing array rather than - // copying it immediately before the pool would discard it. - assertThat(body.array().length > body.remaining(), equalTo(true)); - QueryFormSerializer.acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A2", "V2").finish(); + assertThat(body.array().length, equalTo(body.remaining())); + QueryFormSerializer.acquire(QueryFormSerializer.QueryVariant.AWS_QUERY, "A2", "V2", OutputBuffers.shared()) + .finish(); assertThat(StandardCharsets.UTF_8.decode(body).toString(), equalTo(header() + "&String=" + value)); } @@ -575,7 +575,8 @@ private ByteBuffer serializeStringBuffer(String value) { QueryFormSerializer serializer = QueryFormSerializer.acquire( QueryFormSerializer.QueryVariant.AWS_QUERY, "TestAction", - "2020-01-01"); + "2020-01-01", + OutputBuffers.shared()); serializer.writeString(struct.member("String"), value); return serializer.finish(); } @@ -619,7 +620,8 @@ public T getMemberValue(Schema m) { } }; - QueryFormSerializer s = QueryFormSerializer.acquire(variant, "TestAction", "2020-01-01"); + QueryFormSerializer s = + QueryFormSerializer.acquire(variant, "TestAction", "2020-01-01", OutputBuffers.shared()); s.writeStruct(structSchema, struct); return StandardCharsets.UTF_8.decode(s.finish()).toString(); } diff --git a/aws/client/aws-client-restjson/src/main/java/software/amazon/smithy/java/aws/client/restjson/RestJsonClientProtocol.java b/aws/client/aws-client-restjson/src/main/java/software/amazon/smithy/java/aws/client/restjson/RestJsonClientProtocol.java index 69d1b131ab..3e1c992c8a 100644 --- a/aws/client/aws-client-restjson/src/main/java/software/amazon/smithy/java/aws/client/restjson/RestJsonClientProtocol.java +++ b/aws/client/aws-client-restjson/src/main/java/software/amazon/smithy/java/aws/client/restjson/RestJsonClientProtocol.java @@ -16,6 +16,7 @@ import software.amazon.smithy.java.client.http.HttpErrorDeserializer; import software.amazon.smithy.java.client.http.binding.HttpBindingClientProtocol; import software.amazon.smithy.java.client.http.binding.HttpBindingErrorFactory; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.schema.ApiOperation; import software.amazon.smithy.java.core.schema.SerializableStruct; @@ -48,6 +49,7 @@ public RestJsonClientProtocol(ShapeId service) { .useJsonName(true) .useTimestampFormat(true) .defaultNamespace(service.getNamespace()) + .outputBuffers(OutputBuffers.create()) .build(); this.errorDeserializer = HttpErrorDeserializer.builder() diff --git a/aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java b/aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java index e754cdb2fc..b94ab3f2e7 100644 --- a/aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java +++ b/aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java @@ -18,6 +18,7 @@ import software.amazon.smithy.java.client.http.HttpErrorDeserializer; import software.amazon.smithy.java.client.http.binding.HttpBindingClientProtocol; import software.amazon.smithy.java.client.http.binding.HttpBindingErrorFactory; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.error.CallException; import software.amazon.smithy.java.core.error.ModeledException; @@ -64,6 +65,7 @@ public RestXmlClientProtocol(ShapeId service, Schema serviceSchema) { this.codec = XmlCodec.builder() .defaultNamespace(xmlNamespace) .strictRootElement(false) + .outputBuffers(OutputBuffers.create()) .build(); this.errorDeserializer = HttpErrorDeserializer.builder() .codec(codec) diff --git a/benchmarks/e2e-benchmarks/src/main/java/software/amazon/smithy/java/benchmarks/e2e/BenchmarkClient.java b/benchmarks/e2e-benchmarks/src/main/java/software/amazon/smithy/java/benchmarks/e2e/BenchmarkClient.java index b271b9b40b..015079eb59 100644 --- a/benchmarks/e2e-benchmarks/src/main/java/software/amazon/smithy/java/benchmarks/e2e/BenchmarkClient.java +++ b/benchmarks/e2e-benchmarks/src/main/java/software/amazon/smithy/java/benchmarks/e2e/BenchmarkClient.java @@ -91,7 +91,7 @@ SerializableStruct invoke() throws Throwable { if (outputPayloadMember != null) { Object payload = output.getMemberValue(outputPayloadMember); if (payload instanceof DataStream stream) { - Bodies.drain(stream); + Bodies.drainPayload(stream); } else if (payload instanceof ByteBuffer buffer) { Bodies.consume(buffer); } diff --git a/benchmarks/e2e-benchmarks/src/main/java/software/amazon/smithy/java/benchmarks/e2e/Bodies.java b/benchmarks/e2e-benchmarks/src/main/java/software/amazon/smithy/java/benchmarks/e2e/Bodies.java index de12f2f0e8..d8e6cfc987 100644 --- a/benchmarks/e2e-benchmarks/src/main/java/software/amazon/smithy/java/benchmarks/e2e/Bodies.java +++ b/benchmarks/e2e-benchmarks/src/main/java/software/amazon/smithy/java/benchmarks/e2e/Bodies.java @@ -6,6 +6,7 @@ package software.amazon.smithy.java.benchmarks.e2e; import java.io.IOException; +import java.io.OutputStream; import java.io.UncheckedIOException; import java.nio.ByteBuffer; import software.amazon.smithy.java.io.datastream.DataStream; @@ -13,32 +14,68 @@ /** Reads bodies to the end, as a real transport and a real caller would. */ final class Bodies { - private static final ThreadLocal SCRATCH = ThreadLocal.withInitial(() -> new byte[8192]); + private static final ThreadLocal SCRATCH = ThreadLocal.withInitial(DrainingSink::new); private Bodies() {} - /** Returns the number of bytes the body held. */ + /** + * Keep request and response writeTo call sites separate so HotSpot profiles their body types independently. + */ static long drain(DataStream body) { if (body == null) { return 0; } - byte[] buf = SCRATCH.get(); - long total = 0; - try (var in = body.asInputStream()) { - int n; - while ((n = in.read(buf)) >= 0) { - total += n; - } + DrainingSink sink = SCRATCH.get(); + sink.total = 0; + try { + body.writeTo(sink); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + return sink.total; + } + + static long drainPayload(DataStream payload) { + if (payload == null) { + return 0; + } + DrainingSink sink = SCRATCH.get(); + sink.total = 0; + try { + payload.writeTo(sink); } catch (IOException e) { throw new UncheckedIOException(e); } - return total; + return sink.total; + } + + private static final class DrainingSink extends OutputStream { + private final byte[] scratch = new byte[8192]; + long total; + + @Override + public void write(int b) { + scratch[0] = (byte) b; + total++; + } + + @Override + public void write(byte[] b, int off, int len) { + // Copy each byte to model a socket buffer. + int remaining = len; + while (remaining > 0) { + int n = Math.min(remaining, scratch.length); + System.arraycopy(b, off + (len - remaining), scratch, 0, n); + remaining -= n; + } + total += len; + } } /** Reads every byte because deserialization may leave a payload in a buffer without copying it. */ static void consume(ByteBuffer payload) { ByteBuffer b = payload.duplicate(); - byte[] buf = SCRATCH.get(); + byte[] buf = SCRATCH.get().scratch; while (b.hasRemaining()) { b.get(buf, 0, Math.min(buf.length, b.remaining())); } diff --git a/benchmarks/e2e-benchmarks/src/test/java/software/amazon/smithy/java/benchmarks/e2e/BodiesTest.java b/benchmarks/e2e-benchmarks/src/test/java/software/amazon/smithy/java/benchmarks/e2e/BodiesTest.java new file mode 100644 index 0000000000..859f0f6da1 --- /dev/null +++ b/benchmarks/e2e-benchmarks/src/test/java/software/amazon/smithy/java/benchmarks/e2e/BodiesTest.java @@ -0,0 +1,33 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.benchmarks.e2e; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.ArrayList; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.smithy.java.io.datastream.DataStream; + +class BodiesTest { + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void drainsPooledStorageWithoutEscapingOrReleasingIt(boolean responsePayload) { + byte[] storage = new byte[20_000]; + var recycled = new ArrayList(); + var body = DataStream.ofLeasedBytes(storage, 7, 19_000, null, recycled::add); + + long count = responsePayload ? Bodies.drainPayload(body) : Bodies.drain(body); + + assertThat(count).isEqualTo(19_000); + assertThat(body.escaped()).isFalse(); + assertThat(body.isReleased()).isFalse(); + assertThat(recycled).isEmpty(); + body.release(); + assertThat(recycled).containsExactly(storage); + } +} diff --git a/benchmarks/serde-benchmarks/README.md b/benchmarks/serde-benchmarks/README.md index 427de2aac5..dd0abaef7d 100644 --- a/benchmarks/serde-benchmarks/README.md +++ b/benchmarks/serde-benchmarks/README.md @@ -18,6 +18,11 @@ uses internally, minus the actual HTTP transport. | `restXml` | `RestXmlClientProtocol#createRequest` | `RestXmlClientProtocol#deserializeResponse` | | `rpcv2Cbor` | `RpcV2CborProtocol#createRequest` | `RpcV2CborProtocol#deserializeResponse` | +The `*CodecSerializeBenchmark` classes (AWS JSON 1.0 and RPC v2 CBOR) drive +`Codec#serialize` on the same fixtures instead: the copy-out path that server +responses and other direct codec callers take, as opposed to the leased +request-body path behind `createRequest`. + Each benchmark loop therefore measures the *same* per-iteration work the production code does for that protocol: header / URI label / query string binding (where applicable) plus body serialization on the way out, and the diff --git a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsJson1_0CodecSerializeBenchmark.java b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsJson1_0CodecSerializeBenchmark.java new file mode 100644 index 0000000000..440153d42b --- /dev/null +++ b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsJson1_0CodecSerializeBenchmark.java @@ -0,0 +1,69 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.benchmarks.serde; + +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.infra.Blackhole; +import software.amazon.smithy.java.aws.client.awsjson.AwsJson1Protocol; +import software.amazon.smithy.java.core.serde.Codec; +import software.amazon.smithy.model.shapes.ShapeId; + +@State(Scope.Benchmark) +public class AwsJson1_0CodecSerializeBenchmark { + + private static final String GENERATED_PACKAGE = + "software.amazon.smithy.java.benchmarks.serde.generated.awsjson10.model"; + private static final ShapeId SERVICE_ID = + ShapeId.from("com.amazonaws.sdk.benchmark#AwsJsonRpc10DataPlane"); + + @Param("generic") + public String implementation; + + @Param({ + "awsJson1_0_GetItemInput_Baseline", + "awsJson1_0_HealthcheckRequest_Example", + "awsJson1_0_WideTypesRequest_S", + "awsJson1_0_WideTypesRequest_M", + "awsJson1_0_WideTypesRequest_L", + "awsJson1_0_PutItemRequest_Baseline", + "awsJson1_0_PutItemRequest_ShallowMap_S", + "awsJson1_0_PutItemRequest_ShallowMap_M", + "awsJson1_0_PutItemRequest_ShallowMap_L", + "awsJson1_0_PutItemRequest_Nested_M", + "awsJson1_0_PutItemRequest_Nested_L", + "awsJson1_0_PutItemRequest_MixedItem_S", + "awsJson1_0_PutItemRequest_MixedItem_M", + "awsJson1_0_PutItemRequest_MixedItem_L", + "awsJson1_0_PutItemRequest_BinaryData_S", + "awsJson1_0_PutItemRequest_BinaryData_M", + "awsJson1_0_PutItemRequest_BinaryData_L", + }) + public String testCaseId; + + private Codec codec; + private SerializeState state; + + @Setup + public void setup() { + boolean generated = "generated".equals(implementation); + if (generated && Runtime.version().feature() < 25) { + throw new IllegalStateException("Runtime codegen unavailable on " + Runtime.version()); + } + System.setProperty("smithy-java.runtime-codegen", generated ? "enabled" : "disabled"); + System.setProperty("smithy-java.json-provider", "smithy"); + codec = new AwsJson1Protocol(SERVICE_ID).payloadCodec(); + state = SerializeState.forTestCase(testCaseId, GENERATED_PACKAGE, SERVICE_ID); + } + + @Benchmark + public void serialize(Blackhole bh) { + bh.consume(codec.serialize(state.input)); + } +} diff --git a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsJson1_0SerializeBenchmark.java b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsJson1_0SerializeBenchmark.java index 686f1e2653..73157a7301 100644 --- a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsJson1_0SerializeBenchmark.java +++ b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsJson1_0SerializeBenchmark.java @@ -72,7 +72,11 @@ public void setup() { @Benchmark public void serialize(Blackhole bh) { - bh.consume(protocol.createRequest(operation(), state.input, state.context, state.endpoint)); + var request = protocol.createRequest(operation(), state.input, state.context, state.endpoint); + bh.consume(request); + if (SerializeState.RELEASE_REQUESTS) { + SerializeState.releasePooledBody(protocol, request); + } } @SuppressWarnings({"unchecked", "rawtypes"}) diff --git a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsQuerySerializeBenchmark.java b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsQuerySerializeBenchmark.java index 3bfc31f673..c1434df94c 100644 --- a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsQuerySerializeBenchmark.java +++ b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/AwsQuerySerializeBenchmark.java @@ -58,6 +58,10 @@ public void serialize(Blackhole bh) { @SuppressWarnings({"unchecked", "rawtypes"}) ApiOperation op = (ApiOperation) state.operation; - bh.consume(protocol.createRequest(op, state.input, state.context, state.endpoint)); + var request = protocol.createRequest(op, state.input, state.context, state.endpoint); + bh.consume(request); + if (SerializeState.RELEASE_REQUESTS) { + SerializeState.releasePooledBody(protocol, request); + } } } diff --git a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RestJson1SerializeBenchmark.java b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RestJson1SerializeBenchmark.java index 0e6e0710dc..ca7a1c54bb 100644 --- a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RestJson1SerializeBenchmark.java +++ b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RestJson1SerializeBenchmark.java @@ -64,7 +64,11 @@ public void setup() { @Benchmark public void serialize(Blackhole bh) { - bh.consume(protocol.createRequest(operation(), state.input, state.context, state.endpoint)); + var request = protocol.createRequest(operation(), state.input, state.context, state.endpoint); + bh.consume(request); + if (SerializeState.RELEASE_REQUESTS) { + SerializeState.releasePooledBody(protocol, request); + } } @SuppressWarnings({"unchecked", "rawtypes"}) diff --git a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RestXmlSerializeBenchmark.java b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RestXmlSerializeBenchmark.java index c9010f494b..7b1d6694e8 100644 --- a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RestXmlSerializeBenchmark.java +++ b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RestXmlSerializeBenchmark.java @@ -55,6 +55,10 @@ public void serialize(Blackhole bh) { @SuppressWarnings({"unchecked", "rawtypes"}) ApiOperation op = (ApiOperation) state.operation; - bh.consume(protocol.createRequest(op, state.input, state.context, state.endpoint)); + var request = protocol.createRequest(op, state.input, state.context, state.endpoint); + bh.consume(request); + if (SerializeState.RELEASE_REQUESTS) { + SerializeState.releasePooledBody(protocol, request); + } } } diff --git a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RpcV2CborCodecSerializeBenchmark.java b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RpcV2CborCodecSerializeBenchmark.java new file mode 100644 index 0000000000..ca96780499 --- /dev/null +++ b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RpcV2CborCodecSerializeBenchmark.java @@ -0,0 +1,58 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.benchmarks.serde; + +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.infra.Blackhole; +import software.amazon.smithy.java.client.rpcv2.RpcV2CborProtocol; +import software.amazon.smithy.java.core.serde.Codec; +import software.amazon.smithy.model.shapes.ShapeId; + +@State(Scope.Benchmark) +public class RpcV2CborCodecSerializeBenchmark { + + private static final String GENERATED_PACKAGE = + "software.amazon.smithy.java.benchmarks.serde.generated.rpcv2cbor.model"; + private static final ShapeId SERVICE_ID = + ShapeId.from("com.amazonaws.sdk.benchmark#SmithyRpcV2CborDataPlane"); + + @Param({ + "rpcv2Cbor_WideTypesRequest_S", + "rpcv2Cbor_WideTypesRequest_M", + "rpcv2Cbor_WideTypesRequest_L", + "rpcv2Cbor_PutItemRequest_Baseline", + "rpcv2Cbor_PutItemRequest_ShallowMap_S", + "rpcv2Cbor_PutItemRequest_ShallowMap_M", + "rpcv2Cbor_PutItemRequest_ShallowMap_L", + "rpcv2Cbor_PutItemRequest_Nested_M", + "rpcv2Cbor_PutItemRequest_Nested_L", + "rpcv2Cbor_PutItemRequest_MixedItem_S", + "rpcv2Cbor_PutItemRequest_MixedItem_M", + "rpcv2Cbor_PutItemRequest_MixedItem_L", + "rpcv2Cbor_PutItemRequest_BinaryData_S", + "rpcv2Cbor_PutItemRequest_BinaryData_M", + "rpcv2Cbor_PutItemRequest_BinaryData_L", + }) + public String testCaseId; + + private Codec codec; + private SerializeState state; + + @Setup + public void setup() { + codec = new RpcV2CborProtocol(SERVICE_ID).payloadCodec(); + state = SerializeState.forTestCase(testCaseId, GENERATED_PACKAGE, SERVICE_ID); + } + + @Benchmark + public void serialize(Blackhole bh) { + bh.consume(codec.serialize(state.input)); + } +} diff --git a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RpcV2CborSerializeBenchmark.java b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RpcV2CborSerializeBenchmark.java index 191652d913..9e0be638a6 100644 --- a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RpcV2CborSerializeBenchmark.java +++ b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/RpcV2CborSerializeBenchmark.java @@ -61,6 +61,10 @@ public void serialize(Blackhole bh) { @SuppressWarnings({"unchecked", "rawtypes"}) ApiOperation op = (ApiOperation) state.operation; - bh.consume(protocol.createRequest(op, state.input, state.context, state.endpoint)); + var request = protocol.createRequest(op, state.input, state.context, state.endpoint); + bh.consume(request); + if (SerializeState.RELEASE_REQUESTS) { + SerializeState.releasePooledBody(protocol, request); + } } } diff --git a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/SerializeState.java b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/SerializeState.java index 1aefa96747..b7e28a3159 100644 --- a/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/SerializeState.java +++ b/benchmarks/serde-benchmarks/src/jmh/java/software/amazon/smithy/java/benchmarks/serde/SerializeState.java @@ -7,9 +7,11 @@ import java.lang.reflect.Method; import software.amazon.smithy.java.benchmarks.serde.BenchmarkTestCases.RequestEntry; +import software.amazon.smithy.java.client.http.HttpClientProtocol; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.schema.ApiOperation; import software.amazon.smithy.java.core.schema.SerializableStruct; +import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.io.uri.SmithyUri; import software.amazon.smithy.java.protocoltests.harness.ProtocolTestDocument; import software.amazon.smithy.model.node.Node; @@ -34,6 +36,10 @@ final class SerializeState { private static final SmithyUri ENDPOINT = SmithyUri.of("http://localhost/"); + // Disable request release to measure allocation when pooled arrays never return to the pool. + static final boolean RELEASE_REQUESTS = + Boolean.parseBoolean(System.getProperty("smithy.bench.releaseRequests", "true")); + final ApiOperation operation; final SerializableStruct input; final Context context; @@ -49,6 +55,13 @@ private SerializeState( this.endpoint = ENDPOINT; } + static void releasePooledBody(HttpClientProtocol protocol, HttpRequest request) { + var pooled = protocol.pooledRequestBody(request); + if (pooled != null) { + pooled.release(); + } + } + /** * Build the serialize state from a request test case. * diff --git a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java index 86450c1515..c525874978 100644 --- a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java +++ b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientPipeline.java @@ -32,6 +32,7 @@ import software.amazon.smithy.java.endpoints.Endpoint; import software.amazon.smithy.java.endpoints.EndpointAuthScheme; import software.amazon.smithy.java.endpoints.EndpointResolverParams; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.java.io.uri.SmithyUri; import software.amazon.smithy.java.logging.InternalLogger; import software.amazon.smithy.java.retries.api.AcquireInitialTokenRequest; @@ -122,16 +123,33 @@ O send(ClientCall(call.operation, call.context, input, request); + // Interceptors can replace the request body. Keep the original lease until all attempts complete. + // Keep the lease local because concurrent calls can share a context. + PooledDataStream pooledBody = protocol.pooledRequestBody(request); + try { + var requestHook = new RequestHook<>(call.operation, call.context, input, request); + + // 6. Interceptors: Invoke ReadAfterSerialization. + call.interceptor.readAfterSerialization(requestHook); - // 6. Interceptors: Invoke ReadAfterSerialization. - call.interceptor.readAfterSerialization(requestHook); + // 7. Interceptors: Invoke ModifyBeforeRetryLoop. + requestHook = requestHook.withRequest(call.interceptor.modifyBeforeRetryLoop(requestHook)); - // 7. Interceptors: Invoke ModifyBeforeRetryLoop. - request = call.interceptor.modifyBeforeRetryLoop(requestHook); - requestHook = requestHook.withRequest(request); + return acquireRetryToken(call, requestHook); + } finally { + releasePooledRequestBody(pooledBody); + } + } - return acquireRetryToken(call, requestHook); + private static void releasePooledRequestBody(PooledDataStream pooled) { + if (pooled != null) { + try { + pooled.release(); + } catch (RuntimeException e) { + // A failure to recycle memory must not change the call's result. + LOGGER.debug("Failed to release pooled request body", e); + } + } } private O acquireRetryToken( diff --git a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientProtocol.java b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientProtocol.java index ff21eed2bf..4743e28cc5 100644 --- a/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientProtocol.java +++ b/client/client-core/src/main/java/software/amazon/smithy/java/client/core/ClientProtocol.java @@ -12,6 +12,7 @@ import software.amazon.smithy.java.core.serde.Codec; import software.amazon.smithy.java.core.serde.TypeRegistry; import software.amazon.smithy.java.endpoints.Endpoint; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.java.io.uri.SmithyUri; import software.amazon.smithy.model.shapes.ShapeId; @@ -73,6 +74,19 @@ RequestT createRequ */ RequestT setServiceEndpoint(RequestT request, Endpoint endpoint); + /** + * Returns the request's pooled body, or null if the body is not pooled. + * + *

The pipeline captures this body before interceptors can replace it. + * The pipeline releases it after the call completes, including retries. The default returns null. + * + * @param request request returned by {@link #createRequest} + * @return the pooled body, or null + */ + default PooledDataStream pooledRequestBody(RequestT request) { + return null; + } + /** * Deserializes the output from the transport response or throws a modeled or unmodeled exception. * diff --git a/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientPipelineTest.java b/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientPipelineTest.java index 4170a53b16..e6550883e6 100644 --- a/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientPipelineTest.java +++ b/client/client-core/src/test/java/software/amazon/smithy/java/client/core/ClientPipelineTest.java @@ -36,6 +36,11 @@ import software.amazon.smithy.java.client.http.mock.MockPlugin; import software.amazon.smithy.java.client.http.mock.MockQueue; import software.amazon.smithy.java.context.Context; +import software.amazon.smithy.java.core.error.CallException; +import software.amazon.smithy.java.core.schema.ApiOperation; +import software.amazon.smithy.java.core.schema.SerializableStruct; +import software.amazon.smithy.java.core.serde.Codec; +import software.amazon.smithy.java.core.serde.TypeRegistry; import software.amazon.smithy.java.core.serde.document.Document; import software.amazon.smithy.java.dynamicclient.DynamicClient; import software.amazon.smithy.java.endpoints.Endpoint; @@ -44,6 +49,8 @@ import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.http.api.HttpResponse; import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; +import software.amazon.smithy.java.io.uri.SmithyUri; import software.amazon.smithy.java.retries.api.AcquireInitialTokenRequest; import software.amazon.smithy.java.retries.api.AcquireInitialTokenResponse; import software.amazon.smithy.java.retries.api.RecordSuccessRequest; @@ -681,6 +688,117 @@ public Class identityType() { }; } + @Test + public void releasesThePooledBodyAfterRetriesEvenWhenAnInterceptorReplacesItInPlace() { + var service = ShapeId.from("smithy.example#Sprockets"); + var created = new ArrayList(); + var pooledBodies = new ArrayList(); + var mockQueue = new MockQueue() + .enqueue(HttpResponse.create() + .setStatusCode(429) + .setBody(DataStream.ofString("{\"__type\":\"ExpiredToken\"}")) + .toUnmodifiable()) + .enqueue(HttpResponse.create() + .setStatusCode(200) + .setBody(DataStream.ofString("{\"id\":\"1\"}")) + .toUnmodifiable()); + var mock = MockPlugin.builder().addQueue(mockQueue).build(); + + var client = DynamicClient.builder() + .serviceId(service) + .model(MODEL) + .protocol(new CapturingProtocol(new RestJsonClientProtocol(service), created, pooledBodies)) + .addPlugin(mock) + .addInterceptor(new ClientInterceptor() { + @Override + public RequestT modifyBeforeRetryLoop(RequestHook hook) { + var req = (HttpRequest) hook.request(); + return hook.asRequestType( + req.toModifiable().setBody(DataStream.ofString("{\"id\":\"swapped\"}"))); + } + }) + .endpointResolver(EndpointResolver.staticEndpoint("https://localhost:8081")) + .authSchemeResolver(AuthSchemeResolver.NO_AUTH) + .retryStrategy(retryingStrategy(3)) + .build(); + + var response = client.call("GetSprocket", Document.ofObject(Map.of("id", "1"))); + + assertThat(response.getMember("id").asString(), equalTo("1")); + assertThat(mockQueue.remaining(), is(0)); + assertThat(mock.getRequests(), Matchers.hasSize(2)); + + assertThat(created, Matchers.hasSize(1)); + assertThat(pooledBodies, Matchers.hasSize(1)); + assertThat(created.get(0).body(), Matchers.not(instanceOf(PooledDataStream.class))); + + var pooled = pooledBodies.get(0); + assertThat(pooled.isReleased(), is(true)); + assertThat(pooled.escaped(), is(false)); + Assertions.assertThrows(IllegalStateException.class, pooled::asByteBuffer); + } + + @Test + public void releasesThePooledBodyWhenTheCallFails() { + var service = ShapeId.from("smithy.example#Sprockets"); + var created = new ArrayList(); + var pooledBodies = new ArrayList(); + var mockQueue = new MockQueue() + .enqueue(HttpResponse.create() + .setStatusCode(500) + .setBody(DataStream.ofString("{\"__type\":\"ServiceFooError\"}")) + .toUnmodifiable()); + var mock = MockPlugin.builder().addQueue(mockQueue).build(); + + var client = DynamicClient.builder() + .serviceId(service) + .model(MODEL) + .protocol(new CapturingProtocol(new RestJsonClientProtocol(service), created, pooledBodies)) + .addPlugin(mock) + .endpointResolver(EndpointResolver.staticEndpoint("https://localhost:8081")) + .authSchemeResolver(AuthSchemeResolver.NO_AUTH) + .retryStrategy(retryingStrategy(1)) + .build(); + + Assertions.assertThrows(CallException.class, + () -> client.call("GetSprocket", Document.ofObject(Map.of("id", "1")))); + + assertThat(pooledBodies, Matchers.hasSize(1)); + assertThat(pooledBodies.get(0).isReleased(), is(true)); + } + + private static RetryStrategy retryingStrategy(int maxAttempts) { + return new RetryStrategy() { + @Override + public AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request) { + return new AcquireInitialTokenResponse(new Token(0), Duration.ZERO); + } + + @Override + public RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request) { + if (request.token() instanceof Token t && t.retry + 1 < maxAttempts) { + return new RefreshRetryTokenResponse(new Token(t.retry + 1), Duration.ZERO); + } + throw new TokenAcquisitionFailedException("No more retries"); + } + + @Override + public RecordSuccessResponse recordSuccess(RecordSuccessRequest request) { + return new RecordSuccessResponse(request.token()); + } + + @Override + public int maxAttempts() { + return maxAttempts; + } + + @Override + public Builder toBuilder() { + throw new UnsupportedOperationException(); + } + }; + } + private static final class PropertiesAuthScheme implements AuthScheme { private final ShapeId schemeId; private final Context.Key baseKey; @@ -739,6 +857,74 @@ public Signer signer() { } } + private static final class CapturingProtocol implements ClientProtocol { + private final ClientProtocol delegate; + private final List created; + private final List pooledBodies; + + CapturingProtocol( + ClientProtocol delegate, + List created, + List pooledBodies + ) { + this.delegate = delegate; + this.created = created; + this.pooledBodies = pooledBodies; + } + + @Override + public ShapeId id() { + return delegate.id(); + } + + @Override + public Codec payloadCodec() { + return delegate.payloadCodec(); + } + + @Override + public MessageExchange messageExchange() { + return delegate.messageExchange(); + } + + @Override + public HttpRequest createRequest( + ApiOperation operation, + I input, + Context context, + SmithyUri endpoint + ) { + var request = delegate.createRequest(operation, input, context, endpoint); + created.add(request); + return request; + } + + @Override + public PooledDataStream pooledRequestBody(HttpRequest request) { + var pooled = delegate.pooledRequestBody(request); + if (pooled != null) { + pooledBodies.add(pooled); + } + return pooled; + } + + @Override + public HttpRequest setServiceEndpoint(HttpRequest request, Endpoint endpoint) { + return delegate.setServiceEndpoint(request, endpoint); + } + + @Override + public O deserializeResponse( + ApiOperation operation, + Context context, + TypeRegistry errorRegistry, + HttpRequest request, + HttpResponse response + ) { + return delegate.deserializeResponse(operation, context, errorRegistry, request, response); + } + } + private static final class PropertiesFailingAuthScheme implements AuthScheme { private final ShapeId schemeId; diff --git a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpClientProtocol.java b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpClientProtocol.java index a717de4668..dcc29d2a56 100644 --- a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpClientProtocol.java +++ b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpClientProtocol.java @@ -10,6 +10,7 @@ import software.amazon.smithy.java.endpoints.Endpoint; import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.http.api.HttpResponse; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.model.shapes.ShapeId; /** @@ -33,6 +34,25 @@ public MessageExchange messageExchange() { return HttpMessageExchange.INSTANCE; } + @Override + public PooledDataStream pooledRequestBody(HttpRequest request) { + return request.body() instanceof PooledDataStream pooled ? pooled : null; + } + + /** + * Copies a pooled request body so the request remains readable after the call completes. + * + *

Call this before the pipeline releases the body. + * + * @return an unmodifiable request whose body outlives the call + */ + public static HttpRequest retainRequest(HttpRequest request) { + if (request.body() instanceof PooledDataStream pooled) { + return request.toModifiableCopy().setBody(pooled.copy()).toUnmodifiable(); + } + return request.toUnmodifiable(); + } + @Override public HttpRequest setServiceEndpoint(HttpRequest request, Endpoint endpoint) { var merged = request.uri().withEndpoint(endpoint.uri()); diff --git a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/compression/Gzip.java b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/compression/Gzip.java index 60815a9c92..3e1e850bd5 100644 --- a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/compression/Gzip.java +++ b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/compression/Gzip.java @@ -27,11 +27,10 @@ public DataStream compress(DataStream data) { -1); } - try (var bos = new ByteBufferOutputStream(); - var in = data.asInputStream()) { - var gzip = new GZIPOutputStream(bos); - in.transferTo(gzip); - gzip.close(); + try (var bos = new ByteBufferOutputStream()) { + try (var gzip = new GZIPOutputStream(bos)) { + data.writeTo(gzip); + } return DataStream.ofByteBuffer(bos.toByteBuffer(), data.contentType()); } catch (IOException e) { throw new UncheckedIOException(e); diff --git a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/plugins/HttpChecksumPlugin.java b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/plugins/HttpChecksumPlugin.java index 20a733cf9e..427dd5287c 100644 --- a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/plugins/HttpChecksumPlugin.java +++ b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/plugins/HttpChecksumPlugin.java @@ -16,7 +16,6 @@ import software.amazon.smithy.java.core.schema.TraitKey; import software.amazon.smithy.java.http.api.HeaderName; import software.amazon.smithy.java.http.api.HttpRequest; -import software.amazon.smithy.java.io.ByteBufferUtils; import software.amazon.smithy.model.traits.HttpChecksumRequiredTrait; import software.amazon.smithy.utils.SmithyInternalApi; @@ -50,11 +49,10 @@ public RequestT modifyBeforeTransmit(RequestHook hook static HttpRequest addContentMd5Header(HttpRequest request) { var body = request.body(); if (body != null) { - var buffer = body.asByteBuffer(); - var bytes = ByteBufferUtils.getBytes(buffer); try { - byte[] hash = MessageDigest.getInstance("MD5").digest(bytes); - String base64Hash = Base64.getEncoder().encodeToString(hash); + var md5 = MessageDigest.getInstance("MD5"); + body.forEachBuffer(md5::update); + String base64Hash = Base64.getEncoder().encodeToString(md5.digest()); var modifiable = request.toModifiable(); modifiable.headers().setHeader(HeaderName.CONTENT_MD5, base64Hash); return modifiable; diff --git a/client/client-mock-plugin/src/main/java/software/amazon/smithy/java/client/http/mock/MockPlugin.java b/client/client-mock-plugin/src/main/java/software/amazon/smithy/java/client/http/mock/MockPlugin.java index c2a1cba349..62e3dcc13b 100644 --- a/client/client-mock-plugin/src/main/java/software/amazon/smithy/java/client/http/mock/MockPlugin.java +++ b/client/client-mock-plugin/src/main/java/software/amazon/smithy/java/client/http/mock/MockPlugin.java @@ -18,6 +18,7 @@ import software.amazon.smithy.java.client.core.ClientProtocol; import software.amazon.smithy.java.client.core.ClientTransport; import software.amazon.smithy.java.client.core.MessageExchange; +import software.amazon.smithy.java.client.http.HttpClientProtocol; import software.amazon.smithy.java.client.http.HttpMessageExchange; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.schema.ApiOperation; @@ -28,6 +29,7 @@ import software.amazon.smithy.java.http.api.HttpHeaders; import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.http.api.HttpResponse; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.java.io.uri.SmithyUri; import software.amazon.smithy.java.server.Operation; import software.amazon.smithy.java.server.Service; @@ -213,6 +215,11 @@ public HttpRequest setServiceEndpoint(HttpRequest request, Endpoint endpoint) { return delegate.setServiceEndpoint(request, endpoint); } + @Override + public PooledDataStream pooledRequestBody(HttpRequest request) { + return delegate.pooledRequestBody(request); + } + @Override public O deserializeResponse( ApiOperation operation, @@ -242,7 +249,7 @@ public HttpResponse send(Context context, HttpRequest request) { // Track the request here rather than when the request is created because this method is called // more often than the request creation method when retries happen. if (trackRequests) { - requests.add(currentRequest.request().withRequest(request.toUnmodifiable())); + requests.add(currentRequest.request().withRequest(HttpClientProtocol.retainRequest(request))); } MockedResult result = null; diff --git a/client/client-mock-plugin/src/test/java/software/amazon/smithy/java/client/http/mock/MockPluginTest.java b/client/client-mock-plugin/src/test/java/software/amazon/smithy/java/client/http/mock/MockPluginTest.java index e44ba4a39a..cf66101c1d 100644 --- a/client/client-mock-plugin/src/test/java/software/amazon/smithy/java/client/http/mock/MockPluginTest.java +++ b/client/client-mock-plugin/src/test/java/software/amazon/smithy/java/client/http/mock/MockPluginTest.java @@ -15,6 +15,7 @@ import java.net.URI; import java.net.URISyntaxException; +import java.nio.charset.StandardCharsets; import java.util.Map; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Assertions; @@ -22,13 +23,17 @@ import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeResolver; import software.amazon.smithy.java.client.core.interceptors.ClientInterceptor; import software.amazon.smithy.java.client.core.interceptors.OutputHook; +import software.amazon.smithy.java.client.core.interceptors.RequestHook; import software.amazon.smithy.java.core.error.CallException; import software.amazon.smithy.java.core.error.ErrorFault; import software.amazon.smithy.java.core.serde.document.Document; import software.amazon.smithy.java.dynamicclient.DynamicClient; import software.amazon.smithy.java.endpoints.EndpointResolver; import software.amazon.smithy.java.framework.model.InternalFailureException; +import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.http.api.HttpResponse; +import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.java.retries.api.RetrySafety; import software.amazon.smithy.model.Model; import software.amazon.smithy.model.shapes.ShapeId; @@ -228,4 +233,84 @@ public void returnsResultsBasedOnPredicates() throws URISyntaxException { assertThat(mock.getRequests(), hasSize(2)); } + + @Test + public void releasesPooledRequestBodiesAfterTheCallButRecordsReadableCopies() throws URISyntaxException { + var mockQueue = new MockQueue() + .enqueue(HttpResponse.create() + .setStatusCode(200) + .addHeader("content-type", "application/json") + .setBody(DataStream.ofBytes("{\"id\":\"a\"}".getBytes(StandardCharsets.UTF_8))) + .toUnmodifiable()); + var mock = MockPlugin.builder().addQueue(mockQueue).build(); + var liveBody = new AtomicReference(); + + var client = DynamicClient.builder() + .serviceId(SERVICE) + .model(MODEL) + .addPlugin(mock) + .endpointResolver(EndpointResolver.staticEndpoint(new URI("http://localhost"))) + .authSchemeResolver(AuthSchemeResolver.NO_AUTH) + .addInterceptor(new ClientInterceptor() { + @Override + public void readAfterSigning(RequestHook hook) { + liveBody.set(((HttpRequest) hook.request()).body()); + } + }) + .build(); + + var result = client.call("GetSprocket", Document.of(Map.of("id", Document.of("abc")))); + assertThat(result.getMember("id").asString(), equalTo("a")); + + assertThat(liveBody.get(), instanceOf(PooledDataStream.class)); + var pooled = (PooledDataStream) liveBody.get(); + assertThat(pooled.escaped(), is(false)); + Assertions.assertThrows(IllegalStateException.class, pooled::asByteBuffer); + + assertThat(mock.getRequests(), hasSize(1)); + var recorded = mock.getRequests().get(0).request().body(); + assertThat(recorded, not(instanceOf(PooledDataStream.class))); + assertThat(StandardCharsets.UTF_8.decode(recorded.asByteBuffer()).toString(), equalTo("{\"id\":\"abc\"}")); + } + + @Test + public void recordsReadableCopiesOfMetadataWrappedPooledBodies() throws URISyntaxException { + var mockQueue = new MockQueue() + .enqueue(HttpResponse.create() + .setStatusCode(200) + .addHeader("content-type", "application/json") + .setBody(DataStream.ofBytes("{\"id\":\"a\"}".getBytes(StandardCharsets.UTF_8))) + .toUnmodifiable()); + var mock = MockPlugin.builder().addQueue(mockQueue).build(); + var liveBody = new AtomicReference(); + + var client = DynamicClient.builder() + .serviceId(SERVICE) + .model(MODEL) + .addPlugin(mock) + .endpointResolver(EndpointResolver.staticEndpoint(new URI("http://localhost"))) + .authSchemeResolver(AuthSchemeResolver.NO_AUTH) + .addInterceptor(new ClientInterceptor() { + @Override + public RequestT modifyBeforeRetryLoop(RequestHook hook) { + var req = (HttpRequest) hook.request(); + var wrapped = DataStream.withMetadata(req.body(), "application/x-custom", -1); + liveBody.set(wrapped); + return hook.asRequestType(req.toModifiable().setBody(wrapped)); + } + }) + .build(); + + var result = client.call("GetSprocket", Document.of(Map.of("id", Document.of("abc")))); + assertThat(result.getMember("id").asString(), equalTo("a")); + + assertThat(liveBody.get(), instanceOf(PooledDataStream.class)); + Assertions.assertThrows(IllegalStateException.class, () -> liveBody.get().asByteBuffer()); + + assertThat(mock.getRequests(), hasSize(1)); + var recorded = mock.getRequests().get(0).request().body(); + assertThat(recorded, not(instanceOf(PooledDataStream.class))); + assertThat(recorded.contentType(), equalTo("application/x-custom")); + assertThat(StandardCharsets.UTF_8.decode(recorded.asByteBuffer()).toString(), equalTo("{\"id\":\"abc\"}")); + } } diff --git a/client/client-rpcv2-cbor/src/main/java/software/amazon/smithy/java/client/rpcv2/RpcV2CborProtocol.java b/client/client-rpcv2-cbor/src/main/java/software/amazon/smithy/java/client/rpcv2/RpcV2CborProtocol.java index 896dfb4718..cf947ed619 100644 --- a/client/client-rpcv2-cbor/src/main/java/software/amazon/smithy/java/client/rpcv2/RpcV2CborProtocol.java +++ b/client/client-rpcv2-cbor/src/main/java/software/amazon/smithy/java/client/rpcv2/RpcV2CborProtocol.java @@ -10,6 +10,7 @@ import software.amazon.smithy.java.client.core.ClientProtocol; import software.amazon.smithy.java.client.core.ClientProtocolFactory; import software.amazon.smithy.java.client.core.ProtocolSettings; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.serde.Codec; import software.amazon.smithy.model.shapes.ShapeId; import software.amazon.smithy.protocol.traits.Rpcv2CborTrait; @@ -19,7 +20,7 @@ */ public final class RpcV2CborProtocol extends AbstractRpcV2ClientProtocol { private static final String PAYLOAD_MEDIA_TYPE = "application/cbor"; - private static final Codec CBOR_CODEC = Rpcv2CborCodec.builder().build(); + private final Codec codec = Rpcv2CborCodec.builder().outputBuffers(OutputBuffers.create()).build(); public RpcV2CborProtocol(ShapeId service) { super(Rpcv2CborTrait.ID, service, PAYLOAD_MEDIA_TYPE); @@ -27,7 +28,7 @@ public RpcV2CborProtocol(ShapeId service) { @Override protected Codec codec() { - return CBOR_CODEC; + return codec; } public static final class Factory implements ClientProtocolFactory { diff --git a/client/client-rpcv2-json/src/main/java/software/amazon/smithy/java/client/rpcv2json/RpcV2JsonProtocol.java b/client/client-rpcv2-json/src/main/java/software/amazon/smithy/java/client/rpcv2json/RpcV2JsonProtocol.java index ecc13f72b2..aa9493d956 100644 --- a/client/client-rpcv2-json/src/main/java/software/amazon/smithy/java/client/rpcv2json/RpcV2JsonProtocol.java +++ b/client/client-rpcv2-json/src/main/java/software/amazon/smithy/java/client/rpcv2json/RpcV2JsonProtocol.java @@ -10,6 +10,7 @@ import software.amazon.smithy.java.client.core.ClientProtocolFactory; import software.amazon.smithy.java.client.core.ProtocolSettings; import software.amazon.smithy.java.client.rpcv2.AbstractRpcV2ClientProtocol; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.serde.Codec; import software.amazon.smithy.java.json.JsonCodec; import software.amazon.smithy.model.shapes.ShapeId; @@ -33,6 +34,7 @@ public RpcV2JsonProtocol(ShapeId service) { this.codec = JsonCodec.builder() .defaultNamespace(service.getNamespace()) .useStringForArbitraryPrecision(true) + .outputBuffers(OutputBuffers.create()) .build(); } diff --git a/client/client-rpcv2/src/main/java/software/amazon/smithy/java/client/rpcv2/AbstractRpcV2ClientProtocol.java b/client/client-rpcv2/src/main/java/software/amazon/smithy/java/client/rpcv2/AbstractRpcV2ClientProtocol.java index 4bb7cec7e7..e4de7d0fd2 100644 --- a/client/client-rpcv2/src/main/java/software/amazon/smithy/java/client/rpcv2/AbstractRpcV2ClientProtocol.java +++ b/client/client-rpcv2/src/main/java/software/amazon/smithy/java/client/rpcv2/AbstractRpcV2ClientProtocol.java @@ -159,7 +159,7 @@ private static DataStream bodyDataStream(HttpResponse response) { } private DataStream getBody(SerializableStruct input) { - return DataStream.ofByteBuffer(codec().serialize(input), payloadMediaType); + return codec().serializeToDataStream(input, payloadMediaType); } private EventEncoderFactory getEventEncoderFactory(ApiOperation operation) { diff --git a/codecs/cbor-codec/build.gradle.kts b/codecs/cbor-codec/build.gradle.kts index 0e34a37f8b..38e98e6198 100644 --- a/codecs/cbor-codec/build.gradle.kts +++ b/codecs/cbor-codec/build.gradle.kts @@ -12,7 +12,7 @@ extra["moduleName"] = "software.amazon.smithy.java.cbor" dependencies { api(project(":core")) - implementation(project(":codecs:codec-commons", configuration = "shadow")) + api(project(":codecs:codec-commons", configuration = "shadow")) testFixturesImplementation(libs.assertj.core) testImplementation(project(":codecs:json-codec", configuration = "shadow")) } diff --git a/codecs/cbor-codec/src/jmh/java/software/amazon/smithy/java/cbor/CborStringWriteBenchmark.java b/codecs/cbor-codec/src/jmh/java/software/amazon/smithy/java/cbor/CborStringWriteBenchmark.java index d9f682c5f5..710bc719b0 100644 --- a/codecs/cbor-codec/src/jmh/java/software/amazon/smithy/java/cbor/CborStringWriteBenchmark.java +++ b/codecs/cbor-codec/src/jmh/java/software/amazon/smithy/java/cbor/CborStringWriteBenchmark.java @@ -10,6 +10,7 @@ import org.openjdk.jmh.annotations.Scope; import org.openjdk.jmh.annotations.Setup; import org.openjdk.jmh.annotations.State; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; @State(Scope.Thread) public class CborStringWriteBenchmark { @@ -35,7 +36,7 @@ public class CborStringWriteBenchmark { @Setup public void setup() { value = StringWriteTestCases.value(testCaseId); - serializer = new CborSerializer(); + serializer = new CborSerializer(null, OutputBuffers.shared()); } @Benchmark diff --git a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSerdeProvider.java b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSerdeProvider.java index e1bcc2e2c4..71ad4b591d 100644 --- a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSerdeProvider.java +++ b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSerdeProvider.java @@ -10,6 +10,7 @@ import software.amazon.smithy.java.core.schema.SerializableShape; import software.amazon.smithy.java.core.serde.ShapeDeserializer; import software.amazon.smithy.java.core.serde.ShapeSerializer; +import software.amazon.smithy.java.io.datastream.DataStream; public interface CborSerdeProvider { int getPriority(); @@ -23,4 +24,8 @@ public interface CborSerdeProvider { ShapeSerializer newSerializer(OutputStream sink, CborSettings settings); ByteBuffer serialize(SerializableShape shape, CborSettings settings); + + default DataStream serializeToDataStream(SerializableShape shape, CborSettings settings, String contentType) { + return DataStream.ofByteBuffer(serialize(shape, settings), contentType); + } } diff --git a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSerializer.java b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSerializer.java index aa388a9490..1db5f8322b 100644 --- a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSerializer.java +++ b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSerializer.java @@ -38,6 +38,7 @@ import java.util.Arrays; import java.util.function.BiConsumer; import software.amazon.smithy.java.codecs.commons.CompactStringAccess; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.codecs.commons.StripedPool; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableStruct; @@ -46,6 +47,7 @@ import software.amazon.smithy.java.core.serde.ShapeSerializer; import software.amazon.smithy.java.core.serde.SpecificShapeSerializer; import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.model.shapes.ShapeType; final class CborSerializer implements ShapeSerializer { @@ -58,14 +60,15 @@ final class CborSerializer implements ShapeSerializer { private static final int ARRAY_STREAM = TYPE_ARRAY | INDEFINITE; private static final int DEFAULT_BUF_SIZE = 4096; - private static final int MAX_CACHEABLE_BUF = DEFAULT_BUF_SIZE * 4; - private static final StripedPool POOL = new CborStripedPool(); + private static final StripedPool POOL = new CborStripedPool(); byte[] buf; int pos; + private int sizeHint = DEFAULT_BUF_SIZE; private final OutputStream sink; + private OutputBuffers pool; // Bit i of collectionMask records whether level i was opened as indefinite-length. private long collectionMask = 0L; @@ -78,20 +81,15 @@ final class CborSerializer implements ShapeSerializer { private final CborMapSerializer mapSerializer = new CborMapSerializer(); private SerializeDocumentContents serializeDocumentContents; - CborSerializer() { - this.sink = null; - this.buf = new byte[DEFAULT_BUF_SIZE]; - this.pos = 0; - } - - CborSerializer(OutputStream sink) { + CborSerializer(OutputStream sink, OutputBuffers pool) { this.sink = sink; - this.buf = new byte[DEFAULT_BUF_SIZE]; + this.pool = pool; + this.buf = pool.acquire(DEFAULT_BUF_SIZE); this.pos = 0; } - static CborSerializer acquire() { - return POOL.acquire(null); + static CborSerializer acquire(OutputBuffers pool) { + return POOL.acquire(pool); } static void release(CborSerializer serializer, boolean exception) { @@ -99,9 +97,26 @@ static void release(CborSerializer serializer, boolean exception) { } ByteBuffer extractResult() { + sizeHint = OutputBuffers.nextSizeHint(sizeHint, pos); return ByteBuffer.wrap(Arrays.copyOf(buf, pos)); } + DataStream detach(String contentType) { + byte[] detached = buf; + int length = pos; + sizeHint = OutputBuffers.nextSizeHint(sizeHint, length); + buf = null; + pos = 0; + return DataStream.ofLeasedBytes(detached, 0, length, contentType, pool.recycler()); + } + + private void releaseBuffer() { + if (buf != null) { + pool.release(buf); + buf = null; + } + } + private void ensureCapacity(int needed) { if (pos + needed > buf.length) { grow(needed); @@ -109,7 +124,7 @@ private void ensureCapacity(int needed) { } private void grow(int needed) { - buf = Arrays.copyOf(buf, Math.max(buf.length * 2, pos + needed)); + buf = pool.grow(buf, pos, pos + needed); } @Override @@ -134,6 +149,8 @@ public void close() { } } catch (Exception e) { throw new SerializationException(e); + } finally { + releaseBuffer(); } } @@ -618,30 +635,32 @@ static byte[] encodeMemberName(String name) { return result; } - private static class CborStripedPool extends StripedPool { + private static class CborStripedPool extends StripedPool { @Override - protected CborSerializer create(Void ctx) { - return new CborSerializer(); + protected CborSerializer create(OutputBuffers pool) { + return new CborSerializer(null, pool); } @Override - protected boolean canPool(CborSerializer s) { - return s.buf != null && s.sink == null; + protected void cleanup(CborSerializer s) { + s.releaseBuffer(); + // Clear the client pool reference because idle serializers remain in a JVM-wide pool. + s.pool = null; } @Override - protected void prepareForPool(CborSerializer s) { - if (s.buf.length > MAX_CACHEABLE_BUF) { - s.buf = new byte[DEFAULT_BUF_SIZE]; - } + protected boolean canPool(CborSerializer s) { + return s.sink == null; } @Override - protected boolean reset(CborSerializer s, Void ctx) { + protected boolean reset(CborSerializer s, OutputBuffers pool) { s.pos = 0; s.collectionMask = 0L; s.collectionDepth = 0; s.currentFieldNameTable = null; + s.pool = pool; + s.buf = pool.acquire(s.sizeHint); return true; } } diff --git a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSettings.java b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSettings.java index eabd8e313c..29e05c7b81 100644 --- a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSettings.java +++ b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/CborSettings.java @@ -7,6 +7,7 @@ import java.util.Objects; import java.util.ServiceLoader; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; public final class CborSettings { private static final CborSerdeProvider PROVIDER; @@ -41,10 +42,16 @@ public static CborSettings defaultSettings() { private final String defaultNamespace; private final CborSerdeProvider provider; + private final OutputBuffers outputBuffers; private CborSettings(Builder builder) { this.defaultNamespace = builder.defaultNamespace; this.provider = builder.provider; + this.outputBuffers = builder.outputBuffers; + } + + public OutputBuffers outputBuffers() { + return outputBuffers; } public CborSerdeProvider provider() { @@ -62,6 +69,20 @@ public static Builder builder() { public static final class Builder { private String defaultNamespace; private CborSerdeProvider provider = PROVIDER; + private OutputBuffers outputBuffers = OutputBuffers.shared(); + + /** + * Sets the output buffer pool. The default is {@link OutputBuffers#shared()}. + * + *

Use {@link OutputBuffers#create()} for a pool owned by this codec. + * + * @param outputBuffers pool for serialization output + * @return the builder + */ + public Builder outputBuffers(OutputBuffers outputBuffers) { + this.outputBuffers = Objects.requireNonNull(outputBuffers, "outputBuffers"); + return this; + } /** * Sets the default namespace when attempting to deserialize documents that use a relative shape ID. @@ -90,6 +111,7 @@ Builder overrideSerdeProvider(CborSerdeProvider provider) { public Builder updateBuilder(CborSettings settings) { overrideSerdeProvider(settings.provider()); defaultNamespace(settings.defaultNamespace()); + outputBuffers(settings.outputBuffers()); return this; } diff --git a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/DefaultCborSerdeProvider.java b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/DefaultCborSerdeProvider.java index a52bf57018..38abca02eb 100644 --- a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/DefaultCborSerdeProvider.java +++ b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/DefaultCborSerdeProvider.java @@ -10,6 +10,7 @@ import software.amazon.smithy.java.core.schema.SerializableShape; import software.amazon.smithy.java.core.serde.ShapeDeserializer; import software.amazon.smithy.java.core.serde.ShapeSerializer; +import software.amazon.smithy.java.io.datastream.DataStream; final class DefaultCborSerdeProvider implements CborSerdeProvider { @Override @@ -34,12 +35,12 @@ public ShapeDeserializer newDeserializer(ByteBuffer source, CborSettings setting @Override public ShapeSerializer newSerializer(OutputStream sink, CborSettings settings) { - return new CborSerializer(sink); + return new CborSerializer(sink, settings.outputBuffers()); } @Override public ByteBuffer serialize(SerializableShape shape, CborSettings settings) { - var serializer = CborSerializer.acquire(); + var serializer = CborSerializer.acquire(settings.outputBuffers()); boolean exception = false; try { shape.serialize(serializer); @@ -51,4 +52,19 @@ public ByteBuffer serialize(SerializableShape shape, CborSettings settings) { CborSerializer.release(serializer, exception); } } + + @Override + public DataStream serializeToDataStream(SerializableShape shape, CborSettings settings, String contentType) { + var serializer = CborSerializer.acquire(settings.outputBuffers()); + boolean exception = false; + try { + shape.serialize(serializer); + return serializer.detach(contentType); + } catch (Exception t) { + exception = true; + throw t; + } finally { + CborSerializer.release(serializer, exception); + } + } } diff --git a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/Rpcv2CborCodec.java b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/Rpcv2CborCodec.java index b27d08b980..6a605b9f5f 100644 --- a/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/Rpcv2CborCodec.java +++ b/codecs/cbor-codec/src/main/java/software/amazon/smithy/java/cbor/Rpcv2CborCodec.java @@ -7,10 +7,12 @@ import java.io.OutputStream; import java.nio.ByteBuffer; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.schema.SerializableShape; import software.amazon.smithy.java.core.serde.Codec; import software.amazon.smithy.java.core.serde.ShapeDeserializer; import software.amazon.smithy.java.core.serde.ShapeSerializer; +import software.amazon.smithy.java.io.datastream.DataStream; public final class Rpcv2CborCodec implements Codec { private final CborSettings settings; @@ -28,6 +30,11 @@ public ByteBuffer serialize(SerializableShape shape) { return settings.provider().serialize(shape, settings); } + @Override + public DataStream serializeToDataStream(SerializableShape shape, String contentType) { + return settings.provider().serializeToDataStream(shape, settings, contentType); + } + @Override public ShapeSerializer createSerializer(OutputStream sink) { return settings.provider().newSerializer(sink, settings); @@ -70,6 +77,19 @@ public Builder defaultNamespace(String defaultNamespace) { return this; } + /** + * Sets the output buffer pool. The default is {@link OutputBuffers#shared()}. + * + *

Use {@link OutputBuffers#create()} for a pool owned by this codec. + * + * @param outputBuffers pool for serialization output + * @return the builder + */ + public Builder outputBuffers(OutputBuffers outputBuffers) { + settings().outputBuffers(outputBuffers); + return this; + } + /** * Uses a custom CBOR serde provider. * diff --git a/codecs/cbor-codec/src/test/java/software/amazon/smithy/java/cbor/CborCodecTest.java b/codecs/cbor-codec/src/test/java/software/amazon/smithy/java/cbor/CborCodecTest.java index bab7bb0d0f..0a9d80a5b2 100644 --- a/codecs/cbor-codec/src/test/java/software/amazon/smithy/java/cbor/CborCodecTest.java +++ b/codecs/cbor-codec/src/test/java/software/amazon/smithy/java/cbor/CborCodecTest.java @@ -10,11 +10,13 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import java.io.OutputStream; +import java.lang.ref.WeakReference; import java.math.BigDecimal; import java.math.BigInteger; import java.nio.ByteBuffer; @@ -23,10 +25,12 @@ import java.util.Arrays; import java.util.List; import java.util.function.Consumer; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.schema.PreludeSchemas; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableStruct; @@ -36,10 +40,47 @@ import software.amazon.smithy.java.core.serde.ShapeSerializer; import software.amazon.smithy.java.io.ByteBufferOutputStream; import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.model.shapes.ShapeId; class CborCodecTest { + @Test + void idleSerializersDoNotPinACodecsPool() throws Exception { + WeakReference pool = serializeWithADedicatedPoolThenDropIt(); + for (int i = 0; i < 50 && pool.get() != null; i++) { + System.gc(); + Thread.sleep(20); + } + Assertions.assertNull(pool.get(), + "a codec's pool must be collectible once the codec is gone"); + } + + private static WeakReference serializeWithADedicatedPoolThenDropIt() { + var pool = OutputBuffers.create(); + var bird = new SerializableStruct() { + @Override + public Schema schema() { + return CborTestData.BIRD; + } + + @Override + public void serializeMembers(ShapeSerializer serializer) { + serializer.writeString(CborTestData.BIRD_NAME, "tweety"); + } + + @Override + public T getMemberValue(Schema member) { + return null; + } + }; + try (var codec = Rpcv2CborCodec.builder().outputBuffers(pool).build()) { + ((PooledDataStream) codec.serializeToDataStream(bird, "application/cbor")).release(); + codec.serialize(bird); + } + return new WeakReference<>(pool); + } + private static final CborSettings SETTINGS = CborSettings.defaultSettings(); private static final DefaultCborSerdeProvider CODEC = new DefaultCborSerdeProvider(); @@ -51,10 +92,10 @@ private static ByteBuffer roundTrip(SerializableStruct shape) { class SerializerPoolTest { @Test void acquireAndReleaseReturnsToPool() { - var s1 = CborSerializer.acquire(); + var s1 = CborSerializer.acquire(OutputBuffers.shared()); CborSerializer.release(s1, false); - var s2 = CborSerializer.acquire(); + var s2 = CborSerializer.acquire(OutputBuffers.shared()); // Pool may or may not return the same instance depending on thread, // but it should not throw. assertNotNull(s2); @@ -63,30 +104,32 @@ void acquireAndReleaseReturnsToPool() { @Test void acquireProducesCleanState() { - var s = CborSerializer.acquire(); + var s = CborSerializer.acquire(OutputBuffers.shared()); assertEquals(0, s.pos); CborSerializer.release(s, false); } @Test - void largeBufferIsDownsizedOnRelease() { - var s = CborSerializer.acquire(); - s.buf = new byte[1024 * 1024]; + void releaseReturnsTheWorkingBufferToTheSerializersPool() { + var pool = OutputBuffers.create(); + var s = CborSerializer.acquire(pool); + s.writeString(PreludeSchemas.STRING, "x".repeat(100_000)); + byte[] buf = s.buf; + assertNotNull(buf); + assertTrue(buf.length >= 100_000); + CborSerializer.release(s, false); - var s2 = CborSerializer.acquire(); - // If we got back the same pooled instance, its buffer should have been downsized - // If we got a new instance, it has the default buffer size. Either way, not 1MB. - assertNotNull(s2); - CborSerializer.release(s2, false); + assertNull(s.buf); + assertSame(buf, pool.acquire(buf.length)); } @Test void streamingSerializerNotPooled() { - var s = new CborSerializer(OutputStream.nullOutputStream()); + var s = new CborSerializer(OutputStream.nullOutputStream(), OutputBuffers.shared()); CborSerializer.release(s, false); // Should not throw, and the streaming serializer should not be returned from acquire - var s2 = CborSerializer.acquire(); + var s2 = CborSerializer.acquire(OutputBuffers.shared()); assertNotSame(s, s2); CborSerializer.release(s2, false); } @@ -141,7 +184,7 @@ void unpairedLowSurrogate() { } private byte[] serializeString(String value) { - var s = CborSerializer.acquire(); + var s = CborSerializer.acquire(OutputBuffers.shared()); try { s.writeString(PreludeSchemas.STRING, value); return Arrays.copyOf(s.buf, s.pos); @@ -156,7 +199,7 @@ class CollectionDepthTest { @Test void deeplyNestedCollections() { int depth = 100; - var s = CborSerializer.acquire(); + var s = CborSerializer.acquire(OutputBuffers.shared()); try { for (int i = 0; i < depth; i++) { s.writeList(PreludeSchemas.DOCUMENT, null, -1, ($, inner) -> {}); @@ -169,7 +212,7 @@ void deeplyNestedCollections() { @Test void collectionDepthExceeds64() { int depth = 70; - var s = CborSerializer.acquire(); + var s = CborSerializer.acquire(OutputBuffers.shared()); try { for (int i = 0; i < depth; i++) { boolean indefinite = i % 2 == 0; @@ -680,7 +723,7 @@ public Throwable fillInStackTrace() { private static byte[] write(Consumer consumer) { try ( var stream = new ByteBufferOutputStream(); - var ser = new CborSerializer(stream)) { + var ser = new CborSerializer(stream, OutputBuffers.shared())) { try { consumer.accept(ser); } catch (StopWritingException ignored) {} diff --git a/codecs/cbor-codec/src/test/java/software/amazon/smithy/java/cbor/CborSerializerTest.java b/codecs/cbor-codec/src/test/java/software/amazon/smithy/java/cbor/CborSerializerTest.java index 8585f9750e..ac9efd512c 100644 --- a/codecs/cbor-codec/src/test/java/software/amazon/smithy/java/cbor/CborSerializerTest.java +++ b/codecs/cbor-codec/src/test/java/software/amazon/smithy/java/cbor/CborSerializerTest.java @@ -21,6 +21,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.io.ByteBufferUtils; public class CborSerializerTest { @@ -180,7 +181,7 @@ private static void assertTextStringEncoding(String value) { } System.arraycopy(utf8, 0, expected, expected.length - utf8.length, utf8.length); - var serializer = new CborSerializer(); + var serializer = new CborSerializer(null, OutputBuffers.shared()); serializer.writeString(null, value); assertArrayEquals(expected, ByteBufferUtils.getBytes(serializer.extractResult()), value); } diff --git a/codecs/codec-commons/src/main/java/software/amazon/smithy/java/codecs/commons/OutputBuffers.java b/codecs/codec-commons/src/main/java/software/amazon/smithy/java/codecs/commons/OutputBuffers.java new file mode 100644 index 0000000000..be51e80c0e --- /dev/null +++ b/codecs/codec-commons/src/main/java/software/amazon/smithy/java/codecs/commons/OutputBuffers.java @@ -0,0 +1,159 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.codecs.commons; + +import software.amazon.smithy.java.io.datastream.ByteArrayRecycler; +import software.amazon.smithy.utils.SmithyInternalApi; + +/** + * Pools output arrays in power-of-two sizes from 4 KB to 1 MB. + * + *

Codecs use {@link #shared()} by default. + * Use {@link #create()} for a pool owned by one codec. + */ +public final class OutputBuffers { + + private static final int MIN_SHIFT = 12; + private static final int MAX_SHIFT = 20; + private static final int CLASS_BUDGET = 4 << 20; // bytes of pooled arrays per size class + + private static final class Shared { + static final OutputBuffers INSTANCE = new OutputBuffers(); + } + + private final ArrayPool[] pools = new ArrayPool[MAX_SHIFT - MIN_SHIFT + 1]; + private final ByteArrayRecycler recycler = this::release; + + private OutputBuffers() { + int maxSlots = StripedPool.logicalSlots(4, Runtime.getRuntime().availableProcessors()); + for (int shift = MIN_SHIFT; shift <= MAX_SHIFT; shift++) { + int size = 1 << shift; + int slots = Math.min(maxSlots, Math.max(1, CLASS_BUDGET / size)); + pools[shift - MIN_SHIFT] = new ArrayPool(size, slots); + } + } + + int slotsFor(int capacity) { + return pools[shiftFor(capacity) - MIN_SHIFT].logicalSlotCount(); + } + + /** + * @return the JVM-wide default pool + */ + public static OutputBuffers shared() { + return Shared.INSTANCE; + } + + /** + * @return a new pool + */ + public static OutputBuffers create() { + return new OutputBuffers(); + } + + /** + * Acquires an array from the smallest pooled size that meets {@code minCapacity}. + * Sizes above 1 MB receive an exact-size allocation. + * + * @param minCapacity minimum array length + */ + @SmithyInternalApi + public byte[] acquire(int minCapacity) { + int shift = shiftFor(minCapacity); + if (shift > MAX_SHIFT) { + return new byte[minCapacity]; + } + return pools[shift - MIN_SHIFT].acquire(null); + } + + /** + * Returns an array to the pool. Arrays outside the pooled sizes are discarded. + * + * @param array array from {@link #acquire} + */ + @SmithyInternalApi + public void release(byte[] array) { + int n = array.length; + if (n < (1 << MIN_SHIFT) || n > (1 << MAX_SHIFT) || Integer.bitCount(n) != 1) { + return; + } + pools[Integer.numberOfTrailingZeros(n) - MIN_SHIFT].release(array); + } + + /** + * Replaces the array with a larger array and recycles the original. + * + * @param used bytes to copy + * @param minCapacity minimum replacement length + * @return the replacement containing the first {@code used} bytes + */ + @SmithyInternalApi + public byte[] grow(byte[] array, int used, int minCapacity) { + byte[] bigger = acquire(Math.max(minCapacity, array.length * 2)); + System.arraycopy(array, 0, bigger, 0, used); + release(array); + return bigger; + } + + /** + * @return a thread-safe recycler that returns arrays to this pool + */ + @SmithyInternalApi + public ByteArrayRecycler recycler() { + return recycler; + } + + /** + * Removes all idle arrays from the pool. + * The pool can accept arrays again, including arrays already acquired. + */ + public void clear() { + for (ArrayPool pool : pools) { + pool.clear(); + } + } + + /** + * The hint decreases by a quarter, but stays at least as large as the latest output. + * The result cannot exceed 1 MB. + */ + @SmithyInternalApi + public static int nextSizeHint(int previousHint, int used) { + int decayed = previousHint - (previousHint >>> 2); + return Math.clamp(used, decayed, 1 << MAX_SHIFT); + } + + private static int shiftFor(int minCapacity) { + if (minCapacity <= (1 << MIN_SHIFT)) { + return MIN_SHIFT; + } + return 32 - Integer.numberOfLeadingZeros(minCapacity - 1); + } + + private static final class ArrayPool extends StripedPool { + private final int size; + + ArrayPool(int size, int slots) { + super(slots); + this.size = size; + } + + @Override + protected byte[] create(Void context) { + return new byte[size]; + } + + @Override + protected boolean canPool(byte[] item) { + return item.length == size; + } + + @Override + protected boolean reset(byte[] item, Void context) { + return true; + } + } +} diff --git a/codecs/codec-commons/src/main/java/software/amazon/smithy/java/codecs/commons/StripedPool.java b/codecs/codec-commons/src/main/java/software/amazon/smithy/java/codecs/commons/StripedPool.java index 32d6c1240f..b8ff56be80 100644 --- a/codecs/codec-commons/src/main/java/software/amazon/smithy/java/codecs/commons/StripedPool.java +++ b/codecs/codec-commons/src/main/java/software/amazon/smithy/java/codecs/commons/StripedPool.java @@ -33,15 +33,29 @@ public abstract class StripedPool { private final int slotMask; private final AtomicReferenceArray pool; + /** Sizes the pool from four slots per available processor, rounded up to a power of two. */ protected StripedPool() { - this(4); + this(logicalSlots(4, N_PROCS)); } - protected StripedPool(int multiplier) { - int raw = N_PROCS * multiplier; - int logicalSlots = Integer.highestOneBit(raw - 1) << 1; - this.slotMask = logicalSlots - 1; - this.pool = new AtomicReferenceArray<>(logicalSlots << STRIDE_SHIFT); + /** + * Rounds the slot count up to a power of two, with at least one slot. + * + * @param logicalSlots requested slot count + */ + protected StripedPool(int logicalSlots) { + int slots = Math.max(1, Integer.highestOneBit(Math.max(1, logicalSlots) - 1) << 1); + this.slotMask = slots - 1; + this.pool = new AtomicReferenceArray<>(slots << STRIDE_SHIFT); + } + + final int logicalSlotCount() { + return slotMask + 1; + } + + static int logicalSlots(int multiplier, int processors) { + int raw = Math.max(1, processors * multiplier); + return Math.max(1, Integer.highestOneBit(raw - 1) << 1); } /** @@ -92,6 +106,16 @@ public final T acquire(C context) { return create(context); } + /** + * Removes all idle items from the pool. + * Acquired items remain valid and can return to the pool later. + */ + public final void clear() { + for (int i = 0; i < pool.length(); i++) { + pool.set(i, null); + } + } + public final void release(T item) { cleanup(item); if (!canPool(item)) { diff --git a/codecs/codec-commons/src/test/java/software/amazon/smithy/java/codecs/commons/OutputBuffersTest.java b/codecs/codec-commons/src/test/java/software/amazon/smithy/java/codecs/commons/OutputBuffersTest.java new file mode 100644 index 0000000000..7c3c56cf8f --- /dev/null +++ b/codecs/codec-commons/src/test/java/software/amazon/smithy/java/codecs/commons/OutputBuffersTest.java @@ -0,0 +1,159 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.codecs.commons; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.io.ByteArrayOutputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.io.datastream.DataStream; + +class OutputBuffersTest { + + private static final int MIN = 4 * 1024; + private static final int MAX = 1024 * 1024; + + @Test + void eachSizeClassIsBoundedByBytesAndByCores() { + var pool = OutputBuffers.create(); + int perCore = StripedPool.logicalSlots(4, Runtime.getRuntime().availableProcessors()); + assertThat(pool.slotsFor(MAX)).isEqualTo(Math.min(perCore, 4)); + assertThat(pool.slotsFor(MAX / 2)).isEqualTo(Math.min(perCore, 8)); + assertThat(pool.slotsFor(MAX / 4)).isEqualTo(Math.min(perCore, 16)); + assertThat(pool.slotsFor(MIN)).isEqualTo(Math.min(perCore, 1024)); + long worstCase = 0; + for (int size = MIN; size <= MAX; size <<= 1) { + worstCase += (long) pool.slotsFor(size) * size; + } + assertThat(worstCase).isLessThanOrEqualTo(9L * (4 << 20)); + } + + @Test + void acquireRoundsUpToAPooledSizeClass() { + var pool = OutputBuffers.create(); + assertThat(pool.acquire(0).length).isEqualTo(MIN); + assertThat(pool.acquire(1).length).isEqualTo(MIN); + assertThat(pool.acquire(MIN).length).isEqualTo(MIN); + assertThat(pool.acquire(MIN + 1).length).isEqualTo(MIN * 2); + assertThat(pool.acquire(100_000).length).isEqualTo(131_072); + assertThat(pool.acquire(MAX).length).isEqualTo(MAX); + } + + @Test + void arraysAboveTheLargestClassAreExactlySizedAndNotPooled() { + var pool = OutputBuffers.create(); + byte[] big = pool.acquire(MAX + 1); + assertThat(big.length).isEqualTo(MAX + 1); + pool.release(big); + assertThat(acquireSeveral(pool, MAX + 1)).noneMatch(a -> a == big); + } + + @Test + void releasedArraysComeBackFromAcquire() { + var pool = OutputBuffers.create(); + byte[] a = pool.acquire(8192); + pool.release(a); + assertThat(acquireSeveral(pool, 8192)).anyMatch(b -> b == a); + } + + @Test + void arraysThatAreNotASizeClassAreDropped() { + var pool = OutputBuffers.create(); + pool.release(new byte[100]); + pool.release(new byte[5000]); + pool.release(new byte[MIN - 1]); + pool.release(new byte[MAX * 2]); + assertThat(acquireSeveral(pool, 5000)).allMatch(a -> a.length == 8192); + } + + @Test + void growKeepsTheUsedPrefixAndRecyclesTheOldArray() { + var pool = OutputBuffers.create(); + byte[] a = pool.acquire(MIN); + for (int i = 0; i < 100; i++) { + a[i] = (byte) i; + } + byte[] grown = pool.grow(a, 100, MIN + 1); + assertThat(grown).isNotSameAs(a); + assertThat(grown.length).isGreaterThanOrEqualTo(MIN * 2); + assertThat(Arrays.copyOf(grown, 100)).isEqualTo(Arrays.copyOf(a, 100)); + assertThat(acquireSeveral(pool, MIN)).anyMatch(b -> b == a); + } + + @Test + void growHonorsAMinimumLargerThanDoubling() { + var pool = OutputBuffers.create(); + byte[] grown = pool.grow(pool.acquire(MIN), 0, 100_000); + assertThat(grown.length).isGreaterThanOrEqualTo(100_000); + } + + @Test + void sizeHintFollowsLargeBodiesAtOnceAndDecaysSlowly() { + int hint = 8192; + hint = OutputBuffers.nextSizeHint(hint, 700_000); + assertThat(hint).isEqualTo(700_000); + hint = OutputBuffers.nextSizeHint(hint, 200); + assertThat(hint).isEqualTo(525_000); + assertThat(OutputBuffers.create().acquire(hint).length).isEqualTo(MAX); + for (int i = 0; i < 12; i++) { + hint = OutputBuffers.nextSizeHint(hint, 200); + } + assertThat(hint).isLessThan(64 * 1024); + assertThat(OutputBuffers.nextSizeHint(hint, 5 * MAX)).isEqualTo(MAX); + } + + @Test + void recyclerReturnsAReleasedBodyToThePoolThatLentIt() throws Exception { + var lender = OutputBuffers.create(); + var other = OutputBuffers.create(); + byte[] a = lender.acquire(MIN); + a[0] = 42; + var body = DataStream.ofLeasedBytes(a, 0, 1, "application/octet-stream", lender.recycler()); + var out = new ByteArrayOutputStream(); + body.writeTo(out); + assertThat(out.toByteArray()).containsExactly(42); + body.release(); + assertThat(acquireSeveral(lender, MIN)).anyMatch(b -> b == a); + assertThat(acquireSeveral(other, MIN)).noneMatch(b -> b == a); + } + + @Test + void clearDropsWhatThePoolHoldsAndKeepsWorking() { + var pool = OutputBuffers.create(); + byte[] a = pool.acquire(MIN); + pool.release(a); + + pool.clear(); + + assertThat(acquireSeveral(pool, MIN)).noneMatch(b -> b == a); + byte[] later = pool.acquire(MIN); + pool.release(later); + assertThat(acquireSeveral(pool, MIN)).anyMatch(b -> b == later); + } + + @Test + void poolsAreIndependent() { + var one = OutputBuffers.create(); + var two = OutputBuffers.create(); + byte[] a = one.acquire(8192); + one.release(a); + assertThat(acquireSeveral(two, 8192)).noneMatch(b -> b == a); + assertThat(acquireSeveral(one, 8192)).anyMatch(b -> b == a); + assertThat(OutputBuffers.shared()).isSameAs(OutputBuffers.shared()).isNotSameAs(one); + } + + // Thread-based slot selection can require several acquisitions to find a specific array. + private static List acquireSeveral(OutputBuffers pool, int size) { + List result = new ArrayList<>(); + for (int i = 0; i < 8; i++) { + result.add(pool.acquire(size)); + } + return result; + } +} diff --git a/codecs/codec-commons/src/test/java/software/amazon/smithy/java/codecs/commons/StripedPoolTest.java b/codecs/codec-commons/src/test/java/software/amazon/smithy/java/codecs/commons/StripedPoolTest.java new file mode 100644 index 0000000000..b9ad836dfd --- /dev/null +++ b/codecs/codec-commons/src/test/java/software/amazon/smithy/java/codecs/commons/StripedPoolTest.java @@ -0,0 +1,104 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.codecs.commons; + +import static org.assertj.core.api.Assertions.assertThat; + +import org.junit.jupiter.api.Test; + +class StripedPoolTest { + + @Test + void slotCountIsAPowerOfTwoAndNeverZero() { + assertThat(StripedPool.logicalSlots(1, 1)).isEqualTo(1); + assertThat(StripedPool.logicalSlots(2, 1)).isEqualTo(2); + assertThat(StripedPool.logicalSlots(4, 1)).isEqualTo(4); + assertThat(StripedPool.logicalSlots(1, 2)).isEqualTo(2); + assertThat(StripedPool.logicalSlots(1, 3)).isEqualTo(4); + assertThat(StripedPool.logicalSlots(4, 10)).isEqualTo(64); + } + + @Test + void slotCountRoundsUpToAPowerOfTwoAndIsNeverZero() { + assertThat(slotPool(0).logicalSlotCount()).isEqualTo(1); + assertThat(slotPool(1).logicalSlotCount()).isEqualTo(1); + assertThat(slotPool(3).logicalSlotCount()).isEqualTo(4); + assertThat(slotPool(8).logicalSlotCount()).isEqualTo(8); + assertThat(slotPool(9).logicalSlotCount()).isEqualTo(16); + } + + private static StripedPool slotPool(int slots) { + return new StripedPool<>(slots) { + @Override + protected byte[] create(Void context) { + return new byte[8]; + } + + @Override + protected boolean canPool(byte[] item) { + return true; + } + + @Override + protected boolean reset(byte[] item, Void context) { + return true; + } + }; + } + + @Test + void clearDropsPooledItemsButNotAcquiredOnes() { + var pool = new StripedPool(8) { + @Override + protected byte[] create(Void context) { + return new byte[8]; + } + + @Override + protected boolean canPool(byte[] item) { + return true; + } + + @Override + protected boolean reset(byte[] item, Void context) { + return true; + } + }; + byte[] pooled = pool.acquire(null); + byte[] inUse = pool.acquire(null); + pool.release(pooled); + + pool.clear(); + + assertThat(pool.acquire(null)).isNotSameAs(pooled); + pool.release(inUse); + assertThat(pool.acquire(null)).isSameAs(inUse); + } + + @Test + void singleSlotPoolRoundTrips() { + var pool = new StripedPool(1) { + @Override + protected byte[] create(Void context) { + return new byte[8]; + } + + @Override + protected boolean canPool(byte[] item) { + return true; + } + + @Override + protected boolean reset(byte[] item, Void context) { + return true; + } + }; + byte[] first = pool.acquire(null); + pool.release(first); + assertThat(pool.acquire(null)).isSameAs(first); + assertThat(pool.acquire(null)).isNotSameAs(first); + } +} diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/CodegenJsonSerdeProvider.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/CodegenJsonSerdeProvider.java index 00d3bcea84..fe185fe130 100644 --- a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/CodegenJsonSerdeProvider.java +++ b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/CodegenJsonSerdeProvider.java @@ -14,6 +14,7 @@ import software.amazon.smithy.java.core.serde.MemberSubsetCodec; import software.amazon.smithy.java.core.serde.ShapeDeserializer; import software.amazon.smithy.java.core.serde.ShapeSerializer; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.java.json.smithy.SmithyGeneratedJsonSerde; final class CodegenJsonSerdeProvider implements JsonSerdeProvider { @@ -53,6 +54,21 @@ public ByteBuffer serialize(SerializableShape shape, JsonSettings settings) { return delegate.serialize(shape, settings); } + @Override + public DataStream serializeToDataStream(SerializableShape shape, JsonSettings settings, String contentType) { + if (shape instanceof SerializableStruct struct) { + DataStream result = generated.serializeToDataStream(struct, settings, null, contentType); + if (result != null) { + return result; + } + } else if (settings.strictRuntimeCodegen()) { + throw new IllegalStateException( + "Strict JSON runtime codegen requires a structure or union root, found " + + shape.getClass().getName()); + } + return delegate.serializeToDataStream(shape, settings, contentType); + } + ByteBuffer serializeMemberSubset( SerializableStruct struct, MemberSubsetCodec.MemberSubset subset, @@ -61,6 +77,15 @@ ByteBuffer serializeMemberSubset( return generated.serialize(struct, settings, subset); } + DataStream serializeMemberSubsetToDataStream( + SerializableStruct struct, + MemberSubsetCodec.MemberSubset subset, + JsonSettings settings, + String contentType + ) { + return generated.serializeToDataStream(struct, settings, subset, contentType); + } + boolean deserializeMemberSubset( Schema schema, ShapeBuilder builder, diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonCodec.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonCodec.java index d54c271310..21bbf8ba93 100644 --- a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonCodec.java +++ b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonCodec.java @@ -7,6 +7,7 @@ import java.io.OutputStream; import java.nio.ByteBuffer; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableShape; import software.amazon.smithy.java.core.schema.SerializableStruct; @@ -16,6 +17,7 @@ import software.amazon.smithy.java.core.serde.ShapeDeserializer; import software.amazon.smithy.java.core.serde.ShapeSerializer; import software.amazon.smithy.java.core.serde.TimestampFormatter; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.model.traits.JsonNameTrait; import software.amazon.smithy.model.traits.TimestampFormatTrait; @@ -56,6 +58,11 @@ public ShapeSerializer createSerializer(OutputStream sink) { return settings.provider().newSerializer(sink, settings); } + @Override + public DataStream serializeToDataStream(SerializableShape shape, String contentType) { + return settings.provider().serializeToDataStream(shape, settings, contentType); + } + @Override public ByteBuffer serialize(SerializableStruct struct, MemberSubset subset) { return settings.provider() instanceof CodegenJsonSerdeProvider generated @@ -63,6 +70,13 @@ public ByteBuffer serialize(SerializableStruct struct, MemberSubset subset) { : null; } + @Override + public DataStream serializeToDataStream(SerializableStruct struct, MemberSubset subset, String contentType) { + return settings.provider() instanceof CodegenJsonSerdeProvider generated + ? generated.serializeMemberSubsetToDataStream(struct, subset, settings, contentType) + : null; + } + @Override public boolean deserialize( Schema schema, @@ -264,6 +278,19 @@ public Builder runtimeCodegen(boolean runtimeCodegen) { return this; } + /** + * Sets the output buffer pool. The default is {@link OutputBuffers#shared()}. + * + *

Use {@link OutputBuffers#create()} for a pool owned by this codec. + * + * @param outputBuffers pool for serialization output + * @return the builder + */ + public Builder outputBuffers(OutputBuffers outputBuffers) { + settingsBuilder.outputBuffers(outputBuffers); + return this; + } + /** * Uses a custom JSON serde provider. * diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSerdeProvider.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSerdeProvider.java index 498d2672fe..3d55ada29d 100644 --- a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSerdeProvider.java +++ b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSerdeProvider.java @@ -11,6 +11,7 @@ import software.amazon.smithy.java.core.serde.ShapeDeserializer; import software.amazon.smithy.java.core.serde.ShapeSerializer; import software.amazon.smithy.java.io.ByteBufferOutputStream; +import software.amazon.smithy.java.io.datastream.DataStream; public interface JsonSerdeProvider { @@ -24,6 +25,15 @@ public interface JsonSerdeProvider { ShapeSerializer newSerializer(OutputStream sink, JsonSettings settings); + /** + * Serializes a shape into a body stream. + * Providers can return pooled output. + * See {@link software.amazon.smithy.java.core.serde.Codec#serializeToDataStream} for ownership requirements. + */ + default DataStream serializeToDataStream(SerializableShape shape, JsonSettings settings, String contentType) { + return DataStream.ofByteBuffer(serialize(shape, settings), contentType); + } + /** * Serializes a shape directly to a ByteBuffer. * diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSettings.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSettings.java index 3afda6cb6e..525058e503 100644 --- a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSettings.java +++ b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSettings.java @@ -7,6 +7,7 @@ import java.util.Objects; import java.util.ServiceLoader; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.codecs.commons.internal.codegen.RuntimeCodegenFeature; import software.amazon.smithy.java.core.serde.RuntimeCodegenMode; import software.amazon.smithy.java.core.serde.TimestampFormatter; @@ -50,6 +51,7 @@ public final class JsonSettings { private final boolean prettyPrint; private final boolean useStringForArbitraryPrecision; private final RuntimeCodegenMode runtimeCodegenMode; + private final OutputBuffers outputBuffers; private JsonSettings(Builder builder) { this.timestampResolver = builder.useTimestampFormat @@ -60,6 +62,7 @@ private JsonSettings(Builder builder) { : JsonFieldMapper.UseMemberName.INSTANCE; this.forbidUnknownUnionMembers = builder.forbidUnknownUnionMembers; this.defaultNamespace = builder.defaultNamespace; + this.outputBuffers = builder.outputBuffers; JsonSerdeProvider provider = builder.provider instanceof CodegenJsonSerdeProvider codegen ? codegen.delegate() : builder.provider; @@ -209,6 +212,15 @@ JsonSerdeProvider provider() { return provider; } + /** + * Returns the output buffer pool. Settings equality excludes this pool. + * + * @return the output buffer pool + */ + public OutputBuffers outputBuffers() { + return outputBuffers; + } + void updateBuilder(Builder builder) { builder.forbidUnknownUnionMembers(forbidUnknownUnionMembers); builder.defaultNamespace(defaultNamespace); @@ -224,6 +236,7 @@ void updateBuilder(Builder builder) { builder.prettyPrint(prettyPrint); builder.useStringForArbitraryPrecision(useStringForArbitraryPrecision); builder.runtimeCodegenMode(runtimeCodegenMode); + builder.outputBuffers(outputBuffers); } /** @@ -252,6 +265,7 @@ public static final class Builder { private boolean useStringForArbitraryPrecision = false; private RuntimeCodegenMode runtimeCodegenMode; private boolean providerOverridden; + private OutputBuffers outputBuffers = OutputBuffers.shared(); private Builder() {} @@ -264,6 +278,19 @@ public JsonSettings build() { return new JsonSettings(this); } + /** + * Sets the output buffer pool. The default is {@link OutputBuffers#shared()}. + * + *

Use {@link OutputBuffers#create()} for a pool owned by this codec. + * + * @param outputBuffers pool for serialization output + * @return the builder + */ + public Builder outputBuffers(OutputBuffers outputBuffers) { + this.outputBuffers = Objects.requireNonNull(outputBuffers, "outputBuffers"); + return this; + } + /** * Whether to use the jsonName trait or just the member name. * diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/JsonCodegenWriter.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/JsonCodegenWriter.java index 8c41963976..b0c8cae15f 100644 --- a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/JsonCodegenWriter.java +++ b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/JsonCodegenWriter.java @@ -13,28 +13,34 @@ import java.util.Arrays; import software.amazon.smithy.java.codecs.commons.CompactStringAccess; import software.amazon.smithy.java.codecs.commons.NumberCodec; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.codecs.commons.StripedPool; import software.amazon.smithy.java.core.schema.PreludeSchemas; import software.amazon.smithy.java.core.serde.SerializationException; import software.amazon.smithy.java.core.serde.document.Document; import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.java.json.JsonSettings; final class JsonCodegenWriter { private static final int DEFAULT_CAPACITY = 8192; - private static final int MAX_CACHEABLE_CAPACITY = DEFAULT_CAPACITY * 16; + private static final int MAX_CACHEABLE_SCRATCH = DEFAULT_CAPACITY * 16; private static final int MAX_DEPTH = 64; private static final StripedPool POOL = new WriterPool(); - private byte[] bytes = new byte[DEFAULT_CAPACITY]; + private byte[] bytes; private int position; + private int sizeHint = DEFAULT_CAPACITY; private int depth; private final boolean[] needsComma = new boolean[MAX_DEPTH]; private JsonSettings settings; private OutputStream sink; + private OutputBuffers pool; private JsonCodegenWriter(JsonSettings settings) { this.settings = settings; + this.pool = settings.outputBuffers(); + this.bytes = pool.acquire(DEFAULT_CAPACITY); } static JsonCodegenWriter acquire(JsonSettings settings) { @@ -50,9 +56,26 @@ void sink(OutputStream sink) { } ByteBuffer detach() { + sizeHint = OutputBuffers.nextSizeHint(sizeHint, position); return ByteBuffer.wrap(Arrays.copyOf(bytes, position)); } + DataStream detachLeased(String contentType) { + byte[] detached = bytes; + int length = position; + sizeHint = OutputBuffers.nextSizeHint(sizeHint, length); + bytes = null; + position = 0; + return DataStream.ofLeasedBytes(detached, 0, length, contentType, pool.recycler()); + } + + private void releaseBuffer() { + if (bytes != null) { + pool.release(bytes); + bytes = null; + } + } + int size() { return position; } @@ -358,7 +381,7 @@ private void ensure(int needed) { } private void grow(int needed) { - bytes = Arrays.copyOf(bytes, Math.max(bytes.length * 2, position + needed)); + bytes = pool.grow(bytes, position, position + needed); } private static final class WriterPool extends StripedPool { @@ -372,19 +395,20 @@ protected void cleanup(JsonCodegenWriter writer) { writer.sink = null; writer.position = 0; writer.depth = 0; + writer.releaseBuffer(); + // Clear client settings and pool references because idle writers remain in a JVM-wide pool. + writer.settings = null; + writer.pool = null; } @Override protected boolean canPool(JsonCodegenWriter writer) { - return writer.bytes != null; + return true; } @Override protected void prepareForPool(JsonCodegenWriter writer) { - if (writer.bytes.length > MAX_CACHEABLE_CAPACITY) { - writer.bytes = new byte[DEFAULT_CAPACITY]; - } - if (writer.base64Scratch != null && writer.base64Scratch.length > MAX_CACHEABLE_CAPACITY) { + if (writer.base64Scratch != null && writer.base64Scratch.length > MAX_CACHEABLE_SCRATCH) { writer.base64Scratch = null; } } @@ -392,6 +416,8 @@ protected void prepareForPool(JsonCodegenWriter writer) { @Override protected boolean reset(JsonCodegenWriter writer, JsonSettings settings) { writer.settings = settings; + writer.pool = settings.outputBuffers(); + writer.bytes = writer.pool.acquire(writer.sizeHint); return true; } } diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyGeneratedJsonSerde.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyGeneratedJsonSerde.java index 39bfa9dcad..6d50e58418 100644 --- a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyGeneratedJsonSerde.java +++ b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyGeneratedJsonSerde.java @@ -18,6 +18,7 @@ import software.amazon.smithy.java.core.schema.SerializableStruct; import software.amazon.smithy.java.core.schema.ShapeBuilder; import software.amazon.smithy.java.core.serde.MemberSubsetCodec; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.java.json.JsonSettings; import software.amazon.smithy.utils.SmithyInternalApi; @@ -68,6 +69,39 @@ public ByteBuffer serialize( } } + /** + * Returns pooled serialization output, or null if generated code cannot handle the value. + * + * @see #serialize(SerializableStruct, JsonSettings, MemberSubsetCodec.MemberSubset) + */ + public DataStream serializeToDataStream( + SerializableStruct value, + JsonSettings settings, + MemberSubsetCodec.MemberSubset subset, + String contentType + ) { + if (settings.prettyPrint()) { + rejectStrictFallback(settings, "pretty-printed JSON is not supported"); + return null; + } + if (!generatable(value)) { + rejectStrictFallback(settings, "the value is not the generated class for " + value.schema().id()); + return null; + } + GeneratedJsonCodec codec = codec(value.schema(), settings, subset); + if (codec == null) { + rejectStrictFallback(settings, "no generated codec is available for " + value.schema().id()); + return null; + } + JsonCodegenWriter writer = JsonCodegenWriter.acquire(settings); + try { + codec.write(value, writer); + return writer.detachLeased(contentType); + } finally { + JsonCodegenWriter.release(writer); + } + } + public boolean serializeTo( SerializableStruct value, OutputStream sink, diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyJsonSerdeProvider.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyJsonSerdeProvider.java index 1bfd433190..23dcc0f372 100644 --- a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyJsonSerdeProvider.java +++ b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyJsonSerdeProvider.java @@ -10,6 +10,7 @@ import software.amazon.smithy.java.core.schema.SerializableShape; import software.amazon.smithy.java.core.serde.ShapeDeserializer; import software.amazon.smithy.java.core.serde.ShapeSerializer; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.java.json.JsonSerdeProvider; import software.amazon.smithy.java.json.JsonSettings; import software.amazon.smithy.java.json.jackson.JacksonJsonSerdeProvider; @@ -52,6 +53,24 @@ public ByteBuffer serialize(SerializableShape shape, JsonSettings settings) { } } + @Override + public DataStream serializeToDataStream(SerializableShape shape, JsonSettings settings, String contentType) { + if (settings.prettyPrint()) { + return DataStream.ofByteBuffer(getJacksonFallback().serialize(shape, settings), contentType); + } + var serializer = SmithyJsonSerializer.acquire(settings); + boolean exception = false; + try { + shape.serialize(serializer); + return serializer.detach(contentType); + } catch (Exception t) { + exception = true; + throw t; + } finally { + SmithyJsonSerializer.release(serializer, exception); + } + } + @Override public ShapeSerializer newSerializer(OutputStream sink, JsonSettings settings) { if (settings.prettyPrint()) { diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyJsonSerializer.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyJsonSerializer.java index 3c0a0facaf..d18154d738 100644 --- a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyJsonSerializer.java +++ b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/smithy/SmithyJsonSerializer.java @@ -13,6 +13,7 @@ import java.util.Arrays; import java.util.function.BiConsumer; import software.amazon.smithy.java.codecs.commons.NumberCodec; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.codecs.commons.StripedPool; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableStruct; @@ -23,6 +24,7 @@ import software.amazon.smithy.java.core.serde.TimestampFormatter; import software.amazon.smithy.java.core.serde.document.Document; import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.java.json.JsonFieldMapper; import software.amazon.smithy.java.json.JsonSettings; import software.amazon.smithy.model.shapes.ShapeType; @@ -37,16 +39,18 @@ final class SmithyJsonSerializer implements ShapeSerializer { private static final int MAX_DEPTH = 64; private static final int DEFAULT_BUF_SIZE = 8192; - private static final int MAX_CACHEABLE_BUF = DEFAULT_BUF_SIZE * 16; + private static final int MAX_CACHEABLE_SCRATCH = DEFAULT_BUF_SIZE * 16; private static final StripedPool POOL = new JsonStripedPool(); private byte[] buf; private byte[] base64Scratch; private int pos; + private int sizeHint = DEFAULT_BUF_SIZE; private final OutputStream sink; - private final JsonSettings settings; - private final boolean useJsonName; + private JsonSettings settings; + private OutputBuffers pool; + private boolean useJsonName; // Nesting state for comma insertion private int depth; @@ -65,7 +69,8 @@ final class SmithyJsonSerializer implements ShapeSerializer { this.sink = sink; this.settings = settings; this.useJsonName = settings.fieldMapper() instanceof JsonFieldMapper.UseJsonNameTrait; - this.buf = new byte[DEFAULT_BUF_SIZE]; + this.pool = settings.outputBuffers(); + this.buf = pool.acquire(DEFAULT_BUF_SIZE); this.pos = 0; this.depth = 0; } @@ -78,7 +83,8 @@ final class SmithyJsonSerializer implements ShapeSerializer { this.sink = null; this.settings = settings; this.useJsonName = settings.fieldMapper() instanceof JsonFieldMapper.UseJsonNameTrait; - this.buf = new byte[DEFAULT_BUF_SIZE]; + this.pool = settings.outputBuffers(); + this.buf = pool.acquire(DEFAULT_BUF_SIZE); this.pos = 0; this.depth = 0; } @@ -95,9 +101,26 @@ static void release(SmithyJsonSerializer serializer, boolean exception) { } ByteBuffer extractResult() { + sizeHint = OutputBuffers.nextSizeHint(sizeHint, pos); return ByteBuffer.wrap(Arrays.copyOf(buf, pos)); } + DataStream detach(String contentType) { + byte[] detached = buf; + int length = pos; + sizeHint = OutputBuffers.nextSizeHint(sizeHint, length); + buf = null; + pos = 0; + return DataStream.ofLeasedBytes(detached, 0, length, contentType, pool.recycler()); + } + + private void releaseBuffer() { + if (buf != null) { + pool.release(buf); + buf = null; + } + } + private void ensureCapacity(int needed) { if (pos + needed > buf.length) { grow(needed); @@ -106,7 +129,7 @@ private void ensureCapacity(int needed) { // Separate cold grow method helps JIT inline ensureCapacity's fast path private void grow(int needed) { - buf = Arrays.copyOf(buf, Math.max(buf.length * 2, pos + needed)); + buf = pool.grow(buf, pos, pos + needed); } /** Writes with the exact ASCII reservation, widening only when needed. */ @@ -143,9 +166,10 @@ public void close() { sink.write(buf, 0, pos); pos = 0; } - buf = null; } catch (Exception e) { throw new SerializationException(e); + } finally { + releaseBuffer(); } } @@ -413,29 +437,35 @@ protected SmithyJsonSerializer create(JsonSettings settings) { return new SmithyJsonSerializer(settings); } + @Override + protected void cleanup(SmithyJsonSerializer s) { + s.releaseBuffer(); + // Clear client settings and pool references because idle serializers remain in a JVM-wide pool. + s.settings = null; + s.pool = null; + } + @Override protected boolean canPool(SmithyJsonSerializer s) { - return s.buf != null; + return true; } @Override protected void prepareForPool(SmithyJsonSerializer s) { - if (s.buf.length > MAX_CACHEABLE_BUF) { - s.buf = new byte[DEFAULT_BUF_SIZE]; - } - if (s.base64Scratch != null && s.base64Scratch.length > MAX_CACHEABLE_BUF) { + if (s.base64Scratch != null && s.base64Scratch.length > MAX_CACHEABLE_SCRATCH) { s.base64Scratch = null; } } @Override protected boolean reset(SmithyJsonSerializer s, JsonSettings settings) { - if (!s.settings.equals(settings)) { - return false; - } + s.settings = settings; + s.useJsonName = settings.fieldMapper() instanceof JsonFieldMapper.UseJsonNameTrait; s.pos = 0; s.depth = 0; s.currentFieldNameTable = null; + s.pool = settings.outputBuffers(); + s.buf = s.pool.acquire(s.sizeHint); return true; } } diff --git a/codecs/json-codec/src/test/java/software/amazon/smithy/java/json/JsonSerializerTest.java b/codecs/json-codec/src/test/java/software/amazon/smithy/java/json/JsonSerializerTest.java index a8037b871d..fc5bc71865 100644 --- a/codecs/json-codec/src/test/java/software/amazon/smithy/java/json/JsonSerializerTest.java +++ b/codecs/json-codec/src/test/java/software/amazon/smithy/java/json/JsonSerializerTest.java @@ -12,12 +12,14 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.OutputStream; +import java.lang.ref.WeakReference; import java.math.BigDecimal; import java.math.BigInteger; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.time.Instant; import java.util.ArrayList; +import java.util.Arrays; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -26,11 +28,13 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.schema.PreludeSchemas; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableStruct; import software.amazon.smithy.java.core.serde.ShapeSerializer; import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.java.json.smithy.SmithyJsonSerdeProvider; import software.amazon.smithy.model.shapes.ShapeId; import software.amazon.smithy.model.traits.TimestampFormatTrait; @@ -398,6 +402,55 @@ public T getMemberValue(Schema member) { } } + @Test + public void leasedBodiesReturnToTheCodecsOwnPool() throws IOException { + var pool = OutputBuffers.create(); + try (var codec = JsonCodec.builder().runtimeCodegen(false).outputBuffers(pool).build()) { + var body = Assertions.assertInstanceOf(PooledDataStream.class, + codec.serializeToDataStream(new NestedStruct(), "application/json")); + var out = new ByteArrayOutputStream(); + body.writeTo(out); + String json = out.toString(StandardCharsets.UTF_8); + assertThat(json, equalTo("{\"number\":10}")); + + body.release(); + Assertions.assertTrue(poolHoldsArrayStartingWith(pool, json), "array returned to the codec's pool"); + } + } + + @Test + public void idleSerializersDoNotPinACodecsPool() throws Exception { + WeakReference pool = serializeWithADedicatedPoolThenDropIt(); + for (int i = 0; i < 50 && pool.get() != null; i++) { + System.gc(); + Thread.sleep(20); + } + Assertions.assertNull(pool.get(), "a codec's pool must be collectible once the codec is gone"); + } + + private static WeakReference serializeWithADedicatedPoolThenDropIt() { + var pool = OutputBuffers.create(); + try (var codec = JsonCodec.builder().runtimeCodegen(false).outputBuffers(pool).build()) { + ((PooledDataStream) codec.serializeToDataStream(new NestedStruct(), "application/json")).release(); + codec.serialize(new NestedStruct()); + } + return new WeakReference<>(pool); + } + + private static boolean poolHoldsArrayStartingWith(OutputBuffers pool, String prefix) { + byte[] expected = prefix.getBytes(StandardCharsets.UTF_8); + for (int shift = 12; shift <= 20; shift++) { + for (int i = 0; i < 8; i++) { + byte[] candidate = pool.acquire(1 << shift); + if (candidate.length >= expected.length + && Arrays.equals(candidate, 0, expected.length, expected, 0, expected.length)) { + return true; + } + } + } + return false; + } + private static final class NestedStruct implements SerializableStruct { @Override public Schema schema() { diff --git a/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/LazyXmlSerializer.java b/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/LazyXmlSerializer.java index cbc8032323..0bd02ab80b 100644 --- a/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/LazyXmlSerializer.java +++ b/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/LazyXmlSerializer.java @@ -11,6 +11,7 @@ import java.nio.ByteBuffer; import java.time.Instant; import java.util.function.BiConsumer; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableStruct; import software.amazon.smithy.java.core.serde.MapSerializer; @@ -27,21 +28,31 @@ final class LazyXmlSerializer implements ShapeSerializer { private final XmlNamespaceTrait defaultNamespace; private final XmlInfo xmlInfo; private final OutputStream sink; + private final OutputBuffers pool; private SmithyXmlSerializer delegate; + private boolean closed; - LazyXmlSerializer(XmlNamespaceTrait defaultNamespace, XmlInfo xmlInfo, OutputStream sink) { + LazyXmlSerializer(XmlNamespaceTrait defaultNamespace, XmlInfo xmlInfo, OutputStream sink, OutputBuffers pool) { this.defaultNamespace = defaultNamespace; this.xmlInfo = xmlInfo; this.sink = sink; + this.pool = pool; } private SmithyXmlSerializer delegate() { + if (closed) { + throw new IllegalStateException("Serializer is closed"); + } if (delegate == null) { - delegate = SmithyXmlSerializer.acquire(defaultNamespace, xmlInfo, sink); + delegate = SmithyXmlSerializer.acquire(defaultNamespace, xmlInfo, sink, pool); } return delegate; } + SmithyXmlSerializer delegateForTest() { + return delegate; + } + @Override public void writeStruct(Schema schema, SerializableStruct struct) { delegate().writeStruct(schema, struct); @@ -136,8 +147,12 @@ public void flush() { @Override public void close() { - if (delegate != null) { - delegate.close(); + // Return the delegate to its pool only once. A double release can give the same instance to concurrent callers. + closed = true; + var d = delegate; + if (d != null) { + delegate = null; + d.close(); } } } diff --git a/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/SmithyXmlSerializer.java b/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/SmithyXmlSerializer.java index ecfaa1c0e5..76d0b03439 100644 --- a/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/SmithyXmlSerializer.java +++ b/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/SmithyXmlSerializer.java @@ -12,9 +12,9 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.time.Instant; -import java.util.Arrays; import java.util.function.BiConsumer; import software.amazon.smithy.java.codecs.commons.NumberCodec; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; import software.amazon.smithy.java.codecs.commons.StripedPool; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableStruct; @@ -27,6 +27,7 @@ import software.amazon.smithy.java.core.serde.TimestampFormatter; import software.amazon.smithy.java.core.serde.document.Document; import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.model.traits.TimestampFormatTrait; import software.amazon.smithy.model.traits.XmlNamespaceTrait; @@ -34,16 +35,18 @@ final class SmithyXmlSerializer extends InterceptingSerializer { private static final TimestampFormatTrait.Format DEFAULT_FORMAT = TimestampFormatTrait.Format.DATE_TIME; private static final int DEFAULT_BUF_SIZE = 4096; - private static final int MAX_CACHEABLE_BUF = DEFAULT_BUF_SIZE * 4; + private static final int MAX_CACHEABLE_SCRATCH = DEFAULT_BUF_SIZE * 4; - record AcquireContext(XmlNamespaceTrait defaultNamespace, XmlInfo xmlInfo, OutputStream sink) {} + record AcquireContext(XmlNamespaceTrait defaultNamespace, XmlInfo xmlInfo, OutputStream sink, OutputBuffers pool) {} private static final StripedPool POOL = new XmlStripedPool(); private byte[] buf; private byte[] base64Scratch; private int pos; + private int sizeHint = DEFAULT_BUF_SIZE; private OutputStream sink; + private OutputBuffers pool; private XmlNamespaceTrait defaultNamespace; private XmlInfo xmlInfo; @@ -59,21 +62,43 @@ record AcquireContext(XmlNamespaceTrait defaultNamespace, XmlInfo xmlInfo, Outpu // writeStruct will close it after writing attributes; scalar writes close it immediately. private boolean pendingClose; - private SmithyXmlSerializer(XmlNamespaceTrait defaultNamespace, XmlInfo xmlInfo) { + private SmithyXmlSerializer(XmlNamespaceTrait defaultNamespace, XmlInfo xmlInfo, OutputBuffers pool) { this.defaultNamespace = defaultNamespace; this.xmlInfo = xmlInfo; - this.buf = new byte[DEFAULT_BUF_SIZE]; + this.pool = pool; + this.buf = pool.acquire(DEFAULT_BUF_SIZE); this.pos = 0; } - static SmithyXmlSerializer acquire(XmlNamespaceTrait defaultNamespace, XmlInfo xmlInfo, OutputStream sink) { - return POOL.acquire(new AcquireContext(defaultNamespace, xmlInfo, sink)); + static SmithyXmlSerializer acquire( + XmlNamespaceTrait defaultNamespace, + XmlInfo xmlInfo, + OutputStream sink, + OutputBuffers pool + ) { + return POOL.acquire(new AcquireContext(defaultNamespace, xmlInfo, sink, pool)); } static void release(SmithyXmlSerializer serializer) { POOL.release(serializer); } + DataStream detach(String contentType) { + byte[] detached = buf; + int length = pos; + sizeHint = OutputBuffers.nextSizeHint(sizeHint, length); + buf = null; + pos = 0; + return DataStream.ofLeasedBytes(detached, 0, length, contentType, pool.recycler()); + } + + private void releaseBuffer() { + if (buf != null) { + pool.release(buf); + buf = null; + } + } + @Override public void flush() { try { @@ -108,7 +133,7 @@ private void ensureCapacity(int needed) { } private void grow(int needed) { - buf = Arrays.copyOf(buf, Math.max(buf.length * 2, pos + needed)); + buf = pool.grow(buf, pos, pos + needed); } private byte[] base64Scratch(int dataLen) { @@ -256,7 +281,7 @@ private static String formatTimestamp(Schema schema, Instant value) { private static class XmlStripedPool extends StripedPool { @Override protected SmithyXmlSerializer create(AcquireContext ctx) { - var s = new SmithyXmlSerializer(ctx.defaultNamespace(), ctx.xmlInfo()); + var s = new SmithyXmlSerializer(ctx.defaultNamespace(), ctx.xmlInfo(), ctx.pool()); s.sink = ctx.sink(); return s; } @@ -264,19 +289,19 @@ protected SmithyXmlSerializer create(AcquireContext ctx) { @Override protected void cleanup(SmithyXmlSerializer s) { s.sink = null; + s.releaseBuffer(); + // Clear the client pool reference because idle serializers remain in a JVM-wide pool. + s.pool = null; } @Override protected boolean canPool(SmithyXmlSerializer s) { - return s.buf != null; + return true; } @Override protected void prepareForPool(SmithyXmlSerializer s) { - if (s.buf.length > MAX_CACHEABLE_BUF) { - s.buf = new byte[DEFAULT_BUF_SIZE]; - } - if (s.base64Scratch != null && s.base64Scratch.length > MAX_CACHEABLE_BUF) { + if (s.base64Scratch != null && s.base64Scratch.length > MAX_CACHEABLE_SCRATCH) { s.base64Scratch = null; } } @@ -284,12 +309,16 @@ protected void prepareForPool(SmithyXmlSerializer s) { @Override protected boolean reset(SmithyXmlSerializer s, AcquireContext ctx) { s.pos = 0; + // A failed write can leave a pending '>'. Clear it before the next document. + s.pendingClose = false; s.sink = ctx.sink(); s.defaultNamespace = ctx.defaultNamespace(); s.xmlInfo = ctx.xmlInfo(); s.currentNameTable = null; s.currentIsAttributeTable = null; s.currentIsFlattenedTable = null; + s.pool = ctx.pool(); + s.buf = s.pool.acquire(s.sizeHint); return true; } } diff --git a/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/XmlCodec.java b/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/XmlCodec.java index 2e230ead01..42760f1861 100644 --- a/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/XmlCodec.java +++ b/codecs/xml-codec/src/main/java/software/amazon/smithy/java/xml/XmlCodec.java @@ -8,15 +8,19 @@ import java.io.OutputStream; import java.nio.ByteBuffer; import java.util.List; +import java.util.Objects; import javax.xml.stream.XMLEventFactory; import javax.xml.stream.XMLInputFactory; import javax.xml.stream.XMLOutputFactory; import javax.xml.stream.XMLStreamException; +import software.amazon.smithy.java.codecs.commons.OutputBuffers; +import software.amazon.smithy.java.core.schema.SerializableShape; import software.amazon.smithy.java.core.serde.Codec; import software.amazon.smithy.java.core.serde.SerializationException; import software.amazon.smithy.java.core.serde.ShapeDeserializer; import software.amazon.smithy.java.core.serde.ShapeSerializer; import software.amazon.smithy.java.io.ByteBufferUtils; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.model.traits.XmlNamespaceTrait; /** @@ -37,11 +41,13 @@ public final class XmlCodec implements Codec { private final XmlNamespaceTrait defaultNamespace; private final boolean useNative; private final boolean strictRootElement; + private final OutputBuffers outputBuffers; private XmlCodec(Builder builder) { this.wrapperElements = builder.wrapperElements; this.defaultNamespace = builder.defaultNamespace; this.strictRootElement = builder.strictRootElement; + this.outputBuffers = builder.outputBuffers; this.useNative = builder.useNative != null ? builder.useNative : USE_SMITHY_NATIVE; if (!useNative) { initStax(); @@ -67,10 +73,28 @@ public static Builder builder() { return new Builder(); } + public OutputBuffers outputBuffers() { + return outputBuffers; + } + + @Override + public DataStream serializeToDataStream(SerializableShape shape, String contentType) { + if (!useNative) { + return Codec.super.serializeToDataStream(shape, contentType); + } + var serializer = SmithyXmlSerializer.acquire(defaultNamespace, xmlInfo, null, outputBuffers); + try { + shape.serialize(serializer); + return serializer.detach(contentType); + } finally { + SmithyXmlSerializer.release(serializer); + } + } + @Override public ShapeSerializer createSerializer(OutputStream sink) { if (useNative) { - return new LazyXmlSerializer(defaultNamespace, xmlInfo, sink); + return new LazyXmlSerializer(defaultNamespace, xmlInfo, sink, outputBuffers); } try { return new XmlSerializer(xmlOutputFactory.createXMLStreamWriter(sink), xmlInfo, defaultNamespace); @@ -126,6 +150,7 @@ public static final class Builder { private XmlNamespaceTrait defaultNamespace; private Boolean useNative; private boolean strictRootElement = true; + private OutputBuffers outputBuffers = OutputBuffers.shared(); private Builder() {} @@ -175,6 +200,19 @@ public Builder strictRootElement(boolean strictRootElement) { return this; } + /** + * Sets the output buffer pool. The default is {@link OutputBuffers#shared()}. + * + *

Use {@link OutputBuffers#create()} for a pool owned by this codec. + * + * @param outputBuffers pool for serialization output + * @return the builder + */ + public Builder outputBuffers(OutputBuffers outputBuffers) { + this.outputBuffers = Objects.requireNonNull(outputBuffers, "outputBuffers"); + return this; + } + /** * Override the native provider selection for testing. When set to true, the Smithy * implementation is used regardless of system property. When false, the StAX diff --git a/codecs/xml-codec/src/test/java/software/amazon/smithy/java/xml/XmlCodecTest.java b/codecs/xml-codec/src/test/java/software/amazon/smithy/java/xml/XmlCodecTest.java index 0ca53cbb4c..61cb4b735a 100644 --- a/codecs/xml-codec/src/test/java/software/amazon/smithy/java/xml/XmlCodecTest.java +++ b/codecs/xml-codec/src/test/java/software/amazon/smithy/java/xml/XmlCodecTest.java @@ -14,6 +14,7 @@ import java.io.ByteArrayOutputStream; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.time.Instant; import java.util.ArrayList; import java.util.LinkedHashMap; @@ -170,6 +171,65 @@ public void nativeAndStaxProduceSameResultForBlobMap() { } } + @Test + public void failedSerializationDoesNotPoisonThePooledSerializer() { + var poison = new SerializableStruct() { + @Override + public Schema schema() { + return TestPojo.SCHEMA; + } + + @Override + public void serializeMembers(ShapeSerializer serializer) { + throw new IllegalStateException("boom"); + } + + @Override + public T getMemberValue(Schema member) { + throw new UnsupportedOperationException(); + } + }; + var builder = new TestPojo.Builder(); + builder.name = "Hello"; + var good = builder.build(); + + try (var codec = XmlCodec.builder().useNative(true).build()) { + assertThrows(IllegalStateException.class, () -> codec.serializeToDataStream(poison, "application/xml")); + var leased = codec.serializeToDataStream(good, "application/xml"); + assertThat(StandardCharsets.UTF_8.decode(leased.asByteBuffer()).toString(), + equalTo("Hello")); + + assertThrows(IllegalStateException.class, () -> codec.serialize(poison)); + assertThat(codec.serializeToString(good), equalTo("Hello")); + } + } + + @Test + public void closingTheLazySerializerTwiceReleasesItsDelegateOnce() throws Exception { + try (var codec = XmlCodec.builder().useNative(true).build()) { + var builder = new TestPojo.Builder(); + builder.name = "Hello"; + var pojo = builder.build(); + + var out = new ByteArrayOutputStream(); + var serializer = codec.createSerializer(out); + pojo.serialize(serializer); + serializer.close(); + serializer.close(); + assertThat(out.toString(StandardCharsets.UTF_8), + equalTo("Hello")); + assertThrows(IllegalStateException.class, () -> pojo.serialize(serializer)); + + var first = (LazyXmlSerializer) codec.createSerializer(new ByteArrayOutputStream()); + var second = (LazyXmlSerializer) codec.createSerializer(new ByteArrayOutputStream()); + pojo.serialize(first); + pojo.serialize(second); + assertThat(first.delegateForTest() == second.delegateForTest(), equalTo(false)); + first.close(); + second.close(); + } + } + private static final class MapPojo implements SerializableStruct { private static final ShapeId ID = ShapeId.from("smithy.example#MapPojo"); diff --git a/core/src/main/java/software/amazon/smithy/java/core/serde/Codec.java b/core/src/main/java/software/amazon/smithy/java/core/serde/Codec.java index deca83fd23..3f48212090 100644 --- a/core/src/main/java/software/amazon/smithy/java/core/serde/Codec.java +++ b/core/src/main/java/software/amazon/smithy/java/core/serde/Codec.java @@ -12,6 +12,7 @@ import software.amazon.smithy.java.core.schema.SerializableShape; import software.amazon.smithy.java.core.schema.ShapeBuilder; import software.amazon.smithy.java.io.ByteBufferOutputStream; +import software.amazon.smithy.java.io.datastream.DataStream; /** * Generic shape serialization and deserialization. @@ -70,6 +71,21 @@ default ByteBuffer serialize(SerializableShape shape) { return baos.toByteBuffer(); } + /** + * Serializes a shape into a request or response body. + * + *

Codecs can return a {@link software.amazon.smithy.java.io.datastream.PooledDataStream}. + * The caller must release a pooled stream after all reads complete. + * The client pipeline releases request bodies after the call completes, including retries. + * + * @param shape shape to serialize + * @param contentType content type, or null + * @return a replayable stream containing the serialized shape + */ + default DataStream serializeToDataStream(SerializableShape shape, String contentType) { + return DataStream.ofByteBuffer(serialize(shape), contentType); + } + /** * Create a deserializer from this Codec that deserializes a shape from the source. * diff --git a/core/src/main/java/software/amazon/smithy/java/core/serde/MemberSubsetCodec.java b/core/src/main/java/software/amazon/smithy/java/core/serde/MemberSubsetCodec.java index 6cc7e1833a..fbdcc2ccbb 100644 --- a/core/src/main/java/software/amazon/smithy/java/core/serde/MemberSubsetCodec.java +++ b/core/src/main/java/software/amazon/smithy/java/core/serde/MemberSubsetCodec.java @@ -9,6 +9,7 @@ import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.core.schema.SerializableStruct; import software.amazon.smithy.java.core.schema.ShapeBuilder; +import software.amazon.smithy.java.io.datastream.DataStream; import software.amazon.smithy.utils.SmithyInternalApi; /** Optional {@link Codec} capability for coding selected structure members. */ @@ -30,6 +31,16 @@ default boolean deserialize( /** Returns selected members serialized, or null when unsupported. */ ByteBuffer serialize(SerializableStruct struct, MemberSubset subset); + /** + * Serializes selected members into a body stream. Returns null if unsupported. + * + * @see Codec#serializeToDataStream + */ + default DataStream serializeToDataStream(SerializableStruct struct, MemberSubset subset, String contentType) { + ByteBuffer bytes = serialize(struct, subset); + return bytes == null ? null : DataStream.ofByteBuffer(bytes, contentType); + } + /** * Deserializes only the members {@code subset} includes directly into {@code builder}. * diff --git a/http/http-binding/src/main/java/software/amazon/smithy/java/http/binding/HttpBindingSerializer.java b/http/http-binding/src/main/java/software/amazon/smithy/java/http/binding/HttpBindingSerializer.java index a297b42fb1..b17f8fa554 100644 --- a/http/http-binding/src/main/java/software/amazon/smithy/java/http/binding/HttpBindingSerializer.java +++ b/http/http-binding/src/main/java/software/amazon/smithy/java/http/binding/HttpBindingSerializer.java @@ -76,8 +76,7 @@ final class HttpBindingSerializer extends SpecificShapeSerializer implements Sha // Initialized in writeStruct. Empty until writeStruct is called. private Set namesFromHttpHeader = Set.of(); - // Body bytes from the pooled codec.serialize() path. - private ByteBuffer shapeBodyBuffer; + private DataStream shapeBody; private DataStream httpPayload; private EventStream eventStream; private Schema payloadMember; @@ -150,7 +149,7 @@ public void writeStruct(Schema schema, SerializableStruct struct) { writeBody = allowEmptyStructPayload || resp.writeBody(omitEmptyPayload); if (writeBody && !hasBodyMembers) { // No body members: fixed empty struct bytes, cached on the binding. - shapeBodyBuffer = resp.emptyBody(payloadCodec, schema); + shapeBody = DataStream.ofByteBuffer(resp.emptyBody(payloadCodec, schema), payloadMediaType); } } else { var req = bindings.request(); @@ -165,7 +164,7 @@ public void writeStruct(Schema schema, SerializableStruct struct) { pathSerializer = req.pathSerializer(operationBinding.httpTrait(), schema); writeBody = allowEmptyStructPayload || req.writeBody(omitEmptyPayload); if (writeBody && !hasBodyMembers) { - shapeBodyBuffer = req.emptyBody(payloadCodec, schema); + shapeBody = DataStream.ofByteBuffer(req.emptyBody(payloadCodec, schema), payloadMediaType); } } @@ -183,9 +182,8 @@ public void writeStruct(Schema schema, SerializableStruct struct) { if (writeBody) { if (hasBodyMembers) { - shapeBodyBuffer = serializeBody(schema, struct); + shapeBody = serializeBody(schema, struct); } - // Empty-body case (shapeBodyBuffer was set above from the cached empty bytes). headers.setHeader(HeaderName.CONTENT_TYPE, payloadMediaType); } @@ -295,9 +293,24 @@ void bindNull() { throw new IllegalStateException("Unexpected null value written for an HTTP binding"); } - private ByteBuffer serializeBody(Schema schema, SerializableStruct struct) { + private DataStream serializeBody(Schema schema, SerializableStruct struct) { + if (isResponse) { + // Server response bodies have no lease owner. Return an unpooled copy. + return DataStream.ofByteBuffer(serializeBodyBuffer(schema, struct), payloadMediaType); + } + if (subsetCodec != null && struct.schema() == schema) { + DataStream body = + subsetCodec.serializeToDataStream(struct, BodyMemberSubset.of(false), payloadMediaType); + if (body != null) { + return body; + } + } + return payloadCodec.serializeToDataStream(new StructBodyProxy(struct, activeBindings), payloadMediaType); + } + + private ByteBuffer serializeBodyBuffer(Schema schema, SerializableStruct struct) { if (subsetCodec != null && struct.schema() == schema) { - ByteBuffer body = subsetCodec.serialize(struct, BodyMemberSubset.of(isResponse)); + ByteBuffer body = subsetCodec.serialize(struct, BodyMemberSubset.of(true)); if (body != null) { return body; } @@ -365,14 +378,14 @@ private QueryStringBuilder queryStringParams() { } boolean hasBody() { - return shapeBodyBuffer != null || httpPayload != null; + return shapeBody != null || httpPayload != null; } DataStream getBody() { if (httpPayload != null) { return httpPayload; - } else if (shapeBodyBuffer != null) { - return DataStream.ofByteBuffer(shapeBodyBuffer, payloadMediaType); + } else if (shapeBody != null) { + return shapeBody; } else { return DataStream.ofEmpty(); } diff --git a/http/http-binding/src/test/java/software/amazon/smithy/java/http/binding/HttpBindingWriterTest.java b/http/http-binding/src/test/java/software/amazon/smithy/java/http/binding/HttpBindingWriterTest.java index f1b07c7f3e..096ff32d4f 100644 --- a/http/http-binding/src/test/java/software/amazon/smithy/java/http/binding/HttpBindingWriterTest.java +++ b/http/http-binding/src/test/java/software/amazon/smithy/java/http/binding/HttpBindingWriterTest.java @@ -9,11 +9,14 @@ import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.not; import static org.hamcrest.Matchers.notNullValue; import static org.hamcrest.Matchers.nullValue; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assumptions.assumeTrue; +import java.io.ByteArrayOutputStream; import java.math.BigDecimal; import java.math.BigInteger; import java.net.URLDecoder; @@ -49,6 +52,7 @@ import software.amazon.smithy.java.http.api.HttpHeaders; import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.http.api.HttpResponse; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.java.io.uri.SmithyUri; import software.amazon.smithy.java.json.JsonCodec; import software.amazon.smithy.model.pattern.UriPattern; @@ -404,6 +408,34 @@ private static Bound populated(UnaryOperator customize) { return customize.apply(builder).build(); } + @Test + public void clientRequestBodiesComeFromPooledMemory() throws Exception { + for (boolean codegen : new boolean[] {true, false}) { + var body = request(Bound.OPERATION, fullyPopulated(), codegen).body(); + assertThat("codegen=" + codegen, body, instanceOf(PooledDataStream.class)); + var pooled = (PooledDataStream) body; + assertThat(pooled.contentType(), equalTo("application/json")); + + var out = new ByteArrayOutputStream(); + pooled.writeTo(out); + assertThat(out.toString(StandardCharsets.UTF_8), equalTo("{\"note\":\"hi\"}")); + assertThat(pooled.escaped(), equalTo(false)); + + pooled.release(); + assertThat(pooled.isReleased(), equalTo(true)); + assertThrows(IllegalStateException.class, pooled::asByteBuffer); + } + } + + @Test + public void serverResponseBodiesAreNotPooled() { + for (boolean codegen : new boolean[] {true, false}) { + var body = response(new Status(204, "hi"), codegen).body(); + assertThat("codegen=" + codegen, body, not(instanceOf(PooledDataStream.class))); + assertThat(utf8(body.asByteBuffer()), equalTo("{\"note\":\"hi\"}")); + } + } + private static HttpRequest generated(SerializableStruct input) { return generated(Bound.OPERATION, input); } diff --git a/http/http-client/src/main/java/software/amazon/smithy/java/http/client/DefaultHttpClient.java b/http/http-client/src/main/java/software/amazon/smithy/java/http/client/DefaultHttpClient.java index 18513999da..b5572608ca 100644 --- a/http/http-client/src/main/java/software/amazon/smithy/java/http/client/DefaultHttpClient.java +++ b/http/http-client/src/main/java/software/amazon/smithy/java/http/client/DefaultHttpClient.java @@ -32,6 +32,7 @@ import software.amazon.smithy.java.http.client.connection.HttpConnection; import software.amazon.smithy.java.http.client.connection.Route; import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.java.logging.InternalLogger; /** @@ -194,7 +195,8 @@ private HttpResponse sendForRoute( } private static boolean shouldWriteH2BodyInline(DataStream body) { - return body.isReplayable() && body.hasKnownLength(); + // Write pooled bodies inline so no background reader remains when the caller releases them. + return body instanceof PooledDataStream || (body.isReplayable() && body.hasKnownLength()); } /** @@ -479,12 +481,14 @@ private HttpResponse sendWithTimeout( return future.get(timeout.toMillis(), TimeUnit.MILLISECONDS); } catch (TimeoutException e) { future.cancel(true); + abandonPooledBody(request); throw new IOException(String.format( "Request to `%s` exceeded request timeout of %s seconds", request.uri().getHost(), timeout.toSeconds()), e); } catch (InterruptedException e) { Thread.currentThread().interrupt(); + abandonPooledBody(request); throw new IOException("Interrupted while waiting for HTTP request to complete to `" + request.uri().getHost() + '`', e); } catch (ExecutionException e) { @@ -492,6 +496,13 @@ private HttpResponse sendWithTimeout( } } + // The upload worker may continue after send() returns. Prevent recycling while it can still read. + private static void abandonPooledBody(HttpRequest request) { + if (request.body() instanceof PooledDataStream pooled) { + pooled.abandon(); + } + } + private static IOException unwrap(ExecutionException e) throws IOException { var cause = e.getCause(); return switch (cause) { diff --git a/http/http-client/src/main/java/software/amazon/smithy/java/http/client/h2/H2StreamRequestBody.java b/http/http-client/src/main/java/software/amazon/smithy/java/http/client/h2/H2StreamRequestBody.java index a69e4c9ec9..1ad2221ff5 100644 --- a/http/http-client/src/main/java/software/amazon/smithy/java/http/client/h2/H2StreamRequestBody.java +++ b/http/http-client/src/main/java/software/amazon/smithy/java/http/client/h2/H2StreamRequestBody.java @@ -10,6 +10,7 @@ import java.nio.channels.ReadableByteChannel; import java.util.function.Supplier; import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; /** * Per-stream outbound body state. @@ -54,7 +55,10 @@ void writeRequestBody(DataStream body) throws IOException { return; } - if (body.isReplayable() && body.hasKnownLength() && body.contentLength() <= Integer.MAX_VALUE) { + // Use frame buffers for pooled bodies to avoid creating a private copy. + if (body.isReplayable() && body.hasKnownLength() + && body.contentLength() <= Integer.MAX_VALUE + && !(body instanceof PooledDataStream)) { try { if (body.contentLength() <= DIRECT_REPLAYABLE_UPLOAD_LIMIT) { exchange.writeReplayableBody(body.asByteBuffer(), true); diff --git a/http/http-client/src/test/java/software/amazon/smithy/java/http/client/DefaultHttpClientTest.java b/http/http-client/src/test/java/software/amazon/smithy/java/http/client/DefaultHttpClientTest.java index 8b75f5f433..72159ef96a 100644 --- a/http/http-client/src/test/java/software/amazon/smithy/java/http/client/DefaultHttpClientTest.java +++ b/http/http-client/src/test/java/software/amazon/smithy/java/http/client/DefaultHttpClientTest.java @@ -11,20 +11,25 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.nio.channels.ReadableByteChannel; +import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.UnaryOperator; import java.util.stream.Stream; import javax.net.ssl.SSLSession; import org.junit.jupiter.api.Test; @@ -506,6 +511,113 @@ public int responseStatusCode() throws IOException { } } + @Test + void requestTimeoutAbandonsAPooledBodyStillBeingWritten() throws Exception { + assertTimeoutAbandonsBody(body -> body); + } + + @Test + void requestTimeoutAbandonsAMetadataWrappedPooledBody() throws Exception { + assertTimeoutAbandonsBody(body -> DataStream.withMetadata(body, "application/octet-stream", -1)); + } + + private void assertTimeoutAbandonsBody(UnaryOperator wrap) throws Exception { + var enteredWrite = new CountDownLatch(1); + var workerDone = new CountDownLatch(1); + var written = new AtomicReference(); + var recycled = new AtomicBoolean(); + var pool = new TestConnectionPool() { + @Override + protected HttpExchange createExchange() { + return new TestHttpExchange() { + @Override + public void writeRequestBody(DataStream body) throws IOException { + body.writeTo(new OutputStream() { + @Override + public void write(int b) { + throw new UnsupportedOperationException(); + } + + @Override + public void write(byte[] b, int off, int len) { + // Simulate a kernel write that continues after timeout and ignores interruption. + enteredWrite.countDown(); + long deadline = System.nanoTime() + Duration.ofMillis(300).toNanos(); + while (System.nanoTime() < deadline) { + Thread.onSpinWait(); + } + written.set(new String(b, off, len, StandardCharsets.US_ASCII)); + workerDone.countDown(); + } + }); + } + }; + } + }; + byte[] array = "pooled body".getBytes(StandardCharsets.US_ASCII); + var lease = DataStream.ofLeasedBytes(array, 0, array.length, "text/plain", a -> recycled.set(true)); + try (var client = HttpClient.builder() + .connectionPoolFactory(config -> pool) + .requestTimeout(Duration.ofMillis(50)) + .build()) { + var request = HttpRequest.create() + .setMethod("PUT") + .setUri(SmithyUri.of("http://example.com/test")) + .setBody(wrap.apply(lease)); + + var ex = assertThrows(IOException.class, () -> client.send(request)); + assertTrue(ex.getMessage().contains("exceeded request timeout"), ex.getMessage()); + assertTrue(enteredWrite.await(5, TimeUnit.SECONDS)); + + assertTrue(lease.escaped(), "lease still in use by the worker must be abandoned"); + lease.release(); + assertFalse(recycled.get(), "array must not be recycled while the worker may still read it"); + + assertTrue(workerDone.await(5, TimeUnit.SECONDS)); + assertEquals("pooled body", written.get()); + } + } + + @Test + void pooledBodiesAreWrittenInlineOnH2EvenWhenMarkedNonReplayable() throws Exception { + var writerThread = new AtomicReference(); + var written = new AtomicReference(); + var pool = new TestConnectionPool() { + @Override + protected HttpExchange createExchange() { + return new TestHttpExchange() { + @Override + public boolean supportsBidirectionalStreaming() { + return true; + } + + @Override + public void writeRequestBody(DataStream body) throws IOException { + writerThread.set(Thread.currentThread()); + var out = new ByteArrayOutputStream(); + body.writeTo(out); + written.set(out.toString(StandardCharsets.US_ASCII)); + } + }; + } + }; + byte[] array = "pooled body".getBytes(StandardCharsets.US_ASCII); + var lease = DataStream.ofLeasedBytes(array, 0, array.length, "text/plain", a -> {}); + var body = DataStream.withMetadata(lease, "text/plain", array.length, false); + assertFalse(body.isReplayable()); + + try (var client = HttpClient.builder().connectionPoolFactory(config -> pool).build()) { + var request = HttpRequest.create() + .setMethod("PUT") + .setUri(SmithyUri.of("http://example.com/test")) + .setBody(body); + client.send(request); + } + + assertEquals(Thread.currentThread(), writerThread.get(), "pooled body must be written inline"); + assertEquals("pooled body", written.get()); + } + @Test void requestTimeoutSucceedsWhenFastEnough() throws IOException { var pool = new TestConnectionPool(); diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/ByteArrayRecycler.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/ByteArrayRecycler.java new file mode 100644 index 0000000000..da0b961af9 --- /dev/null +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/ByteArrayRecycler.java @@ -0,0 +1,21 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.io.datastream; + +/** + * Receives an array after its pooled stream releases ownership. + * + *

Implementations must be thread-safe and must not throw. + */ +@FunctionalInterface +public interface ByteArrayRecycler { + /** + * Takes ownership of an array with no remaining readers. + * + * @param array array whose contents can be overwritten + */ + void recycle(byte[] array); +} diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/DataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/DataStream.java index 25b880c24b..dc19841a06 100644 --- a/io/src/main/java/software/amazon/smithy/java/io/datastream/DataStream.java +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/DataStream.java @@ -19,6 +19,7 @@ import java.nio.file.Path; import java.util.List; import java.util.concurrent.Flow; +import java.util.function.Consumer; /** * Abstraction for reading streams of data. @@ -188,6 +189,22 @@ default boolean hasByteBuffer() { return false; } + /** + * Visits the content as ByteBuffer views in order. + * + *

Views remain valid only during the callback. + * Do not retain or change a view. Read only between its position and limit. + * The backing array can contain bytes outside that range. + * Visiting a pooled stream requires no copy. Visiting a non-replayable stream consumes it. + * + *

The default visits {@link #asByteBuffer()} once. + * + * @param action callback for each view + */ + default void forEachBuffer(Consumer action) { + action.accept(asByteBuffer()); + } + @Override default void subscribe(Flow.Subscriber subscriber) { HttpRequest.BodyPublishers.ofInputStream(this::asInputStream).subscribe(subscriber); @@ -389,6 +406,30 @@ static DataStream ofByteBuffer(ByteBuffer buffer, String contentType) { return new ByteBufferDataStream(buffer, contentType); } + /** + * Creates a replayable stream that owns a borrowed array until {@link PooledDataStream#release()}. + * + *

Release returns the array to {@code recycler}, unless the stream was abandoned. + * See {@link PooledDataStream} for read and lifecycle requirements. + * + * @param array borrowed array + * @param offset first content byte + * @param length number of content bytes + * @param contentType content type, or null + * @param recycler receives the array on release unless abandoned + * @return the pooled stream + * @throws IndexOutOfBoundsException if the range exceeds the array + */ + static PooledDataStream ofLeasedBytes( + byte[] array, + int offset, + int length, + String contentType, + ByteArrayRecycler recycler + ) { + return new LeasedDataStream(array, offset, length, contentType, recycler); + } + /** * Create a DataStream from a file on disk. * @@ -495,10 +536,11 @@ static DataStream withMetadata(DataStream ds, String contentType, long contentLe return ds; } - return new WrappedDataStream( - ds, - changeContentLength ? contentLength : ds.contentLength(), - changeContentType ? contentType : ds.contentType(), - isReplayable); + long length = changeContentLength ? contentLength : ds.contentLength(); + String type = changeContentType ? contentType : ds.contentType(); + if (ds instanceof PooledDataStream pooled) { + return new WrappedPooledDataStream(pooled, length, type, isReplayable); + } + return new WrappedDataStream(ds, length, type, isReplayable); } } diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/LeasedDataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/LeasedDataStream.java new file mode 100644 index 0000000000..ad7a77662f --- /dev/null +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/LeasedDataStream.java @@ -0,0 +1,174 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.io.datastream; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.channels.Channels; +import java.nio.channels.ReadableByteChannel; +import java.nio.channels.WritableByteChannel; +import java.util.Arrays; +import java.util.Objects; +import java.util.concurrent.Flow; +import java.util.function.Consumer; + +final class LeasedDataStream implements PooledDataStream { + private final byte[] array; + private final int offset; + private final int length; + private final String contentType; + private final ByteArrayRecycler recycler; + private byte[] aliasCopy; + private boolean abandoned; + private boolean released; + + LeasedDataStream(byte[] array, int offset, int length, String contentType, ByteArrayRecycler recycler) { + this.array = Objects.requireNonNull(array, "array"); + Objects.checkFromIndexSize(offset, length, array.length); + this.offset = offset; + this.length = length; + this.contentType = contentType; + this.recycler = Objects.requireNonNull(recycler, "recycler"); + } + + @Override + public long contentLength() { + return length; + } + + @Override + public boolean hasKnownLength() { + return true; + } + + @Override + public String contentType() { + return contentType; + } + + @Override + public boolean isReplayable() { + return true; + } + + @Override + public boolean isAvailable() { + return !released; + } + + @Override + public boolean hasByteBuffer() { + return true; + } + + @Override + public boolean isReleased() { + return released; + } + + @Override + public boolean escaped() { + return aliasCopy != null || abandoned; + } + + @Override + public ByteBuffer asByteBuffer() { + return ByteBuffer.wrap(aliasCopy()); + } + + @Override + public InputStream asInputStream() { + return new ByteArrayInputStream(aliasCopy()); + } + + @Override + public ReadableByteChannel asChannel() { + return Channels.newChannel(asInputStream()); + } + + @Override + public void subscribe(Flow.Subscriber subscriber) { + DataStream.ofByteBuffer(asByteBuffer(), contentType).subscribe(subscriber); + } + + private byte[] aliasCopy() { + checkNotReleased(); + byte[] copy = aliasCopy; + if (copy == null) { + copy = aliasCopy = Arrays.copyOfRange(array, offset, offset + length); + } + return copy; + } + + @Override + public void writeTo(OutputStream out) throws IOException { + checkNotReleased(); + out.write(array, offset, length); + } + + @Override + public void writeTo(WritableByteChannel channel) throws IOException { + checkNotReleased(); + ByteBuffer view = ByteBuffer.wrap(array, offset, length); + while (view.hasRemaining()) { + channel.write(view); + } + } + + @Override + public void forEachBuffer(Consumer action) { + checkNotReleased(); + action.accept(ByteBuffer.wrap(array, offset, length).slice()); + } + + @Override + public DataStream copy() { + checkNotReleased(); + return DataStream.ofBytes(Arrays.copyOfRange(array, offset, offset + length), contentType); + } + + @Override + public void discard() {} + + @Override + public void abandon() { + if (!released) { + abandoned = true; + } + } + + @Override + public void release() { + if (released) { + return; + } + released = true; + if (!abandoned) { + recycler.recycle(array); + } + } + + private void checkNotReleased() { + if (released) { + throw new IllegalStateException( + "This pooled request body was released when its call completed and cannot be read again. " + + "Code that keeps a request past its call must copy the body while the call is in progress."); + } + } + + @Override + public String toString() { + return "LeasedDataStream{length=" + length + + ", contentType=" + contentType + + ", aliased=" + (aliasCopy != null) + + ", abandoned=" + abandoned + + ", released=" + released + + '}'; + } +} diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/MultiByteBufferDataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/MultiByteBufferDataStream.java index d258be388a..59da077844 100644 --- a/io/src/main/java/software/amazon/smithy/java/io/datastream/MultiByteBufferDataStream.java +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/MultiByteBufferDataStream.java @@ -15,6 +15,7 @@ import java.util.List; import java.util.concurrent.Flow; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; import software.amazon.smithy.java.io.ByteBufferUtils; /** @@ -77,6 +78,13 @@ public String contentType() { return contentType; } + @Override + public void forEachBuffer(Consumer action) { + for (var buffer : buffers) { + action.accept(buffer.duplicate()); + } + } + @Override public boolean hasByteBuffer() { // The contiguous buffer is materialized lazily by asByteBuffer(). diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/PooledDataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/PooledDataStream.java new file mode 100644 index 0000000000..a7e1d43c88 --- /dev/null +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/PooledDataStream.java @@ -0,0 +1,66 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.io.datastream; + +/** + * A replayable in-memory stream with pooled storage. + * + *

The client pipeline owns request leases and releases them after the call completes, including retries. + * Call {@link #copy()} before release to retain the content. + * Reads after release throw {@link IllegalStateException}. + * + *

Read methods use storage as follows: + *

    + *
  • {@code writeTo}, {@link #forEachBuffer}, and {@link #copy()} read pooled storage directly.
    • + *
  • {@link #asByteBuffer()}, {@link #asInputStream()}, {@link #asChannel()}, and {@code subscribe} + * use a private copy created on the first such read and reused for later reads. + * These views remain valid after release.
    • + *
+ * + *

Pooled streams are not thread-safe. + * Complete all reads before release. If a reader can continue after the transport returns, call {@link #abandon()}. + * + *

Wrappers that read pooled storage must implement this interface and forward its lifecycle methods. + */ +public interface PooledDataStream extends DataStream { + /** + * Ends the lease and returns its storage to the pool unless abandoned. + * Complete all reads before calling this method. Repeated calls are safe. + */ + void release(); + + /** + * @return true if the stream was released and can no longer be read + */ + boolean isReleased(); + + /** + * Prevents recycling when an uncontrolled reader can still access pooled storage. + * Call this before returning from a transport while an upload worker can still read. + * Calls after release have no effect. + */ + void abandon(); + + /** + * @return true if a read created a private copy or {@link #abandon()} was called before release + */ + boolean escaped(); + + /** + * Copies the content into an unpooled stream. + * + * @return a replayable copy with the same content type + * @throws IllegalStateException if this stream was released + */ + DataStream copy(); + + /** + * Does nothing so the body remains available for retries. + * Use {@link #release()} to end the lease. + */ + @Override + default void close() {} +} diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedDataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedDataStream.java index 1fd3d4b843..09dcc9772d 100644 --- a/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedDataStream.java +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedDataStream.java @@ -12,8 +12,9 @@ import java.nio.channels.ReadableByteChannel; import java.nio.channels.WritableByteChannel; import java.util.concurrent.Flow; +import java.util.function.Consumer; -final class WrappedDataStream implements DataStream { +class WrappedDataStream implements DataStream { private final DataStream delegate; private final String contentType; @@ -47,6 +48,11 @@ public void writeTo(OutputStream out) throws IOException { delegate.writeTo(out); } + @Override + public void forEachBuffer(Consumer action) { + delegate.forEachBuffer(action); + } + @Override public void discard() throws IOException { delegate.discard(); diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedPooledDataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedPooledDataStream.java new file mode 100644 index 0000000000..20c796ba7b --- /dev/null +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedPooledDataStream.java @@ -0,0 +1,50 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.io.datastream; + +final class WrappedPooledDataStream extends WrappedDataStream implements PooledDataStream { + + private final PooledDataStream pooled; + + WrappedPooledDataStream(PooledDataStream delegate, long contentLength, String contentType, boolean isReplayable) { + super(delegate, contentLength, contentType, isReplayable); + this.pooled = delegate; + } + + @Override + public void release() { + pooled.release(); + } + + @Override + public boolean isReleased() { + return pooled.isReleased(); + } + + @Override + public void abandon() { + pooled.abandon(); + } + + @Override + public boolean escaped() { + return pooled.escaped(); + } + + @Override + public DataStream copy() { + return DataStream.withMetadata(pooled.copy(), contentType(), contentLength(), isReplayable()); + } + + @Override + public void close() {} + + @Override + public String toString() { + return "WrappedPooledDataStream{contentType=" + contentType() + ", contentLength=" + contentLength() + + ", delegate=" + pooled + '}'; + } +} diff --git a/io/src/test/java/software/amazon/smithy/java/io/datastream/LeasedDataStreamTest.java b/io/src/test/java/software/amazon/smithy/java/io/datastream/LeasedDataStreamTest.java new file mode 100644 index 0000000000..4e164da7b5 --- /dev/null +++ b/io/src/test/java/software/amazon/smithy/java/io/datastream/LeasedDataStreamTest.java @@ -0,0 +1,331 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.io.datastream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.channels.Channels; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Flow; +import org.junit.jupiter.api.Test; + +class LeasedDataStreamTest { + + private static final String CONTENT = "HELLO WORLD"; + + private final List recycled = new ArrayList<>(); + private final ByteArrayRecycler recycler = recycled::add; + private final byte[] array = ("xx" + CONTENT + "yyy").getBytes(StandardCharsets.US_ASCII); + + private PooledDataStream stream() { + return DataStream.ofLeasedBytes(array, 2, CONTENT.length(), "text/plain", recycler); + } + + private static String ascii(ByteBuffer buffer) { + return StandardCharsets.US_ASCII.decode(buffer).toString(); + } + + @Test + void reportsMetadata() { + var s = stream(); + assertThat(s.contentLength()).isEqualTo(11); + assertThat(s.hasKnownLength()).isTrue(); + assertThat(s.contentType()).isEqualTo("text/plain"); + assertThat(s.isReplayable()).isTrue(); + assertThat(s.isAvailable()).isTrue(); + assertThat(s.hasByteBuffer()).isTrue(); + assertThat(s.isReleased()).isFalse(); + assertThat(s.escaped()).isFalse(); + assertThat(s.toString()).contains("length=11", "aliased=false", "abandoned=false", "released=false"); + } + + @Test + void rejectsRangesOutsideTheArray() { + assertThatThrownBy(() -> DataStream.ofLeasedBytes(array, 10, 10, null, recycler)) + .isInstanceOf(IndexOutOfBoundsException.class); + assertThatThrownBy(() -> DataStream.ofLeasedBytes(array, -1, 1, null, recycler)) + .isInstanceOf(IndexOutOfBoundsException.class); + assertThatThrownBy(() -> DataStream.ofLeasedBytes(array, 0, -1, null, recycler)) + .isInstanceOf(IndexOutOfBoundsException.class); + } + + @Test + void writeToReadsThePooledArrayAndReleaseRecyclesOnce() throws IOException { + var s = stream(); + var out = new ByteArrayOutputStream(); + s.writeTo(out); + s.writeTo(out); + assertThat(out.toString(StandardCharsets.US_ASCII)).isEqualTo(CONTENT + CONTENT); + assertThat(s.escaped()).isFalse(); + + s.release(); + assertThat(recycled).containsExactly(array); + assertThat(s.isReleased()).isTrue(); + assertThat(s.isAvailable()).isFalse(); + + s.release(); + assertThat(recycled).hasSize(1); + } + + @Test + void channelWriteAndForEachBufferReadThePooledArray() throws IOException { + var s = stream(); + var out = new ByteArrayOutputStream(); + s.writeTo(Channels.newChannel(out)); + assertThat(out.toString(StandardCharsets.US_ASCII)).isEqualTo(CONTENT); + + long[] seen = {0}; + s.forEachBuffer(b -> { + assertThat(b.hasArray()).isTrue(); + assertThat(b.array()).isSameAs(array); + assertThat(b.arrayOffset()).isEqualTo(2); + assertThat(b.position()).isZero(); + assertThat(b.remaining()).isEqualTo(11); + assertThat(ascii(b.duplicate())).isEqualTo(CONTENT); + seen[0] += b.remaining(); + }); + assertThat(seen[0]).isEqualTo(11); + assertThat(s.escaped()).isFalse(); + + s.release(); + assertThat(recycled).containsExactly(array); + } + + @Test + void everyReadThrowsAfterRelease() { + var s = stream(); + s.release(); + assertThatThrownBy(() -> s.writeTo(new ByteArrayOutputStream())).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(() -> s.writeTo(Channels.newChannel(new ByteArrayOutputStream()))) + .isInstanceOf(IllegalStateException.class); + assertThatThrownBy(() -> s.forEachBuffer(b -> {})).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(s::asByteBuffer).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(s::asInputStream).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(s::asChannel).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(() -> s.subscribe(new NoopSubscriber())).isInstanceOf(IllegalStateException.class); + assertThatThrownBy(s::copy).isInstanceOf(IllegalStateException.class); + assertThat(recycled).containsExactly(array); + assertThat(s.escaped()).isFalse(); + } + + @Test + void aliasingReadsAreServedFromAPrivateCopyAndTheArrayIsStillRecycled() throws IOException { + var s = stream(); + ByteBuffer view = s.asByteBuffer(); + var in = s.asInputStream(); + assertThat(view.isReadOnly()).isFalse(); + assertThat(view.hasArray()).isTrue(); + assertThat(view.remaining()).isEqualTo(11); + assertThat(s.escaped()).isTrue(); + assertThat(s.toString()).contains("aliased=true"); + + s.release(); + assertThat(recycled).containsExactly(array); + Arrays.fill(array, (byte) 'J'); + assertThat(ascii(view)).isEqualTo(CONTENT); + assertThat(new String(in.readAllBytes(), StandardCharsets.US_ASCII)).isEqualTo(CONTENT); + + assertThatThrownBy(s::asByteBuffer).isInstanceOf(IllegalStateException.class); + } + + @Test + void channelAndSubscriptionAreServedFromThePrivateCopyToo() throws IOException { + var c = stream(); + var channel = c.asChannel(); + assertThat(c.escaped()).isTrue(); + c.release(); + Arrays.fill(array, (byte) 'J'); + var buf = ByteBuffer.allocate(32); + while (channel.read(buf) != -1) {} + buf.flip(); + assertThat(ascii(buf)).isEqualTo(CONTENT); + assertThat(recycled).containsExactly(array); + } + + @Test + void subscriptionIsServedFromThePrivateCopy() { + var t = stream(); + var received = new ArrayList(); + var done = new CountDownLatch(1); + t.subscribe(new Flow.Subscriber<>() { + @Override + public void onSubscribe(Flow.Subscription subscription) { + subscription.request(Long.MAX_VALUE); + } + + @Override + public void onNext(ByteBuffer item) { + received.add(item); + } + + @Override + public void onError(Throwable throwable) { + done.countDown(); + } + + @Override + public void onComplete() { + done.countDown(); + } + }); + assertThat(done.getCount()).isZero(); + assertThat(received).hasSize(1); + assertThat(t.escaped()).isTrue(); + t.release(); + assertThat(recycled).containsExactly(array); + Arrays.fill(array, (byte) 'J'); + assertThat(ascii(received.get(0))).isEqualTo(CONTENT); + } + + @Test + void nonAliasingReadsKeepUsingTheArrayAfterAnAlias() throws IOException { + var s = stream(); + s.asByteBuffer(); + s.forEachBuffer(b -> assertThat(b.array()).isSameAs(array)); + var out = new ByteArrayOutputStream(); + s.writeTo(out); + assertThat(out.toString(StandardCharsets.US_ASCII)).isEqualTo(CONTENT); + s.release(); + assertThat(recycled).containsExactly(array); + } + + @Test + void closeIsNotRelease() throws IOException { + var s = stream(); + s.close(); + assertThat(s.isReleased()).isFalse(); + assertThat(recycled).isEmpty(); + var out = new ByteArrayOutputStream(); + s.writeTo(out); + assertThat(out.toString(StandardCharsets.US_ASCII)).isEqualTo(CONTENT); + s.release(); + assertThat(recycled).containsExactly(array); + } + + @Test + void copyOutlivesRelease() { + var s = stream(); + DataStream copy = s.copy(); + assertThat(s.escaped()).isFalse(); + assertThat(copy).isNotInstanceOf(PooledDataStream.class); + s.release(); + assertThat(recycled).containsExactly(array); + Arrays.fill(array, (byte) 'J'); + assertThat(copy.contentType()).isEqualTo("text/plain"); + assertThat(copy.contentLength()).isEqualTo(11); + assertThat(ascii(copy.asByteBuffer())).isEqualTo(CONTENT); + } + + @Test + void metadataWrapperVisitsAndWritesWithoutAliasing() throws IOException { + var s = stream(); + DataStream wrapped = DataStream.withMetadata(s, "application/octet-stream", -1); + assertThat(wrapped).isNotSameAs(s); + var out = new ByteArrayOutputStream(); + wrapped.writeTo(out); + wrapped.forEachBuffer(b -> out.write(b.get())); + assertThat(out.toString(StandardCharsets.US_ASCII)).isEqualTo(CONTENT + "H"); + assertThat(s.escaped()).isFalse(); + s.release(); + assertThat(recycled).containsExactly(array); + } + + @Test + void metadataWrapperOverAPooledBodyForwardsTheLifecycle() throws IOException { + var s = stream(); + DataStream wrapped = DataStream.withMetadata(s, "application/octet-stream", -1); + assertThat(wrapped).isInstanceOf(PooledDataStream.class); + var pooledWrapper = (PooledDataStream) wrapped; + assertThat(wrapped.contentType()).isEqualTo("application/octet-stream"); + assertThat(wrapped.contentLength()).isEqualTo(11); + + DataStream copy = pooledWrapper.copy(); + assertThat(copy).isNotInstanceOf(PooledDataStream.class); + assertThat(copy.contentType()).isEqualTo("application/octet-stream"); + assertThat(ascii(copy.asByteBuffer())).isEqualTo(CONTENT); + + pooledWrapper.close(); + assertThat(s.isReleased()).isFalse(); + + pooledWrapper.abandon(); + assertThat(s.escaped()).isTrue(); + assertThat(pooledWrapper.escaped()).isTrue(); + pooledWrapper.release(); + assertThat(s.isReleased()).isTrue(); + assertThat(pooledWrapper.isReleased()).isTrue(); + assertThat(recycled).as("abandoned through the wrapper").isEmpty(); + assertThatThrownBy(() -> wrapped.writeTo(new ByteArrayOutputStream())) + .isInstanceOf(IllegalStateException.class); + } + + @Test + void abandonLeavesTheArrayToTheGc() throws IOException { + var s = stream(); + s.abandon(); + assertThat(s.escaped()).isTrue(); + var out = new ByteArrayOutputStream(); + s.writeTo(out); + assertThat(out.toString(StandardCharsets.US_ASCII)).isEqualTo(CONTENT); + s.release(); + assertThat(recycled).isEmpty(); + + var u = stream(); + u.asByteBuffer(); + u.abandon(); + u.release(); + assertThat(recycled).isEmpty(); + + var t = stream(); + t.release(); + t.abandon(); + assertThat(recycled).containsExactly(array); + } + + @Test + void abandonMidWritePreventsRecycling() throws IOException { + var s = stream(); + s.writeTo(new OutputStream() { + @Override + public void write(int b) { + throw new UnsupportedOperationException(); + } + + @Override + public void write(byte[] b, int off, int len) { + s.abandon(); + s.release(); + } + }); + assertThat(s.escaped()).isTrue(); + assertThat(s.isReleased()).isTrue(); + assertThat(recycled).isEmpty(); + } + + private static final class NoopSubscriber implements Flow.Subscriber { + @Override + public void onSubscribe(Flow.Subscription subscription) { + subscription.request(Long.MAX_VALUE); + } + + @Override + public void onNext(ByteBuffer item) {} + + @Override + public void onError(Throwable throwable) {} + + @Override + public void onComplete() {} + } +} diff --git a/protocol-test-harness/src/main/java/software/amazon/smithy/java/protocoltests/harness/EventStreamClientTestsProtocolTestProvider.java b/protocol-test-harness/src/main/java/software/amazon/smithy/java/protocoltests/harness/EventStreamClientTestsProtocolTestProvider.java index bc30e89b4f..c97df4ff4e 100644 --- a/protocol-test-harness/src/main/java/software/amazon/smithy/java/protocoltests/harness/EventStreamClientTestsProtocolTestProvider.java +++ b/protocol-test-harness/src/main/java/software/amazon/smithy/java/protocoltests/harness/EventStreamClientTestsProtocolTestProvider.java @@ -24,6 +24,7 @@ import software.amazon.smithy.java.client.core.MessageExchange; import software.amazon.smithy.java.client.core.RequestOverrideConfig; import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeResolver; +import software.amazon.smithy.java.client.http.HttpClientProtocol; import software.amazon.smithy.java.client.http.HttpMessageExchange; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.schema.ApiOperation; @@ -324,7 +325,7 @@ public HttpRequest getCapturedRequest() { @Override public HttpResponse send(Context context, HttpRequest request) { - this.capturedRequest = request; + this.capturedRequest = HttpClientProtocol.retainRequest(request); return DUMMY_RESPONSE; } diff --git a/protocol-test-harness/src/main/java/software/amazon/smithy/java/protocoltests/harness/HttpClientRequestProtocolTestProvider.java b/protocol-test-harness/src/main/java/software/amazon/smithy/java/protocoltests/harness/HttpClientRequestProtocolTestProvider.java index 04aa287255..e73fe3f697 100644 --- a/protocol-test-harness/src/main/java/software/amazon/smithy/java/protocoltests/harness/HttpClientRequestProtocolTestProvider.java +++ b/protocol-test-harness/src/main/java/software/amazon/smithy/java/protocoltests/harness/HttpClientRequestProtocolTestProvider.java @@ -20,6 +20,7 @@ import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeOption; import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeResolver; import software.amazon.smithy.java.client.core.plugins.InjectIdempotencyTokenPlugin; +import software.amazon.smithy.java.client.http.HttpClientProtocol; import software.amazon.smithy.java.client.http.HttpMessageExchange; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.schema.ApiOperation; @@ -216,7 +217,7 @@ public HttpRequest getCapturedRequest() { @Override public HttpResponse send(Context context, HttpRequest request) { - this.capturedRequest = request; + this.capturedRequest = HttpClientProtocol.retainRequest(request); return exceptionalResponse; } From 641f796218afa32d6c832391dfc6a299ecdedb52 Mon Sep 17 00:00:00 2001 From: Adwait Kumar Singh Date: Tue, 13 Oct 2026 02:23:57 +0530 Subject: [PATCH 2/2] Fixes after rebase --- .../S3ExpressChecksumInterceptor.java | 2 +- .../scheme/sigv4/AwsChunkedDataStream.java | 255 ++++++----- .../sigv4/AwsChunkedDataStreamTest.java | 427 +++++++++++++++--- .../auth/scheme/sigv4/SigV4CaseGenerator.java | 23 +- .../auth/scheme/sigv4/SigV4Differential.java | 4 + .../auth/scheme/sigv4/SigV4SignerTest.java | 48 ++ .../sigv4/SigningResourcesPropertyTest.java | 6 + .../aws/client/awsjson/AwsJsonProtocol.java | 1 + .../aws-client-awsquery/build.gradle.kts | 2 + .../awsquery/AwsQueryClientProtocol.java | 1 + .../awsquery/Ec2QueryClientProtocol.java | 3 +- .../QueryResponseBodyReleaseTest.java | 207 +++++++++ .../java/client/http/HttpClientProtocol.java | 15 + .../client/http/JavaHttpClientTransport.java | 11 +- .../java/client/http/JavaHttpClientTest.java | 45 +- .../rpcv2/AbstractRpcV2ClientProtocol.java | 1 + .../ForwardingPooledDataStream.java | 38 ++ .../java/io/datastream/PooledDataStream.java | 3 +- .../datastream/WrappedPooledDataStream.java | 21 +- 19 files changed, 907 insertions(+), 206 deletions(-) create mode 100644 aws/client/aws-client-awsquery/src/test/java/software/amazon/smithy/java/aws/client/awsquery/QueryResponseBodyReleaseTest.java create mode 100644 io/src/main/java/software/amazon/smithy/java/io/datastream/ForwardingPooledDataStream.java diff --git a/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressChecksumInterceptor.java b/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressChecksumInterceptor.java index d0ac64eca2..d7fd246cee 100644 --- a/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressChecksumInterceptor.java +++ b/aws/aws-sigv4-s3express/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/s3express/S3ExpressChecksumInterceptor.java @@ -87,7 +87,7 @@ public RequestT modifyBeforeSigning(RequestHook hook) return hook.request(); } - var chunked = new AwsChunkedDataStream(body); + var chunked = AwsChunkedDataStream.of(body); var modifiable = req.toModifiable(); modifiable.setBody(chunked); var headers = modifiable.headers(); diff --git a/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStream.java b/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStream.java index 4a5b223dd4..8dcbcec793 100644 --- a/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStream.java +++ b/aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStream.java @@ -6,18 +6,20 @@ package software.amazon.smithy.java.aws.client.auth.scheme.sigv4; import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; -import java.io.UncheckedIOException; import java.nio.ByteBuffer; +import java.nio.channels.WritableByteChannel; import java.nio.charset.StandardCharsets; import java.util.Base64; +import java.util.List; import java.util.concurrent.Flow; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; import java.util.zip.CRC32; import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.ForwardingPooledDataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; /** * Wraps a {@link DataStream} so its bytes are emitted in {@code aws-chunked} content-encoding @@ -38,36 +40,43 @@ * the body when {@code Content-Encoding: aws-chunked} and {@code x-amz-trailer: x-amz-checksum-*} * are set on the request. * - *

This implementation only supports replayable, in-memory bodies (single-buffer). Streaming - * bodies would emit multiple chunks; we punt on that until needed. + *

Supports only replayable, in-memory bodies encoded as one chunk. + * A {@link PooledDataStream} wrapper forwards lifecycle methods to the lease. + * {@code writeTo} and {@link #forEachBuffer} read leased storage directly. + * Reads after release throw before emitting bytes. */ -public final class AwsChunkedDataStream implements DataStream { +public sealed class AwsChunkedDataStream implements DataStream { private static final byte[] CRLF = {'\r', '\n'}; private static final byte[] TERMINATOR = {'0', '\r', '\n'}; private static final byte[] TRAILER_NAME = "x-amz-checksum-crc32:".getBytes(StandardCharsets.US_ASCII); + // CRC32 is 4 bytes, so its base64 form is always 8 chars. + private static final int FOOTER_LENGTH = CRLF.length + TERMINATOR.length + + TRAILER_NAME.length + + 8 + + CRLF.length + + CRLF.length; private final DataStream inner; + private final byte[] header; private final long encodedLength; private final long decodedLength; + private byte[] footer; + private byte[] encoded; - public AwsChunkedDataStream(DataStream inner) { + public static AwsChunkedDataStream of(DataStream inner) { + return inner instanceof PooledDataStream pooled ? new Pooled(pooled) : new AwsChunkedDataStream(inner); + } + + private AwsChunkedDataStream(DataStream inner) { if (!inner.isReplayable() || !inner.hasByteBuffer()) { throw new IllegalArgumentException("AwsChunkedDataStream requires a replayable, " + "byte-buffer-backed inner stream; got " + inner.getClass().getSimpleName()); } this.inner = inner; this.decodedLength = inner.contentLength(); - // Single chunk: \r\n\r\n + 0\r\n + trailer-line\r\n + \r\n - // CRC32 is 4 bytes → base64 is always 8 chars (with padding) for any 4-byte input. - int hexSizeLen = hexLen(decodedLength); - int trailerLineLen = TRAILER_NAME.length + 8 /* base64(crc32) */ + CRLF.length; - this.encodedLength = (long) hexSizeLen + CRLF.length - + decodedLength - + CRLF.length - + TERMINATOR.length - + trailerLineLen - + CRLF.length; + this.header = header(decodedLength); + this.encodedLength = header.length + decodedLength + FOOTER_LENGTH; } public long decodedLength() { @@ -96,127 +105,147 @@ public boolean isAvailable() { } @Override - public InputStream asInputStream() { - // Materialize the encoded bytes once. Used by tests / sync transports; the publisher path - // (subscribe) is what JavaHttpClientTransport uses on the hot benchmark path and avoids - // this allocation. - var out = new ByteArrayOutputStream((int) encodedLength); - try { - writeTo(out); - } catch (IOException e) { - throw new UncheckedIOException(e); - } - return new ByteArrayInputStream(out.toByteArray()); + public void writeTo(OutputStream out) throws IOException { + checkAvailable(); + byte[] footer = footer(); + out.write(header); + inner.writeTo(out); + out.write(footer); } @Override - public void writeTo(OutputStream out) throws IOException { - var crc = new CRC32(); - inner.forEachBuffer(crc::update); - out.write(hexBytes(decodedLength)); - out.write(CRLF); - inner.writeTo(out); - out.write(CRLF); - out.write(TERMINATOR); - out.write(TRAILER_NAME); - out.write(crc32Base64(crc)); - out.write(CRLF); - out.write(CRLF); + public void writeTo(WritableByteChannel channel) throws IOException { + checkAvailable(); + byte[] footer = footer(); + writeFully(channel, ByteBuffer.wrap(header)); + inner.writeTo(channel); + writeFully(channel, ByteBuffer.wrap(footer)); } @Override - public void subscribe(Flow.Subscriber subscriber) { - subscriber.onSubscribe(new Subscription(subscriber)); + public void forEachBuffer(Consumer action) { + checkAvailable(); + byte[] footer = footer(); + action.accept(ByteBuffer.wrap(header)); + inner.forEachBuffer(action); + action.accept(ByteBuffer.wrap(footer)); + } + + @Override + public ByteBuffer asByteBuffer() { + checkAvailable(); + return ByteBuffer.wrap(encoded()); } - private final class Subscription implements Flow.Subscription { - private final AtomicBoolean done = new AtomicBoolean(); - private final Flow.Subscriber subscriber; + @Override + public InputStream asInputStream() { + checkAvailable(); + return new ByteArrayInputStream(encoded()); + } - Subscription(Flow.Subscriber subscriber) { - this.subscriber = subscriber; + /** Published buffers share storage with the cached encoded body or the wrapped stream. */ + @Override + public void subscribe(Flow.Subscriber subscriber) { + checkAvailable(); + byte[] encoded = this.encoded; + if (encoded != null) { + DataStream.ofByteBuffer(ByteBuffer.wrap(encoded), contentType()).subscribe(subscriber); + return; } + byte[] footer = footer(); + var frames = List.of(ByteBuffer.wrap(header), inner.asByteBuffer(), ByteBuffer.wrap(footer)); + DataStream.ofByteBuffers(frames, encodedLength, contentType()).subscribe(subscriber); + } - @Override - public void request(long n) { - if (n <= 0 || !done.compareAndSet(false, true)) { - return; - } - try { - var body = inner.asByteBuffer(); - var crc = new CRC32(); - crc.update(body.duplicate()); - - // Frame: \r\n - ByteBuffer header = ByteBuffer.wrap(joinHeader(body.remaining())); - subscriber.onNext(header); - - // Frame: - subscriber.onNext(body.duplicate()); - - // Frame: \r\n then 0\r\n then trailer line then \r\n - subscriber.onNext(ByteBuffer.wrap(joinFooter(crc))); - - subscriber.onComplete(); - } catch (RuntimeException e) { - subscriber.onError(e); - } + @Override + public void close() { + inner.close(); + } + + // Check availability before every read, including cached reads, to prevent output after release. + private void checkAvailable() { + if (!inner.isAvailable()) { + throw new IllegalStateException("This aws-chunked body wraps a pooled request body that was released " + + "when its call completed and cannot be read again."); } + } - @Override - public void cancel() { - done.set(true); + private byte[] footer() { + byte[] bytes = footer; + if (bytes == null) { + var crc = new CRC32(); + inner.forEachBuffer(crc::update); + bytes = footer = footer(crc); } + return bytes; } - private static byte[] joinHeader(int chunkSize) { - byte[] hex = hexBytes(chunkSize); - byte[] out = new byte[hex.length + CRLF.length]; - System.arraycopy(hex, 0, out, 0, hex.length); - System.arraycopy(CRLF, 0, out, hex.length, CRLF.length); - return out; + private byte[] encoded() { + byte[] bytes = encoded; + if (bytes == null) { + bytes = encoded = encode(); + } + return bytes; } - private static byte[] joinFooter(CRC32 crc) { - byte[] crcB64 = crc32Base64(crc); - // \r\n + 0\r\n + trailer-name + base64-crc + \r\n + \r\n - int len = CRLF.length + TERMINATOR.length + TRAILER_NAME.length + crcB64.length + CRLF.length + CRLF.length; - byte[] out = new byte[len]; - int p = 0; - System.arraycopy(CRLF, 0, out, p, CRLF.length); - p += CRLF.length; - System.arraycopy(TERMINATOR, 0, out, p, TERMINATOR.length); - p += TERMINATOR.length; - System.arraycopy(TRAILER_NAME, 0, out, p, TRAILER_NAME.length); - p += TRAILER_NAME.length; - System.arraycopy(crcB64, 0, out, p, crcB64.length); - p += crcB64.length; - System.arraycopy(CRLF, 0, out, p, CRLF.length); - p += CRLF.length; - System.arraycopy(CRLF, 0, out, p, CRLF.length); - return out; + private byte[] encode() { + byte[] footer = footer(); + ByteBuffer dst = ByteBuffer.allocate((int) encodedLength); + dst.put(header); + inner.forEachBuffer(dst::put); + dst.put(footer); + return dst.array(); } - private static byte[] crc32Base64(CRC32 crc) { - long value = crc.getValue(); - byte[] crcBytes = { - (byte) ((value >>> 24) & 0xff), - (byte) ((value >>> 16) & 0xff), - (byte) ((value >>> 8) & 0xff), - (byte) (value & 0xff)}; - return Base64.getEncoder().encode(crcBytes); + private static void writeFully(WritableByteChannel channel, ByteBuffer buffer) throws IOException { + while (buffer.hasRemaining()) { + channel.write(buffer); + } } - private static byte[] hexBytes(long size) { - return Long.toHexString(size).getBytes(StandardCharsets.US_ASCII); + private static byte[] header(long size) { + byte[] hex = Long.toHexString(size).getBytes(StandardCharsets.US_ASCII); + return ByteBuffer.allocate(hex.length + CRLF.length).put(hex).put(CRLF).array(); } - private static int hexLen(long size) { - return Long.toHexString(size).length(); + private static byte[] footer(CRC32 crc) { + byte[] crcB64 = Base64.getEncoder().encode(ByteBuffer.allocate(4).putInt((int) crc.getValue()).array()); + return ByteBuffer.allocate(FOOTER_LENGTH) + .put(CRLF) + .put(TERMINATOR) + .put(TRAILER_NAME) + .put(crcB64) + .put(CRLF) + .put(CRLF) + .array(); } - @Override - public void close() { - inner.close(); + private static final class Pooled extends AwsChunkedDataStream implements ForwardingPooledDataStream { + Pooled(PooledDataStream lease) { + super(lease); + } + + @Override + public PooledDataStream lease() { + return (PooledDataStream) super.inner; + } + + @Override + public boolean escaped() { + return lease().escaped() || super.encoded != null; + } + + // A publisher can outlive the lease. Publish a private copy to keep its buffers valid after release. + @Override + public void subscribe(Flow.Subscriber subscriber) { + DataStream.ofByteBuffer(asByteBuffer(), contentType()).subscribe(subscriber); + } + + @Override + public DataStream copy() { + super.checkAvailable(); + byte[] encoded = super.encoded; + return DataStream.ofBytes(encoded != null ? encoded : super.encode(), contentType()); + } } } diff --git a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStreamTest.java b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStreamTest.java index f453b36f50..35137cb4f1 100644 --- a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStreamTest.java +++ b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/AwsChunkedDataStreamTest.java @@ -7,20 +7,28 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.nullValue; +import static org.hamcrest.Matchers.sameInstance; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; import java.nio.ByteBuffer; +import java.nio.channels.Channels; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Base64; import java.util.List; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.Flow; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import java.util.zip.CRC32; import org.junit.jupiter.api.Test; import software.amazon.smithy.java.io.datastream.DataStream; @@ -36,7 +44,7 @@ void writeToProducesAwsChunkedFraming() throws Exception { payload[i] = (byte) ('a' + (i % 8)); } var inner = DataStream.ofBytes(payload, "application/octet-stream"); - var chunked = new AwsChunkedDataStream(inner); + var chunked = AwsChunkedDataStream.of(inner); var out = new ByteArrayOutputStream(); chunked.writeTo(out); @@ -63,7 +71,7 @@ void writeToProducesAwsChunkedFraming() throws Exception { void contentLengthMatchesActualEncoding() throws Exception { for (int size : new int[] {1, 15, 16, 255, 256, 1024, 16384, 262144}) { byte[] payload = new byte[size]; - var chunked = new AwsChunkedDataStream(DataStream.ofBytes(payload, null)); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(payload, null)); var out = new ByteArrayOutputStream(); chunked.writeTo(out); @@ -75,49 +83,54 @@ void contentLengthMatchesActualEncoding() throws Exception { @Test void subscribeProducesSameBytesAsWriteTo() throws Exception { byte[] payload = "hello world".getBytes(StandardCharsets.UTF_8); - var chunked = new AwsChunkedDataStream(DataStream.ofBytes(payload, null)); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(payload, null)); var sync = new ByteArrayOutputStream(); chunked.writeTo(sync); - var collected = new ByteArrayOutputStream(); - var done = new CompletableFuture(); - chunked.subscribe(new Flow.Subscriber<>() { - @Override - public void onSubscribe(Flow.Subscription s) { - s.request(Long.MAX_VALUE); - } + var frames = subscribe(chunked); - @Override - public void onNext(ByteBuffer buf) { - byte[] tmp = new byte[buf.remaining()]; - buf.duplicate().get(tmp); - try { - collected.write(tmp); - } catch (Exception e) { - throw new RuntimeException(e); - } - } + assertArrayEquals(sync.toByteArray(), concat(frames.frames)); + assertThat(frames.calls, equalTo(List.of("next", "next", "next", "complete"))); + } - @Override - public void onError(Throwable t) { - done.completeExceptionally(t); - } + @Test + void subscribeDoesNotCopyAnUnpooledBody() { + byte[] payload = "zero copy".getBytes(StandardCharsets.UTF_8); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(payload, null)); - @Override - public void onComplete() { - done.complete(null); - } + var frames = subscribe(chunked).frames; + + assertThat(frames.size(), equalTo(3)); + assertThat(frames.get(1).array(), sameInstance(payload)); + } + + @Test + void forEachBufferAndAsByteBufferMirrorWriteTo() throws Exception { + byte[] payload = "hello buffers".getBytes(StandardCharsets.UTF_8); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(payload, null)); + + var written = new ByteArrayOutputStream(); + chunked.writeTo(written); + + var visited = new ByteArrayOutputStream(); + chunked.forEachBuffer(b -> { + byte[] tmp = new byte[b.remaining()]; + b.get(tmp); + visited.writeBytes(tmp); }); - done.get(); + assertArrayEquals(written.toByteArray(), visited.toByteArray()); - assertArrayEquals(sync.toByteArray(), collected.toByteArray()); + var first = chunked.asByteBuffer(); + assertArrayEquals(written.toByteArray(), first.array()); + assertThat(chunked.asByteBuffer().array(), sameInstance(first.array())); + assertArrayEquals(written.toByteArray(), chunked.asInputStream().readAllBytes()); } @Test void replayableProducesIdenticalEncodingAcrossCalls() throws Exception { byte[] payload = "stable".getBytes(StandardCharsets.UTF_8); - var chunked = new AwsChunkedDataStream(DataStream.ofBytes(payload, null)); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(payload, null)); var first = new ByteArrayOutputStream(); chunked.writeTo(first); @@ -130,7 +143,7 @@ void replayableProducesIdenticalEncodingAcrossCalls() throws Exception { @Test void emptyBodyComputesAllZeroCrc() throws Exception { - var chunked = new AwsChunkedDataStream(DataStream.ofBytes(new byte[0], null)); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(new byte[0], null)); var out = new ByteArrayOutputStream(); chunked.writeTo(out); // CRC32 of empty bytes is 0 → base64("\0\0\0\0") == "AAAAAA==". @@ -141,37 +154,22 @@ void emptyBodyComputesAllZeroCrc() throws Exception { @Test void rejectsNonReplayableInput() { var nonReplayable = DataStream.ofInputStream(new ByteArrayInputStream("x".getBytes())); - assertThrows(IllegalArgumentException.class, () -> new AwsChunkedDataStream(nonReplayable)); + assertThrows(IllegalArgumentException.class, () -> AwsChunkedDataStream.of(nonReplayable)); } @Test void decodedLengthReflectsInnerLength() { - var chunked = new AwsChunkedDataStream(DataStream.ofBytes(new byte[42], null)); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(new byte[42], null)); assertThat(chunked.decodedLength(), equalTo(42L)); } - @Test - void asInputStreamMirrorsWriteTo() throws Exception { - byte[] payload = new byte[256]; - var chunked = new AwsChunkedDataStream(DataStream.ofBytes(payload, null)); - - var fromStream = chunked.asInputStream().readAllBytes(); - var fromWriteTo = new ByteArrayOutputStream(); - chunked.writeTo(fromWriteTo); - - assertArrayEquals(fromWriteTo.toByteArray(), fromStream); - } - @Test void singleSubscriberCannotBeRequestedTwice() { - var chunked = new AwsChunkedDataStream(DataStream.ofBytes(new byte[4], null)); - var calls = new AtomicReference<>(new ArrayList()); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(new byte[4], null)); + var calls = new ArrayList(); chunked.subscribe(new Flow.Subscriber<>() { - Flow.Subscription sub; - @Override public void onSubscribe(Flow.Subscription s) { - sub = s; s.request(Long.MAX_VALUE); // Second request after first should be a no-op. s.request(Long.MAX_VALUE); @@ -179,38 +177,341 @@ public void onSubscribe(Flow.Subscription s) { @Override public void onNext(ByteBuffer item) { - calls.get().add("next"); + calls.add("next"); } @Override public void onError(Throwable throwable) { - calls.get().add("error"); + calls.add("error"); } @Override public void onComplete() { - calls.get().add("complete"); + calls.add("complete"); } }); // Frames: header + body + footer + complete = 3 onNext + 1 onComplete. - assertThat(calls.get(), equalTo(List.of("next", "next", "next", "complete"))); + assertThat(calls, equalTo(List.of("next", "next", "next", "complete"))); + } + + @Test + void plainInnerIsNotPooled() { + assertThat(AwsChunkedDataStream.of(DataStream.ofBytes(new byte[4])), not(instanceOf(PooledDataStream.class))); } @Test - public void writeToDoesNotAliasAPooledBody() throws Exception { + void writeToAndForEachBufferDoNotAliasAPooledBody() throws Exception { byte[] payload = "hello chunked world".getBytes(StandardCharsets.UTF_8); var recycled = new AtomicReference(); PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, "text/plain", recycled::set); - var encoded = new ByteArrayOutputStream(); - new AwsChunkedDataStream(pooled).writeTo(encoded); + var chunked = AwsChunkedDataStream.of(pooled); + assertThat(chunked, instanceOf(PooledDataStream.class)); + assertThat(chunked.contentType(), equalTo("text/plain")); - var expected = new ByteArrayOutputStream(); - new AwsChunkedDataStream(DataStream.ofBytes(payload)).writeTo(expected); - assertArrayEquals(expected.toByteArray(), encoded.toByteArray()); + var written = new ByteArrayOutputStream(); + chunked.writeTo(written); + assertArrayEquals(encode(payload), written.toByteArray()); + var visited = new ByteArrayOutputStream(); + chunked.forEachBuffer(b -> { + byte[] tmp = new byte[b.remaining()]; + b.get(tmp); + visited.writeBytes(tmp); + }); + assertArrayEquals(encode(payload), visited.toByteArray()); + + var lease = (PooledDataStream) chunked; + assertThat(lease.escaped(), is(false)); assertThat(pooled.escaped(), is(false)); + assertThat(lease.isReleased(), is(false)); + + lease.release(); + assertThat(lease.isReleased(), is(true)); + assertThat(pooled.isReleased(), is(true)); + assertThat(recycled.get(), sameInstance(payload)); + } + + @Test + void readsAfterReleaseThrowBeforeEmittingAnyByte() throws Exception { + byte[] payload = "late".getBytes(StandardCharsets.UTF_8); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, null, array -> {}); + var chunked = AwsChunkedDataStream.of(pooled); + var out = new ByteArrayOutputStream(); + chunked.writeTo(out); + pooled.release(); + + assertThrows(IllegalStateException.class, () -> chunked.writeTo(out)); + assertArrayEquals(encode(payload), out.toByteArray()); + + var visited = new ArrayList(); + assertThrows(IllegalStateException.class, () -> chunked.forEachBuffer(visited::add)); + assertThat(visited.isEmpty(), is(true)); + assertThrows(IllegalStateException.class, chunked::asByteBuffer); + assertThrows(IllegalStateException.class, () -> chunked.subscribe(new CollectingSubscriber())); + assertThrows(IllegalStateException.class, ((PooledDataStream) chunked)::copy); + } + + @Test + void subscribeOnAPooledBodyPublishesTheSharedPrivateCopy() { + byte[] payload = "publish me".getBytes(StandardCharsets.UTF_8); + var recycled = new AtomicReference(); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, null, recycled::set); + var chunked = AwsChunkedDataStream.of(pooled); + + var frames = subscribe(chunked).frames; + + assertArrayEquals(encode(payload), concat(frames)); + assertThat(frames.size(), equalTo(1)); + assertThat(frames.get(0).array(), sameInstance(chunked.asByteBuffer().array())); + assertThat(pooled.escaped(), is(false)); + assertThat(((PooledDataStream) chunked).escaped(), is(true)); + + ((PooledDataStream) chunked).release(); + assertThat(recycled.get(), sameInstance(payload)); + } + + @Test + void asByteBufferOnAPooledBodyIsOnePrivateCopy() throws Exception { + byte[] payload = "copy me once".getBytes(StandardCharsets.UTF_8); + var recycled = new AtomicReference(); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, null, recycled::set); + var chunked = AwsChunkedDataStream.of(pooled); + + var encoded = chunked.asByteBuffer(); + assertArrayEquals(encode(payload), encoded.array()); + assertThat(chunked.asByteBuffer().array(), sameInstance(encoded.array())); + assertArrayEquals(encode(payload), chunked.asInputStream().readAllBytes()); + assertThat(((PooledDataStream) chunked).escaped(), is(true)); + assertThat(pooled.escaped(), is(false)); + + ((PooledDataStream) chunked).release(); + assertThat(recycled.get(), sameInstance(payload)); + } + + @Test + void cachedReadsAlsoThrowAfterRelease() { + byte[] payload = "strict".getBytes(StandardCharsets.UTF_8); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, null, array -> {}); + var chunked = AwsChunkedDataStream.of(pooled); + var view = chunked.asByteBuffer(); + pooled.release(); + + assertArrayEquals(encode(payload), view.array()); + assertThrows(IllegalStateException.class, chunked::asByteBuffer); + assertThrows(IllegalStateException.class, chunked::asInputStream); + assertThrows(IllegalStateException.class, () -> chunked.subscribe(new CollectingSubscriber())); + assertThrows(IllegalStateException.class, ((PooledDataStream) chunked)::copy); + } + + @Test + void copyOfAPooledBodyIsEncodedAndUnpooled() throws Exception { + byte[] payload = "retain me".getBytes(StandardCharsets.UTF_8); + var recycled = new AtomicReference(); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, "text/plain", recycled::set); + var chunked = (PooledDataStream) AwsChunkedDataStream.of(pooled); + + var copy = chunked.copy(); + assertThat(chunked.escaped(), is(false)); + chunked.release(); + + assertThat(copy, not(instanceOf(PooledDataStream.class))); + assertThat(copy.contentType(), equalTo("text/plain")); + assertThat(copy.contentLength(), equalTo(chunked.contentLength())); + var out = new ByteArrayOutputStream(); + copy.writeTo(out); + assertArrayEquals(encode(payload), out.toByteArray()); + assertThat(recycled.get(), sameInstance(payload)); + } + + @Test + void copyReusesAnExistingEncodedCopy() { + byte[] payload = "already encoded".getBytes(StandardCharsets.UTF_8); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, "text/plain", array -> {}); + var chunked = AwsChunkedDataStream.of(pooled); + + var encoded = chunked.asByteBuffer(); + var copy = ((PooledDataStream) chunked).copy(); + + assertThat(copy, not(instanceOf(PooledDataStream.class))); + assertThat(copy.contentType(), equalTo("text/plain")); + assertThat(copy.contentLength(), equalTo(chunked.contentLength())); + assertThat(copy.asByteBuffer().array(), sameInstance(encoded.array())); + } + + @Test + void subscribePublishesTheCachedCopyWhenPresent() { + byte[] payload = "cached".getBytes(StandardCharsets.UTF_8); + var chunked = AwsChunkedDataStream.of(DataStream.ofBytes(payload, null)); + + var encoded = chunked.asByteBuffer(); + var frames = subscribe(chunked).frames; + + assertThat(frames.size(), equalTo(1)); + assertThat(frames.get(0).array(), sameInstance(encoded.array())); + } + + @Test + void writeToChannelDoesNotAliasAPooledBody() throws Exception { + byte[] payload = "down the channel".getBytes(StandardCharsets.UTF_8); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, null, array -> {}); + var chunked = AwsChunkedDataStream.of(pooled); + + var out = new ByteArrayOutputStream(); + chunked.writeTo(Channels.newChannel(out)); + + assertArrayEquals(encode(payload), out.toByteArray()); + assertThat(((PooledDataStream) chunked).escaped(), is(false)); + pooled.release(); - assertThat(recycled.get(), is(payload)); + var late = new ByteArrayOutputStream(); + assertThrows(IllegalStateException.class, () -> chunked.writeTo(Channels.newChannel(late))); + assertThat(late.size(), equalTo(0)); + } + + @Test + void crcIsComputedOnceAcrossReads() throws Exception { + var inner = new CountingBody("count me".getBytes(StandardCharsets.UTF_8)); + var chunked = AwsChunkedDataStream.of(inner); + + chunked.writeTo(new ByteArrayOutputStream()); + chunked.writeTo(new ByteArrayOutputStream()); + chunked.forEachBuffer(b -> {}); + chunked.asByteBuffer(); + subscribe(chunked); + + // The three visits compute the CRC, call forEachBuffer, and copy the encoded body. + assertThat(inner.visits, equalTo(3)); + } + + @Test + void abandonForwardsToThePooledBody() { + byte[] payload = "abandon".getBytes(StandardCharsets.UTF_8); + var recycled = new AtomicReference(); + PooledDataStream pooled = DataStream.ofLeasedBytes(payload, 0, payload.length, null, recycled::set); + var chunked = (PooledDataStream) AwsChunkedDataStream.of(pooled); + + chunked.abandon(); + assertThat(chunked.escaped(), is(true)); + chunked.release(); + + assertThat(chunked.isReleased(), is(true)); + assertThat(recycled.get(), nullValue()); + } + + private static final class CountingBody implements DataStream { + private final DataStream delegate; + int visits; + + CountingBody(byte[] bytes) { + this.delegate = DataStream.ofBytes(bytes); + } + + @Override + public long contentLength() { + return delegate.contentLength(); + } + + @Override + public String contentType() { + return null; + } + + @Override + public boolean isReplayable() { + return true; + } + + @Override + public boolean isAvailable() { + return true; + } + + @Override + public boolean hasByteBuffer() { + return true; + } + + @Override + public ByteBuffer asByteBuffer() { + return delegate.asByteBuffer(); + } + + @Override + public InputStream asInputStream() { + return delegate.asInputStream(); + } + + @Override + public void writeTo(OutputStream out) throws IOException { + delegate.writeTo(out); + } + + @Override + public void forEachBuffer(Consumer action) { + visits++; + delegate.forEachBuffer(action); + } + } + + private static CollectingSubscriber subscribe(DataStream stream) { + var subscriber = new CollectingSubscriber(); + stream.subscribe(subscriber); + assertThat(subscriber.calls.contains("error"), is(false)); + return subscriber; + } + + private static final class CollectingSubscriber implements Flow.Subscriber { + final List calls = new ArrayList<>(); + final List frames = new ArrayList<>(); + + @Override + public void onSubscribe(Flow.Subscription s) { + s.request(Long.MAX_VALUE); + } + + @Override + public void onNext(ByteBuffer item) { + calls.add("next"); + frames.add(item); + } + + @Override + public void onError(Throwable throwable) { + calls.add("error"); + } + + @Override + public void onComplete() { + calls.add("complete"); + } + } + + private static byte[] concat(List frames) { + var out = new ByteArrayOutputStream(); + for (var frame : frames) { + var view = frame.duplicate(); + byte[] tmp = new byte[view.remaining()]; + view.get(tmp); + out.writeBytes(tmp); + } + return out.toByteArray(); + } + + private static byte[] encode(byte[] payload) { + var crc = new CRC32(); + crc.update(payload); + long v = crc.getValue(); + byte[] crcBytes = { + (byte) ((v >>> 24) & 0xff), + (byte) ((v >>> 16) & 0xff), + (byte) ((v >>> 8) & 0xff), + (byte) (v & 0xff)}; + var out = new ByteArrayOutputStream(); + out.writeBytes((Long.toHexString(payload.length) + "\r\n").getBytes(StandardCharsets.US_ASCII)); + out.writeBytes(payload); + out.writeBytes(("\r\n0\r\nx-amz-checksum-crc32:" + Base64.getEncoder().encodeToString(crcBytes) + "\r\n\r\n") + .getBytes(StandardCharsets.US_ASCII)); + return out.toByteArray(); } } diff --git a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4CaseGenerator.java b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4CaseGenerator.java index 62e27a006c..8bd8fd0fe4 100644 --- a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4CaseGenerator.java +++ b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4CaseGenerator.java @@ -444,30 +444,45 @@ private static String body(Source s, ModifiableHttpRequest request) { request.setBody(DataStream.ofBytes(new byte[0])); return "empty"; } - if (k < 60) { + if (k < 55) { request.setBody(DataStream.ofBytes(bytes)); return "heap[" + bytes.length + "]"; } - if (k < 70) { + if (k < 63) { var direct = ByteBuffer.allocateDirect(bytes.length); direct.put(bytes).flip(); request.setBody(DataStream.ofByteBuffer(direct)); return "direct[" + bytes.length + "]"; } - if (k < 80) { + if (k < 71) { int pad = s.nextInt(10); byte[] padded = new byte[bytes.length + pad + 5]; System.arraycopy(bytes, 0, padded, pad, bytes.length); request.setBody(DataStream.ofByteBuffer(ByteBuffer.wrap(padded, pad, bytes.length).asReadOnlyBuffer())); return "readOnly[" + bytes.length + "@" + pad + "]"; } - if (k < 90) { + if (k < 79) { int pad = s.nextInt(10); byte[] padded = new byte[bytes.length + pad + 5]; System.arraycopy(bytes, 0, padded, pad, bytes.length); request.setBody(DataStream.ofByteBuffer(ByteBuffer.wrap(padded, pad, bytes.length).slice())); return "slice[" + bytes.length + "@" + pad + "]"; } + if (k < 87) { + int pad = s.nextInt(10); + byte[] padded = new byte[bytes.length + pad + 5]; + System.arraycopy(bytes, 0, padded, pad, bytes.length); + request.setBody(DataStream.ofLeasedBytes(padded, pad, bytes.length, null, array -> {})); + return "leased[" + bytes.length + "@" + pad + "]"; + } + if (k < 94) { + int split = bytes.length == 0 ? 0 : s.nextInt(bytes.length + 1); + var parts = List.of( + ByteBuffer.wrap(bytes, 0, split), + ByteBuffer.wrap(bytes, split, bytes.length - split)); + request.setBody(DataStream.ofByteBuffers(parts, bytes.length, null)); + return "multi[" + split + "+" + (bytes.length - split) + "]"; + } request.setBody(DataStream.ofInputStream(new ByteArrayInputStream(bytes))); return "unknownLength[" + bytes.length + "]"; } diff --git a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Differential.java b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Differential.java index f2d0156c8e..7b13247788 100644 --- a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Differential.java +++ b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Differential.java @@ -13,6 +13,7 @@ import java.util.concurrent.Callable; import software.amazon.smithy.java.auth.api.SignResult; import software.amazon.smithy.java.http.api.HttpRequest; +import software.amazon.smithy.java.io.datastream.PooledDataStream; final class SigV4Differential { @@ -33,6 +34,9 @@ static void check(SigV4CaseGenerator.SignCase c, SigV4Signer signer) { Outcome first = Outcome.of(() -> signer.sign(c.request(), c.identity(), c.properties())); Outcome second = Outcome.of(() -> signer.sign(c.request(), c.identity(), c.properties())); assertThat(second).as("re-signing with the same signer: %s", c).isEqualTo(first); + if (c.request().body() instanceof PooledDataStream pooled) { + assertThat(pooled.escaped()).as("signing must not alias a pooled body: %s", c).isFalse(); + } Outcome reference = Outcome.of(() -> StringSigV4Reference.sign(c.request(), c.identity(), c.properties())); assertThat(first).as("string reference: %s", c).isEqualTo(reference); diff --git a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4SignerTest.java b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4SignerTest.java index aeff0fc2dd..29530a1a69 100644 --- a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4SignerTest.java +++ b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4SignerTest.java @@ -8,20 +8,28 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.is; import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.sameInstance; import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; import java.time.Clock; import java.time.Instant; import java.time.ZoneOffset; +import java.util.HexFormat; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import software.amazon.smithy.java.aws.auth.api.identity.AwsCredentialsIdentity; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.http.api.HttpHeaders; import software.amazon.smithy.java.http.api.HttpRequest; import software.amazon.smithy.java.http.api.HttpVersion; +import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; class SigV4SignerTest { @@ -92,4 +100,44 @@ void ignoresPreexistingAuthorizationHeader() { assertThat(staleAuth, not(containsString("authorization"))); assertThat(staleAuth, equalTo(cleanAuth)); } + + @Test + void signingReadsAPooledBodyWithoutAliasingIt() throws Exception { + byte[] payload = "{\"hello\":\"world\"}".getBytes(StandardCharsets.UTF_8); + var recycled = new AtomicReference(); + PooledDataStream pooled = + DataStream.ofLeasedBytes(payload, 0, payload.length, "application/json", recycled::set); + var context = Context.create() + .put(SigV4Settings.REGION, "us-east-1") + .put(SigV4Settings.SIGNING_NAME, "s3") + .put(SigV4Settings.CLOCK, Clock.fixed(Instant.parse("2026-06-01T23:59:59Z"), ZoneOffset.UTC)); + var identity = AwsCredentialsIdentity.create("AKID", "secret"); + + HttpRequest signedPooled; + HttpRequest signedPlain; + try (var signer = SigV4Signer.create()) { + signedPooled = signer.sign(bodyRequest(pooled), identity, context).signedRequest(); + signedPlain = signer.sign(bodyRequest(DataStream.ofBytes(payload, "application/json")), identity, context) + .signedRequest(); + } + + assertThat(pooled.escaped(), is(false)); + assertThat(pooled.isReleased(), is(false)); + assertThat(recycled.get(), is((byte[]) null)); + assertThat(signedPooled.body(), sameInstance(pooled)); + var sha256 = HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(payload)); + assertThat(signedPooled.headers().firstValue("x-amz-content-sha256"), equalTo(sha256)); + assertThat(signedPooled.headers().firstValue("authorization"), + equalTo(signedPlain.headers().firstValue("authorization"))); + } + + private static HttpRequest bodyRequest(DataStream body) { + return HttpRequest.create() + .setMethod("PUT") + .setHttpVersion(HttpVersion.HTTP_1_1) + .setUri(URI.create("https://bucket.s3.amazonaws.com/key")) + .setHeaders(HttpHeaders.of(Map.of("Host", List.of("bucket.s3.amazonaws.com")))) + .setBody(body) + .toUnmodifiable(); + } } diff --git a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResourcesPropertyTest.java b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResourcesPropertyTest.java index f058125b5c..fb7de08df9 100644 --- a/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResourcesPropertyTest.java +++ b/aws/aws-sigv4/src/test/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResourcesPropertyTest.java @@ -155,6 +155,12 @@ static void assertHash(SigningResources r, byte[] data, int pad, int kind) throw r.hash(stream, viaStream); assertThat(viaStream).as("hash via stream kind %d size %d", kind, data.length).isEqualTo(expected); assertThat(first.position()).as("stream view position must be untouched").isEqualTo(position); + + var leased = DataStream.ofLeasedBytes(padded, pad, data.length, null, array -> {}); + byte[] viaLease = new byte[32]; + r.hash(leased, viaLease); + assertThat(viaLease).as("hash via lease size %d", data.length).isEqualTo(expected); + assertThat(leased.escaped()).as("hashing must not alias the lease").isFalse(); } @Test diff --git a/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java b/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java index 142cd2a34e..2e532cd101 100644 --- a/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java +++ b/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java @@ -122,6 +122,7 @@ public O deserializ // If the payload is empty, then use "{}" as the default empty payload. if (content.contentLength() == 0) { + discardBody(response); return codec.deserializeShape(EMPTY_PAYLOAD, builder); } diff --git a/aws/client/aws-client-awsquery/build.gradle.kts b/aws/client/aws-client-awsquery/build.gradle.kts index b182c8ffdc..698a42d121 100644 --- a/aws/client/aws-client-awsquery/build.gradle.kts +++ b/aws/client/aws-client-awsquery/build.gradle.kts @@ -18,6 +18,8 @@ dependencies { // Protocol test dependencies testImplementation(libs.smithy.aws.protocol.tests) + testImplementation(project(":client:dynamic-client")) + testImplementation(project(":client:client-mock-plugin")) } protocolTestRuns { diff --git a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java index 67c3792c29..3afe0c427d 100644 --- a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java +++ b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java @@ -108,6 +108,7 @@ public O deserializ var content = response.body(); if (content.contentLength() == 0 || operation.outputSchema().hasTrait(TraitKey.UNIT_TYPE_TRAIT)) { + discardBody(response); return builder.build(); } diff --git a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java index 1a89910f42..dfef0a878c 100644 --- a/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java +++ b/aws/client/aws-client-awsquery/src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java @@ -107,7 +107,8 @@ public O deserializ var builder = operation.outputBuilder(); var content = response.body(); - if (content.contentLength() == 0) { + if (content.contentLength() == 0 || operation.outputSchema().hasTrait(TraitKey.UNIT_TYPE_TRAIT)) { + discardBody(response); return builder.build(); } diff --git a/aws/client/aws-client-awsquery/src/test/java/software/amazon/smithy/java/aws/client/awsquery/QueryResponseBodyReleaseTest.java b/aws/client/aws-client-awsquery/src/test/java/software/amazon/smithy/java/aws/client/awsquery/QueryResponseBodyReleaseTest.java new file mode 100644 index 0000000000..1e05b9fc48 --- /dev/null +++ b/aws/client/aws-client-awsquery/src/test/java/software/amazon/smithy/java/aws/client/awsquery/QueryResponseBodyReleaseTest.java @@ -0,0 +1,207 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.aws.client.awsquery; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.is; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.util.Map; +import org.junit.jupiter.api.Test; +import software.amazon.smithy.java.client.core.ClientProtocol; +import software.amazon.smithy.java.client.core.auth.scheme.AuthSchemeResolver; +import software.amazon.smithy.java.client.http.mock.MockPlugin; +import software.amazon.smithy.java.client.http.mock.MockQueue; +import software.amazon.smithy.java.core.serde.document.Document; +import software.amazon.smithy.java.dynamicclient.DynamicClient; +import software.amazon.smithy.java.endpoints.EndpointResolver; +import software.amazon.smithy.java.http.api.HttpRequest; +import software.amazon.smithy.java.http.api.HttpResponse; +import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.model.Model; +import software.amazon.smithy.model.shapes.ShapeId; + +/** + * A pooled HTTP connection is only returned when the response body is consumed or discarded, so every + * path that skips deserialization must still discard the body. + */ +class QueryResponseBodyReleaseTest { + + private static final ShapeId METRICS = ShapeId.from("smithy.example#Metrics"); + private static final ShapeId INSTANCES = ShapeId.from("smithy.example#Instances"); + private static final Model MODEL = Model.assembler() + .addUnparsedModel("test.smithy", """ + $version: "2" + namespace smithy.example + + use aws.protocols#awsQuery + use aws.protocols#ec2Query + + @awsQuery + @xmlNamespace(uri: "https://example.com") + service Metrics { + version: "2010-08-01" + operations: [PutMetric, GetMetric] + } + + @ec2Query + @xmlNamespace(uri: "https://example.com") + service Instances { + version: "2016-11-15" + operations: [PutMetric, GetMetric] + } + + operation PutMetric { + input := { + name: String + } + output: Unit + } + + operation GetMetric { + input := { + name: String + } + output := { + value: Integer + } + } + """) + .discoverModels() + .assemble() + .unwrap(); + + private static final String PUT_RESPONSE = "" + + "1"; + + @Test + void awsQueryDiscardsTheBodyOfAUnitOutput() { + var body = new RecordingBody(DataStream.ofString(PUT_RESPONSE)); + + call(METRICS, new AwsQueryClientProtocol(METRICS, "2010-08-01"), "PutMetric", body); + + assertThat(body.discarded, is(true)); + assertThat(body.reads, equalTo(0)); + } + + @Test + void awsQueryDiscardsAnEmptyBody() { + var body = new RecordingBody(DataStream.ofBytes(new byte[0])); + + call(METRICS, new AwsQueryClientProtocol(METRICS, "2010-08-01"), "GetMetric", body); + + assertThat(body.discarded, is(true)); + } + + @Test + void awsQueryConsumesTheBodyOfARealOutput() { + var body = new RecordingBody(DataStream.ofString("" + + "7")); + + var output = call(METRICS, new AwsQueryClientProtocol(METRICS, "2010-08-01"), "GetMetric", body); + + assertThat(output.getMember("value").asInteger(), equalTo(7)); + assertThat(body.reads, equalTo(1)); + } + + @Test + void ec2QueryDiscardsTheBodyOfAUnitOutput() { + var body = new RecordingBody(DataStream.ofString(PUT_RESPONSE)); + + call(INSTANCES, new Ec2QueryClientProtocol(INSTANCES, "2016-11-15"), "PutMetric", body); + + assertThat(body.discarded, is(true)); + assertThat(body.reads, equalTo(0)); + } + + @Test + void ec2QueryDiscardsAnEmptyBody() { + var body = new RecordingBody(DataStream.ofBytes(new byte[0])); + + call(INSTANCES, new Ec2QueryClientProtocol(INSTANCES, "2016-11-15"), "GetMetric", body); + + assertThat(body.discarded, is(true)); + } + + private static Document call( + ShapeId service, + ClientProtocol protocol, + String operation, + DataStream body + ) { + var queue = new MockQueue().enqueue(HttpResponse.create() + .setStatusCode(200) + .setBody(body) + .toUnmodifiable()); + var client = DynamicClient.builder() + .serviceId(service) + .model(MODEL) + .protocol(protocol) + .addPlugin(MockPlugin.builder().addQueue(queue).build()) + .endpointResolver(EndpointResolver.staticEndpoint("https://localhost:8081")) + .authSchemeResolver(AuthSchemeResolver.NO_AUTH) + .build(); + var output = client.call(operation, Document.ofObject(Map.of("name", "n"))); + assertThat(queue.remaining(), equalTo(0)); + return output; + } + + private static final class RecordingBody implements DataStream { + private final DataStream delegate; + boolean discarded; + int reads; + + RecordingBody(DataStream delegate) { + this.delegate = delegate; + } + + @Override + public long contentLength() { + return delegate.contentLength(); + } + + @Override + public String contentType() { + return "text/xml"; + } + + @Override + public boolean isReplayable() { + return true; + } + + @Override + public boolean isAvailable() { + return true; + } + + @Override + public boolean hasByteBuffer() { + return true; + } + + @Override + public ByteBuffer asByteBuffer() { + reads++; + return delegate.asByteBuffer(); + } + + @Override + public InputStream asInputStream() { + reads++; + return delegate.asInputStream(); + } + + @Override + public void discard() throws IOException { + discarded = true; + delegate.discard(); + } + } +} diff --git a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpClientProtocol.java b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpClientProtocol.java index dcc29d2a56..4e064f9cfa 100644 --- a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpClientProtocol.java +++ b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpClientProtocol.java @@ -5,6 +5,8 @@ package software.amazon.smithy.java.client.http; +import java.io.IOException; +import java.io.UncheckedIOException; import software.amazon.smithy.java.client.core.ClientProtocol; import software.amazon.smithy.java.client.core.MessageExchange; import software.amazon.smithy.java.endpoints.Endpoint; @@ -34,6 +36,19 @@ public MessageExchange messageExchange() { return HttpMessageExchange.INSTANCE; } + /** + * Releases a response body that will not be read, so a pooled connection returns to its pool. + * + *

Call this on every path that completes deserialization without consuming the body. + */ + protected static void discardBody(HttpResponse response) { + try { + response.body().discard(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + @Override public PooledDataStream pooledRequestBody(HttpRequest request) { return request.body() instanceof PooledDataStream pooled ? pooled : null; diff --git a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/JavaHttpClientTransport.java b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/JavaHttpClientTransport.java index 294e6d5a5a..9238ac519e 100644 --- a/client/client-http/src/main/java/software/amazon/smithy/java/client/http/JavaHttpClientTransport.java +++ b/client/client-http/src/main/java/software/amazon/smithy/java/client/http/JavaHttpClientTransport.java @@ -29,6 +29,7 @@ import software.amazon.smithy.java.http.api.HttpResponse; import software.amazon.smithy.java.http.api.HttpVersion; import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.java.logging.InternalLogger; /** @@ -160,7 +161,7 @@ public MessageExchange messageExchange() { @Override public HttpResponse send(Context context, HttpRequest request) { - return sendRequest(createJavaRequest(context, request)); + return sendRequest(request, createJavaRequest(context, request)); } /** @@ -181,12 +182,16 @@ private java.net.http.HttpRequest createJavaRequest(Context context, HttpRequest return httpRequestBuilder.build(); } - private HttpResponse sendRequest(java.net.http.HttpRequest request) { + private HttpResponse sendRequest(HttpRequest request, java.net.http.HttpRequest javaRequest) { java.net.http.HttpResponse res = null; try { - res = client.send(request, responseBodyHandler); + res = client.send(javaRequest, responseBodyHandler); return createSmithyResponse(res); } catch (IOException | InterruptedException | RuntimeException e) { + // The JDK can still subscribe to the body on its own threads after send() fails. Prevent recycling. + if (request.body() instanceof PooledDataStream pooled) { + pooled.abandon(); + } if (res != null) { try { res.body().close(); diff --git a/client/client-http/src/test/java/software/amazon/smithy/java/client/http/JavaHttpClientTest.java b/client/client-http/src/test/java/software/amazon/smithy/java/client/http/JavaHttpClientTest.java index 9c66f638c5..91ff862990 100644 --- a/client/client-http/src/test/java/software/amazon/smithy/java/client/http/JavaHttpClientTest.java +++ b/client/client-http/src/test/java/software/amazon/smithy/java/client/http/JavaHttpClientTest.java @@ -10,11 +10,14 @@ import static org.hamcrest.Matchers.empty; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.nullValue; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.ByteArrayOutputStream; +import java.io.IOException; import java.net.Authenticator; import java.net.CookieHandler; import java.net.ProxySelector; @@ -24,7 +27,9 @@ import java.net.http.HttpHeaders; import java.net.http.HttpRequest; import java.net.http.HttpResponse; +import java.net.http.HttpTimeoutException; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; @@ -34,6 +39,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.concurrent.Flow; +import java.util.concurrent.atomic.AtomicReference; import javax.net.ssl.SSLContext; import javax.net.ssl.SSLParameters; import javax.net.ssl.SSLSession; @@ -42,10 +48,12 @@ import org.junit.jupiter.api.Test; import software.amazon.smithy.java.aws.client.awsjson.AwsJson1Protocol; import software.amazon.smithy.java.client.core.ClientConfig; +import software.amazon.smithy.java.client.core.error.TransportException; import software.amazon.smithy.java.context.Context; import software.amazon.smithy.java.core.schema.Schema; import software.amazon.smithy.java.endpoints.EndpointResolver; import software.amazon.smithy.java.io.datastream.DataStream; +import software.amazon.smithy.java.io.datastream.PooledDataStream; import software.amazon.smithy.model.shapes.ShapeId; public class JavaHttpClientTest { @@ -196,6 +204,26 @@ public void publishesBodyViaSubscribeForKnownLengthBodies() { assertThat(client.capturedBytes, equalTo(payload)); } + @Test + public void abandonsAPooledBodyWhenTheJdkClientFails() { + byte[] payload = "{}".getBytes(StandardCharsets.UTF_8); + var recycled = new AtomicReference(); + PooledDataStream body = DataStream.ofLeasedBytes(payload, 0, payload.length, "application/json", recycled::set); + var transport = new JavaHttpClientTransport(new FailingHttpClient()); + var request = software.amazon.smithy.java.http.api.HttpRequest.create() + .setUri(URI.create("http://localhost/test")) + .setMethod("POST") + .setBody(body) + .toUnmodifiable(); + + assertThrows(TransportException.class, () -> transport.send(Context.create(), request)); + + assertThat(body.escaped(), equalTo(true)); + body.release(); + assertThat(body.isReleased(), equalTo(true)); + assertThat(recycled.get(), nullValue()); + } + @Test public void streamsResponseBodyThroughCustomBodySubscriber() throws Exception { byte[] payload = "streamed response body".getBytes(); @@ -302,7 +330,8 @@ Optional capturedVersion() { } @Override - public HttpResponse send(HttpRequest request, HttpResponse.BodyHandler responseBodyHandler) { + public HttpResponse send(HttpRequest request, HttpResponse.BodyHandler responseBodyHandler) + throws IOException { capturedBytes = collect(request.bodyPublisher().orElseThrow()); capturedVersion = request.version(); var responseInfo = new HttpResponse.ResponseInfo() { @@ -380,7 +409,11 @@ public CompletableFuture> sendAsync( HttpRequest request, HttpResponse.BodyHandler responseBodyHandler ) { - return CompletableFuture.completedFuture(send(request, responseBodyHandler)); + try { + return CompletableFuture.completedFuture(send(request, responseBodyHandler)); + } catch (IOException e) { + return CompletableFuture.failedFuture(e); + } } @Override @@ -465,6 +498,14 @@ public void onComplete() {} } } + private static final class FailingHttpClient extends CapturingHttpClient { + @Override + public HttpResponse send(HttpRequest request, HttpResponse.BodyHandler responseBodyHandler) + throws IOException { + throw new HttpTimeoutException("request timed out"); + } + } + private static final class StreamingResponseHttpClient extends HttpClient { private final byte[] responseBytes; diff --git a/client/client-rpcv2/src/main/java/software/amazon/smithy/java/client/rpcv2/AbstractRpcV2ClientProtocol.java b/client/client-rpcv2/src/main/java/software/amazon/smithy/java/client/rpcv2/AbstractRpcV2ClientProtocol.java index e4de7d0fd2..e19538161b 100644 --- a/client/client-rpcv2/src/main/java/software/amazon/smithy/java/client/rpcv2/AbstractRpcV2ClientProtocol.java +++ b/client/client-rpcv2/src/main/java/software/amazon/smithy/java/client/rpcv2/AbstractRpcV2ClientProtocol.java @@ -145,6 +145,7 @@ public O deserializ var builder = operation.outputBuilder(); var content = response.body(); if (content.contentLength() == 0) { + discardBody(response); return builder.build(); } diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/ForwardingPooledDataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/ForwardingPooledDataStream.java new file mode 100644 index 0000000000..2d7f327787 --- /dev/null +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/ForwardingPooledDataStream.java @@ -0,0 +1,38 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.java.io.datastream; + +/** + * A pooled stream that reads another pooled stream's storage and forwards the lease to it. + * + *

Override {@link #escaped()} when the wrapper keeps a private copy of its own. + */ +public interface ForwardingPooledDataStream extends PooledDataStream { + /** + * @return the pooled stream whose storage this stream reads + */ + PooledDataStream lease(); + + @Override + default void release() { + lease().release(); + } + + @Override + default boolean isReleased() { + return lease().isReleased(); + } + + @Override + default void abandon() { + lease().abandon(); + } + + @Override + default boolean escaped() { + return lease().escaped(); + } +} diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/PooledDataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/PooledDataStream.java index a7e1d43c88..9fcda54619 100644 --- a/io/src/main/java/software/amazon/smithy/java/io/datastream/PooledDataStream.java +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/PooledDataStream.java @@ -23,7 +23,8 @@ *

Pooled streams are not thread-safe. * Complete all reads before release. If a reader can continue after the transport returns, call {@link #abandon()}. * - *

Wrappers that read pooled storage must implement this interface and forward its lifecycle methods. + *

Wrappers that read pooled storage must implement this interface. {@link ForwardingPooledDataStream} + * forwards the lifecycle methods to the wrapped lease. */ public interface PooledDataStream extends DataStream { /** diff --git a/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedPooledDataStream.java b/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedPooledDataStream.java index 20c796ba7b..0b9a8df29f 100644 --- a/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedPooledDataStream.java +++ b/io/src/main/java/software/amazon/smithy/java/io/datastream/WrappedPooledDataStream.java @@ -5,7 +5,7 @@ package software.amazon.smithy.java.io.datastream; -final class WrappedPooledDataStream extends WrappedDataStream implements PooledDataStream { +final class WrappedPooledDataStream extends WrappedDataStream implements ForwardingPooledDataStream { private final PooledDataStream pooled; @@ -15,23 +15,8 @@ final class WrappedPooledDataStream extends WrappedDataStream implements PooledD } @Override - public void release() { - pooled.release(); - } - - @Override - public boolean isReleased() { - return pooled.isReleased(); - } - - @Override - public void abandon() { - pooled.abandon(); - } - - @Override - public boolean escaped() { - return pooled.escaped(); + public PooledDataStream lease() { + return pooled; } @Override