From 39ab2f2f1ea326457709691a8f5b2f809053ff0d Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 21 Jul 2023 09:24:01 -0700 Subject: [PATCH 01/14] Refactor to common class AwsChunkedInputStream --- .../AwsSignedChunkedEncodingInputStream.java | 1 - .../io/AwsChunkedEncodingInputStream.java | 90 +++---------------- .../internal/io/AwsChunkedInputStream.java | 90 +++++++++++++++++++ ...uffer.java => UnderlyingStreamBuffer.java} | 6 +- 4 files changed, 104 insertions(+), 83 deletions(-) create mode 100644 core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsChunkedInputStream.java rename core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/{DecodedStreamBuffer.java => UnderlyingStreamBuffer.java} (93%) diff --git a/core/auth/src/main/java/software/amazon/awssdk/auth/signer/internal/chunkedencoding/AwsSignedChunkedEncodingInputStream.java b/core/auth/src/main/java/software/amazon/awssdk/auth/signer/internal/chunkedencoding/AwsSignedChunkedEncodingInputStream.java index 636fad74f9fc..3174eb7c6caa 100644 --- a/core/auth/src/main/java/software/amazon/awssdk/auth/signer/internal/chunkedencoding/AwsSignedChunkedEncodingInputStream.java +++ b/core/auth/src/main/java/software/amazon/awssdk/auth/signer/internal/chunkedencoding/AwsSignedChunkedEncodingInputStream.java @@ -40,7 +40,6 @@ @SdkInternalApi public final class AwsSignedChunkedEncodingInputStream extends AwsChunkedEncodingInputStream { - private static final String CRLF = "\r\n"; private static final String CHUNK_SIGNATURE_HEADER = ";chunk-signature="; private static final String CHECKSUM_SIGNATURE_HEADER = "x-amz-trailer-signature:"; private String previousChunkSignature; diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsChunkedEncodingInputStream.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsChunkedEncodingInputStream.java index f382bd5ced40..ec4870f5e686 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsChunkedEncodingInputStream.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsChunkedEncodingInputStream.java @@ -22,8 +22,6 @@ import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.core.checksums.SdkChecksum; import software.amazon.awssdk.core.internal.chunked.AwsChunkedEncodingConfig; -import software.amazon.awssdk.core.io.SdkInputStream; -import software.amazon.awssdk.utils.Logger; import software.amazon.awssdk.utils.Validate; /** @@ -37,37 +35,18 @@ * the wrapped stream. */ @SdkInternalApi -public abstract class AwsChunkedEncodingInputStream extends SdkInputStream { +public abstract class AwsChunkedEncodingInputStream extends AwsChunkedInputStream { - public static final int DEFAULT_CHUNK_SIZE = 128 * 1024; - protected static final int SKIP_BUFFER_SIZE = 256 * 1024; protected static final String CRLF = "\r\n"; protected static final byte[] FINAL_CHUNK = new byte[0]; protected static final String HEADER_COLON_SEPARATOR = ":"; - private static final Logger log = Logger.loggerFor(AwsChunkedEncodingInputStream.class); protected byte[] calculatedChecksum = null; protected final String checksumHeaderForTrailer; protected boolean isTrailingTerminated = true; - private InputStream is = null; private final int chunkSize; private final int maxBufferSize; private final SdkChecksum sdkChecksum; private boolean isLastTrailingCrlf; - /** - * Iterator on the current chunk. - */ - private ChunkContentIterator currentChunkIterator; - - /** - * Iterator on the buffer of the decoded stream, - * Null if the wrapped stream is marksupported, - * otherwise it will be initialized when this wrapper is marked. - */ - private DecodedStreamBuffer decodedStreamBuffer; - - private boolean isAtStart = true; - private boolean isTerminating = false; - /** * Creates a chunked encoding input stream initialized with the originating stream. The configuration allows @@ -89,10 +68,10 @@ protected AwsChunkedEncodingInputStream(InputStream in, AwsChunkedEncodingInputStream originalChunkedStream = (AwsChunkedEncodingInputStream) in; providedMaxBufferSize = Math.max(originalChunkedStream.maxBufferSize, providedMaxBufferSize); is = originalChunkedStream.is; - decodedStreamBuffer = originalChunkedStream.decodedStreamBuffer; + underlyingStreamBuffer = originalChunkedStream.underlyingStreamBuffer; } else { is = in; - decodedStreamBuffer = null; + underlyingStreamBuffer = null; } this.chunkSize = awsChunkedEncodingConfig.chunkSize(); this.maxBufferSize = providedMaxBufferSize; @@ -153,19 +132,6 @@ public T checksumHeaderForTrailer(String checksumHeaderForTrailer) { } - @Override - public int read() throws IOException { - byte[] tmp = new byte[1]; - int count = read(tmp, 0, 1); - if (count > 0) { - log.debug(() -> "One byte read from the stream."); - int unsignedByte = (int) tmp[0] & 0xFF; - return unsignedByte; - } else { - return count; - } - } - @Override public int read(byte[] b, int off, int len) throws IOException { abortIfNeeded(); @@ -211,32 +177,6 @@ private boolean setUpTrailingChunks() { return true; } - @Override - public long skip(long n) throws IOException { - if (n <= 0) { - return 0; - } - long remaining = n; - int toskip = (int) Math.min(SKIP_BUFFER_SIZE, n); - byte[] temp = new byte[toskip]; - while (remaining > 0) { - int count = read(temp, 0, toskip); - if (count < 0) { - break; - } - remaining -= count; - } - return n - remaining; - } - - /** - * @see java.io.InputStream#markSupported() - */ - @Override - public boolean markSupported() { - return true; - } - /** * The readlimit parameter is ignored. */ @@ -256,7 +196,7 @@ public void mark(int readlimit) { } else { log.debug(() -> "AwsChunkedEncodingInputStream marked at the start of the stream " + "(initializing the buffer since the wrapped stream is not mark-supported)."); - decodedStreamBuffer = new DecodedStreamBuffer(maxBufferSize); + underlyingStreamBuffer = new UnderlyingStreamBuffer(maxBufferSize); } } @@ -280,8 +220,8 @@ public void reset() throws IOException { is.reset(); } else { log.debug(() -> "AwsChunkedEncodingInputStream reset (will use the buffer of the decoded stream)."); - Validate.notNull(decodedStreamBuffer, "Cannot reset the stream because the mark is not set."); - decodedStreamBuffer.startReadBuffer(); + Validate.notNull(underlyingStreamBuffer, "Cannot reset the stream because the mark is not set."); + underlyingStreamBuffer.startReadBuffer(); } isAtStart = true; isTerminating = false; @@ -298,14 +238,14 @@ private boolean setUpNextChunk() throws IOException { int chunkSizeInBytes = 0; while (chunkSizeInBytes < chunkSize) { /** Read from the buffer of the decoded stream */ - if (null != decodedStreamBuffer && decodedStreamBuffer.hasNext()) { - chunkData[chunkSizeInBytes++] = decodedStreamBuffer.next(); + if (null != underlyingStreamBuffer && underlyingStreamBuffer.hasNext()) { + chunkData[chunkSizeInBytes++] = underlyingStreamBuffer.next(); } else { /** Read from the wrapped stream */ int bytesToRead = chunkSize - chunkSizeInBytes; int count = is.read(chunkData, chunkSizeInBytes, bytesToRead); if (count != -1) { - if (null != decodedStreamBuffer) { - decodedStreamBuffer.buffer(chunkData, chunkSizeInBytes, count); + if (null != underlyingStreamBuffer) { + underlyingStreamBuffer.buffer(chunkData, chunkSizeInBytes, count); } chunkSizeInBytes += count; } else { @@ -333,13 +273,6 @@ private boolean setUpNextChunk() throws IOException { } } - - @Override - protected InputStream getWrappedInputStream() { - return is; - } - - /** * The final chunk. * @@ -361,5 +294,4 @@ protected InputStream getWrappedInputStream() { * @return ChecksumChunkHeader in bytes based on the Header name field. */ protected abstract byte[] createChecksumChunkHeader(); - -} \ No newline at end of file +} diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsChunkedInputStream.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsChunkedInputStream.java new file mode 100644 index 000000000000..11beb216f16f --- /dev/null +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsChunkedInputStream.java @@ -0,0 +1,90 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.core.internal.io; + +import java.io.IOException; +import java.io.InputStream; +import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.core.io.SdkInputStream; +import software.amazon.awssdk.utils.Logger; + +/** + * A wrapper of InputStream that implements streaming in chunks. + */ +@SdkInternalApi +public abstract class AwsChunkedInputStream extends SdkInputStream { + public static final int DEFAULT_CHUNK_SIZE = 128 * 1024; + protected static final int SKIP_BUFFER_SIZE = 256 * 1024; + protected static final Logger log = Logger.loggerFor(AwsChunkedInputStream.class); + protected InputStream is; + /** + * Iterator on the current chunk. + */ + protected ChunkContentIterator currentChunkIterator; + + /** + * Iterator on the buffer of the underlying stream, + * Null if the wrapped stream is marksupported, + * otherwise it will be initialized when this wrapper is marked. + */ + protected UnderlyingStreamBuffer underlyingStreamBuffer; + protected boolean isAtStart = true; + protected boolean isTerminating = false; + + @Override + public int read() throws IOException { + byte[] tmp = new byte[1]; + int count = read(tmp, 0, 1); + if (count > 0) { + log.debug(() -> "One byte read from the stream."); + int unsignedByte = (int) tmp[0] & 0xFF; + return unsignedByte; + } else { + return count; + } + } + + @Override + public long skip(long n) throws IOException { + if (n <= 0) { + return 0; + } + long remaining = n; + int toskip = (int) Math.min(SKIP_BUFFER_SIZE, n); + byte[] temp = new byte[toskip]; + while (remaining > 0) { + int count = read(temp, 0, toskip); + if (count < 0) { + break; + } + remaining -= count; + } + return n - remaining; + } + + /** + * @see InputStream#markSupported() + */ + @Override + public boolean markSupported() { + return true; + } + + @Override + protected InputStream getWrappedInputStream() { + return is; + } +} diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/DecodedStreamBuffer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/UnderlyingStreamBuffer.java similarity index 93% rename from core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/DecodedStreamBuffer.java rename to core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/UnderlyingStreamBuffer.java index f6d3c47c0c1e..6fc086983fda 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/DecodedStreamBuffer.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/UnderlyingStreamBuffer.java @@ -20,8 +20,8 @@ import software.amazon.awssdk.utils.Logger; @SdkInternalApi -class DecodedStreamBuffer { - private static final Logger log = Logger.loggerFor(DecodedStreamBuffer.class); +class UnderlyingStreamBuffer { + private static final Logger log = Logger.loggerFor(UnderlyingStreamBuffer.class); private byte[] bufferArray; private int maxBufferSize; @@ -29,7 +29,7 @@ class DecodedStreamBuffer { private int pos = -1; private boolean bufferSizeOverflow; - DecodedStreamBuffer(int maxBufferSize) { + UnderlyingStreamBuffer(int maxBufferSize) { bufferArray = new byte[maxBufferSize]; this.maxBufferSize = maxBufferSize; } From 4cb38f71d0c071c7486137efc3844e2978299f32 Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 21 Jul 2023 09:24:14 -0700 Subject: [PATCH 02/14] Sync streaming compression --- .../traits/RequestCompressionTrait.java | 10 +- .../pipeline/stages/CompressRequestStage.java | 20 ++- .../io/AwsCompressionInputStream.java | 170 ++++++++++++++++++ .../CompressionContentStreamProvider.java | 55 ++++++ 4 files changed, 247 insertions(+), 8 deletions(-) create mode 100644 core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsCompressionInputStream.java create mode 100644 core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/sync/CompressionContentStreamProvider.java diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/traits/RequestCompressionTrait.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/traits/RequestCompressionTrait.java index 122e38d730e8..5168b2e36ed9 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/traits/RequestCompressionTrait.java +++ b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/traits/RequestCompressionTrait.java @@ -42,13 +42,9 @@ public static CodeBlock create(OperationModel operationModel, IntermediateModel return CodeBlock.of(""); } - // TODO : remove once request compression for streaming operations is supported - if (operationModel.isStreaming()) { - throw new IllegalStateException("Request compression for streaming operations is not yet supported in the AWS SDK " - + "for Java."); - } - - // TODO : remove once S3 checksum interceptors are moved to occur after CompressRequestStage + // TODO : remove once: + // 1) S3 checksum interceptors are moved to occur after CompressRequestStage + // 2) Transfer-Encoding:chunked is supported in S3 if (model.getMetadata().getServiceName().equals("S3")) { throw new IllegalStateException("Request compression for S3 is not yet supported in the AWS SDK for Java."); } diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java index 859f8394ac91..1eadb88d32db 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java @@ -35,6 +35,7 @@ import software.amazon.awssdk.core.internal.http.HttpClientDependencies; import software.amazon.awssdk.core.internal.http.RequestExecutionContext; import software.amazon.awssdk.core.internal.http.pipeline.MutableRequestToRequestPipeline; +import software.amazon.awssdk.core.internal.sync.CompressionContentStreamProvider; import software.amazon.awssdk.http.ContentStreamProvider; import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.utils.IoUtils; @@ -67,10 +68,21 @@ public SdkHttpFullRequest.Builder execute(SdkHttpFullRequest.Builder input, Requ compressEntirePayload(input, compressor); updateContentEncodingHeader(input, compressor); updateContentLengthHeader(input); + return input; + } + + if (!isTransferEncodingChunked(input)) { + return input; } - // TODO : streaming - sync & async + if (context.requestProvider() == null) { + // sync streaming + input.contentStreamProvider(new CompressionContentStreamProvider(input.contentStreamProvider(), compressor)); + } + + // TODO : streaming - async + updateContentEncodingHeader(input, compressor); return input; } @@ -123,6 +135,12 @@ private void updateContentLengthHeader(SdkHttpFullRequest.Builder input) { } } + private boolean isTransferEncodingChunked(SdkHttpFullRequest.Builder input) { + return input.firstMatchingHeader("Transfer-Encoding") + .map(headerValue -> headerValue.equals("chunked")) + .orElse(false); + } + private Compressor resolveCompressorType(ExecutionAttributes executionAttributes) { List encodings = executionAttributes.getAttribute(SdkInternalExecutionAttribute.REQUEST_COMPRESSION).getEncodings(); diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsCompressionInputStream.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsCompressionInputStream.java new file mode 100644 index 000000000000..93642bad8c47 --- /dev/null +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/AwsCompressionInputStream.java @@ -0,0 +1,170 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.core.internal.io; + +import java.io.IOException; +import java.io.InputStream; +import java.util.Arrays; +import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.core.internal.compression.Compressor; +import software.amazon.awssdk.utils.Validate; + +/** + * A wrapper class of InputStream that implements compression in chunks. + */ +@SdkInternalApi +public final class AwsCompressionInputStream extends AwsChunkedInputStream { + private final Compressor compressor; + + private AwsCompressionInputStream(InputStream in, Compressor compressor) { + this.compressor = compressor; + if (in instanceof AwsCompressionInputStream) { + // This could happen when the request is retried. + AwsCompressionInputStream originalCompressionStream = (AwsCompressionInputStream) in; + this.is = originalCompressionStream.is; + this.underlyingStreamBuffer = originalCompressionStream.underlyingStreamBuffer; + } else { + this.is = in; + this.underlyingStreamBuffer = null; + } + } + + public static Builder builder() { + return new Builder(); + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + abortIfNeeded(); + Validate.notNull(b, "buff"); + if (off < 0 || len < 0 || len > b.length - off) { + throw new IndexOutOfBoundsException(); + } else if (len == 0) { + return 0; + } + + if (currentChunkIterator == null || !currentChunkIterator.hasNext()) { + if (isTerminating) { + return -1; + } + isTerminating = setUpNextChunk(); + } + + int count = currentChunkIterator.read(b, off, len); + if (count > 0) { + isAtStart = false; + log.trace(() -> count + " byte read from the stream."); + } + return count; + } + + private boolean setUpNextChunk() throws IOException { + byte[] chunkData = new byte[DEFAULT_CHUNK_SIZE]; + int chunkSizeInBytes = 0; + while (chunkSizeInBytes < DEFAULT_CHUNK_SIZE) { + /** Read from the buffer of the uncompressed stream */ + if (underlyingStreamBuffer != null && underlyingStreamBuffer.hasNext()) { + chunkData[chunkSizeInBytes++] = underlyingStreamBuffer.next(); + } else { /** Read from the wrapped stream */ + int bytesToRead = DEFAULT_CHUNK_SIZE - chunkSizeInBytes; + int count = is.read(chunkData, chunkSizeInBytes, bytesToRead); + if (count != -1) { + if (underlyingStreamBuffer != null) { + underlyingStreamBuffer.buffer(chunkData, chunkSizeInBytes, count); + } + chunkSizeInBytes += count; + } else { + break; + } + } + } + if (chunkSizeInBytes == 0) { + return true; + } + + if (chunkSizeInBytes < chunkData.length) { + chunkData = Arrays.copyOf(chunkData, chunkSizeInBytes); + } + // Compress the chunk + byte[] compressedChunkData = compressor.compress(chunkData); + currentChunkIterator = new ChunkContentIterator(compressedChunkData); + return false; + } + + /** + * The readlimit parameter is ignored. + */ + @Override + public void mark(int readlimit) { + abortIfNeeded(); + if (!isAtStart) { + throw new UnsupportedOperationException("Compression stream only supports mark() at the start of the stream."); + } + if (is.markSupported()) { + log.debug(() -> "AwsCompressionInputStream marked at the start of the stream " + + "(will directly mark the wrapped stream since it's mark-supported)."); + is.mark(readlimit); + } else { + log.debug(() -> "AwsCompressionInputStream marked at the start of the stream " + + "(initializing the buffer since the wrapped stream is not mark-supported)."); + underlyingStreamBuffer = new UnderlyingStreamBuffer(SKIP_BUFFER_SIZE); + } + } + + /** + * Reset the stream, either by resetting the wrapped stream or using the + * buffer created by this class. + */ + @Override + public void reset() throws IOException { + abortIfNeeded(); + // Clear up any encoded data + currentChunkIterator = null; + // Reset the wrapped stream if it is mark-supported, + // otherwise use our buffered data. + if (is.markSupported()) { + log.debug(() -> "AwsCompressionInputStream reset " + + "(will reset the wrapped stream because it is mark-supported)."); + is.reset(); + } else { + log.debug(() -> "AwsCompressionInputStream reset (will use the buffer of the decoded stream)."); + Validate.notNull(underlyingStreamBuffer, "Cannot reset the stream because the mark is not set."); + underlyingStreamBuffer.startReadBuffer(); + } + isAtStart = true; + isTerminating = false; + } + + public static final class Builder { + InputStream inputStream; + Compressor compressor; + + public AwsCompressionInputStream build() { + return new AwsCompressionInputStream( + this.inputStream, this.compressor); + } + + public Builder inputStream(InputStream inputStream) { + this.inputStream = inputStream; + return this; + } + + public Builder compressor(Compressor compressor) { + this.compressor = compressor; + return this; + } + } +} diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/sync/CompressionContentStreamProvider.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/sync/CompressionContentStreamProvider.java new file mode 100644 index 000000000000..52a222bc372c --- /dev/null +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/sync/CompressionContentStreamProvider.java @@ -0,0 +1,55 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.core.internal.sync; + +import java.io.InputStream; +import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.core.internal.compression.Compressor; +import software.amazon.awssdk.core.internal.io.AwsCompressionInputStream; +import software.amazon.awssdk.http.ContentStreamProvider; +import software.amazon.awssdk.utils.IoUtils; + +/** + * {@link ContentStreamProvider} implementation for compression. + */ +@SdkInternalApi +public class CompressionContentStreamProvider implements ContentStreamProvider { + private final ContentStreamProvider underlyingInputStreamProvider; + private InputStream currentStream; + private final Compressor compressor; + + public CompressionContentStreamProvider(ContentStreamProvider underlyingInputStreamProvider, Compressor compressor) { + this.underlyingInputStreamProvider = underlyingInputStreamProvider; + this.compressor = compressor; + } + + @Override + public InputStream newStream() { + closeCurrentStream(); + currentStream = AwsCompressionInputStream.builder() + .inputStream(underlyingInputStreamProvider.newStream()) + .compressor(compressor) + .build(); + return currentStream; + } + + private void closeCurrentStream() { + if (currentStream != null) { + IoUtils.closeQuietly(currentStream, null); + currentStream = null; + } + } +} From da42b41a4eac374ba8b571b6f1a5eab0e1221773 Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 21 Jul 2023 09:33:33 -0700 Subject: [PATCH 03/14] Sync streaming compression functional tests --- .../customresponsemetadata/service-2.json | 24 +++++ .../services/RequestCompressionTest.java | 93 +++++++++++++++++++ 2 files changed, 117 insertions(+) diff --git a/test/codegen-generated-classes-test/src/main/resources/codegen-resources/customresponsemetadata/service-2.json b/test/codegen-generated-classes-test/src/main/resources/codegen-resources/customresponsemetadata/service-2.json index b1f994fd1d44..8cdb71614e38 100644 --- a/test/codegen-generated-classes-test/src/main/resources/codegen-resources/customresponsemetadata/service-2.json +++ b/test/codegen-generated-classes-test/src/main/resources/codegen-resources/customresponsemetadata/service-2.json @@ -289,6 +289,19 @@ "encodings": ["gzip"] } }, + "PutOperationWithStreamingRequestCompression":{ + "name":"PutOperationWithStreamingRequestCompression", + "http":{ + "method":"PUT", + "requestUri":"/" + }, + "input":{"shape":"RequestCompressionStructureWithStreaming"}, + "output":{"shape":"RequestCompressionStructureWithStreaming"}, + "requestCompression": { + "encodings": ["gzip"] + }, + "authtype":"v4-unsigned-body" + }, "GetOperationWithChecksum":{ "name":"GetOperationWithChecksum", "http":{ @@ -1030,6 +1043,17 @@ } }, "payload":"Body" + }, + "RequestCompressionStructureWithStreaming":{ + "type":"structure", + "members":{ + "Body":{ + "shape":"Body", + "documentation":"

Object data.

", + "streaming":true + } + }, + "payload":"Body" } } } diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java index e66f5f47bd1e..29664c5f53f3 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java @@ -16,9 +16,16 @@ package software.amazon.awssdk.services; import static org.assertj.core.api.Assertions.assertThat; +import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely; +import java.io.ByteArrayInputStream; +import java.io.FilterInputStream; +import java.io.IOException; import java.io.InputStream; +import java.nio.charset.StandardCharsets; import java.time.Duration; +import java.util.ArrayList; +import java.util.List; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.AfterEach; @@ -26,6 +33,9 @@ import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.core.internal.compression.Compressor; import software.amazon.awssdk.core.internal.compression.GzipCompressor; +import software.amazon.awssdk.core.sync.RequestBody; +import software.amazon.awssdk.core.sync.ResponseTransformer; +import software.amazon.awssdk.http.ContentStreamProvider; import software.amazon.awssdk.http.HttpExecuteResponse; import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.http.SdkHttpResponse; @@ -33,6 +43,7 @@ import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient; import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient; import software.amazon.awssdk.services.protocolrestjson.model.PutOperationWithRequestCompressionRequest; +import software.amazon.awssdk.services.protocolrestjson.model.PutOperationWithStreamingRequestCompressionRequest; import software.amazon.awssdk.testutils.service.http.MockAsyncHttpClient; import software.amazon.awssdk.testutils.service.http.MockSyncHttpClient; @@ -46,6 +57,7 @@ public class RequestCompressionTest { private ProtocolRestJsonClient syncClient; private ProtocolRestJsonAsyncClient asyncClient; private Compressor compressor; + private RequestBody requestBody; @BeforeEach public void setUp() { @@ -65,6 +77,8 @@ public void setUp() { byte[] compressedBodyBytes = compressor.compress(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)).asByteArray(); compressedLen = compressedBodyBytes.length; compressedBody = new String(compressedBodyBytes); + TestContentProvider provider = new TestContentProvider(UNCOMPRESSED_BODY.getBytes(StandardCharsets.UTF_8)); + requestBody = RequestBody.fromContentProvider(provider, "binary/octet-stream"); } @AfterEach @@ -118,6 +132,24 @@ public void async_nonStreaming_compression_compressesCorrectly() { assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); } + @Test + public void sync_streaming_compression_compressesCorrectly() { + mockHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + PutOperationWithStreamingRequestCompressionRequest request = + PutOperationWithStreamingRequestCompressionRequest.builder().build(); + syncClient.putOperationWithStreamingRequestCompression(request, requestBody, ResponseTransformer.toBytes()); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockHttpClient.getLastRequest(); + InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); + String loggedBody = new String(SdkBytes.fromInputStream(loggedStream).asByteArray()); + + assertThat(loggedBody).isEqualTo(compressedBody); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + assertThat(loggedRequest.matchingHeaders("Content-Length")).isEmpty(); + assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); + } + @Test public void sync_nonStreaming_compression_withRetry_compressesCorrectly() { mockHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); @@ -165,6 +197,25 @@ public void async_nonStreaming_compression_withRetry_compressesCorrectly() { assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); } + @Test + public void sync_streaming_compression_withRetry_compressesCorrectly() { + mockHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); + mockHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + PutOperationWithStreamingRequestCompressionRequest request = + PutOperationWithStreamingRequestCompressionRequest.builder().build(); + syncClient.putOperationWithStreamingRequestCompression(request, requestBody, ResponseTransformer.toBytes()); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockHttpClient.getLastRequest(); + InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); + String loggedBody = new String(SdkBytes.fromInputStream(loggedStream).asByteArray()); + + assertThat(loggedBody).isEqualTo(compressedBody); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + assertThat(loggedRequest.matchingHeaders("Content-Length")).isEmpty(); + assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); + } + private HttpExecuteResponse mockResponse() { return HttpExecuteResponse.builder() .response(SdkHttpResponse.builder().statusCode(200).build()) @@ -176,4 +227,46 @@ private HttpExecuteResponse mockErrorResponse() { .response(SdkHttpResponse.builder().statusCode(500).build()) .build(); } + + private static final class TestContentProvider implements ContentStreamProvider { + private final byte[] content; + private final List createdStreams = new ArrayList<>(); + private CloseTrackingInputStream currentStream; + + private TestContentProvider(byte[] content) { + this.content = content; + } + + @Override + public InputStream newStream() { + if (currentStream != null) { + invokeSafely(currentStream::close); + } + currentStream = new CloseTrackingInputStream(new ByteArrayInputStream(content)); + createdStreams.add(currentStream); + return currentStream; + } + + List getCreatedStreams() { + return createdStreams; + } + } + + private static class CloseTrackingInputStream extends FilterInputStream { + private boolean isClosed = false; + + CloseTrackingInputStream(InputStream in) { + super(in); + } + + @Override + public void close() throws IOException { + super.close(); + isClosed = true; + } + + boolean isClosed() { + return isClosed; + } + } } From 0e561b4de31744de55f9b3ddb8ff7229fdd15231 Mon Sep 17 00:00:00 2001 From: hdavidh Date: Thu, 20 Jul 2023 16:28:28 -0700 Subject: [PATCH 04/14] Sync streaming compression integ tests --- services/mediastoredata/pom.xml | 6 + .../MediaStoreDataIntegrationTestBase.java | 155 ++++++++++++++++ ...stCompressionStreamingIntegrationTest.java | 168 ++++++++++++++++++ ...ransferEncodingChunkedIntegrationTest.java | 122 +------------ .../src/it/resources/log4j2.properties | 38 ++++ 5 files changed, 369 insertions(+), 120 deletions(-) create mode 100644 services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java create mode 100644 services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java create mode 100644 services/mediastoredata/src/it/resources/log4j2.properties diff --git a/services/mediastoredata/pom.xml b/services/mediastoredata/pom.xml index 32d52843ee66..3e5e81a0a216 100644 --- a/services/mediastoredata/pom.xml +++ b/services/mediastoredata/pom.xml @@ -74,5 +74,11 @@ commons-lang3 test + + software.amazon.awssdk + mediastore + ${awsjavasdk.version} + test + diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java new file mode 100644 index 000000000000..bb65e4ff7579 --- /dev/null +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java @@ -0,0 +1,155 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.services.mediastoredata; + +import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely; + +import io.reactivex.Flowable; +import java.io.ByteArrayInputStream; +import java.io.FilterInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.net.URI; +import java.nio.ByteBuffer; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.reactivestreams.Subscriber; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.http.ContentStreamProvider; +import software.amazon.awssdk.http.apache.ApacheHttpClient; +import software.amazon.awssdk.services.mediastore.MediaStoreClient; +import software.amazon.awssdk.services.mediastore.model.Container; +import software.amazon.awssdk.services.mediastore.model.ContainerStatus; +import software.amazon.awssdk.services.mediastore.model.DescribeContainerResponse; +import software.amazon.awssdk.testutils.Waiter; +import software.amazon.awssdk.testutils.service.AwsIntegrationTestBase; + +/** + * Base class for MediaStoreData integration tests. Used for Transfer-Encoding and Request Compression testing. + */ +public class MediaStoreDataIntegrationTestBase extends AwsIntegrationTestBase { + protected static final String CONTAINER_NAME = "java-sdk-test-mediastoredata-" + Instant.now().toEpochMilli(); + protected static AwsCredentialsProvider credentialsProvider; + protected static MediaStoreClient mediaStoreClient; + protected static URI uri; + + @BeforeAll + public static void init() { + credentialsProvider = getCredentialsProvider(); + mediaStoreClient = MediaStoreClient.builder() + .credentialsProvider(credentialsProvider) + .httpClient(ApacheHttpClient.builder().build()) + .build(); + uri = URI.create(createContainer().endpoint()); + } + + @AfterEach + public static void reset() { + CaptureTransferEncodingHeaderInterceptor.reset(); + } + + private static Container createContainer() { + mediaStoreClient.createContainer(r -> r.containerName(CONTAINER_NAME)); + DescribeContainerResponse response = waitContainerToBeActive(); + return response.container(); + } + + private static DescribeContainerResponse waitContainerToBeActive() { + return Waiter.run(() -> mediaStoreClient.describeContainer(r -> r.containerName(CONTAINER_NAME))) + .until(r -> r.container().status() == ContainerStatus.ACTIVE) + .orFailAfter(Duration.ofMinutes(3)); + } + + protected AsyncRequestBody customAsyncRequestBodyWithoutContentLength() { + return new AsyncRequestBody() { + @Override + public Optional contentLength() { + return Optional.empty(); + } + + @Override + public void subscribe(Subscriber s) { + Flowable.fromPublisher(AsyncRequestBody.fromBytes("Random text".getBytes())) + .subscribe(s); + } + }; + } + + protected static class CaptureTransferEncodingHeaderInterceptor implements ExecutionInterceptor { + public static boolean isChunked; + + public static void reset() { + isChunked = false; + } + + @Override + public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) { + isChunked = context.httpRequest().matchingHeaders("Transfer-Encoding").contains("chunked"); + } + } + + protected static class TestContentProvider implements ContentStreamProvider { + private final byte[] content; + private final List createdStreams = new ArrayList<>(); + private CloseTrackingInputStream currentStream; + + protected TestContentProvider(byte[] content) { + this.content = content.clone(); + } + + @Override + public InputStream newStream() { + if (currentStream != null) { + invokeSafely(currentStream::close); + } + currentStream = new CloseTrackingInputStream(new ByteArrayInputStream(content)); + createdStreams.add(currentStream); + return currentStream; + } + + List getCreatedStreams() { + return Collections.unmodifiableList(createdStreams); + } + } + + protected static class CloseTrackingInputStream extends FilterInputStream { + private boolean isClosed = false; + + CloseTrackingInputStream(InputStream in) { + super(in); + } + + @Override + public void close() throws IOException { + super.close(); + isClosed = true; + } + + boolean isClosed() { + return isClosed; + } + } +} diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java new file mode 100644 index 000000000000..5e9b29038249 --- /dev/null +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java @@ -0,0 +1,168 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.services.mediastoredata; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.core.RequestCompressionConfiguration; +import software.amazon.awssdk.core.ResponseInputStream; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute; +import software.amazon.awssdk.core.internal.compression.Compressor; +import software.amazon.awssdk.core.internal.compression.GzipCompressor; +import software.amazon.awssdk.core.internal.interceptor.trait.RequestCompression; +import software.amazon.awssdk.core.sync.RequestBody; +import software.amazon.awssdk.http.apache.ApacheHttpClient; +import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient; +import software.amazon.awssdk.services.mediastoredata.model.DeleteObjectRequest; +import software.amazon.awssdk.services.mediastoredata.model.GetObjectRequest; +import software.amazon.awssdk.services.mediastoredata.model.GetObjectResponse; +import software.amazon.awssdk.services.mediastoredata.model.ObjectNotFoundException; +import software.amazon.awssdk.services.mediastoredata.model.PutObjectRequest; +import software.amazon.awssdk.testutils.Waiter; + +/** + * Integration test to verify Request Compression functionalities for streaming operations. Do not delete. + */ +public class RequestCompressionStreamingIntegrationTest extends MediaStoreDataIntegrationTestBase { + private static final String UNCOMPRESSED_BODY = + "RequestCompressionTest-RequestCompressionTest-RequestCompressionTest-RequestCompressionTest-RequestCompressionTest"; + private static String compressedBody; + private static MediaStoreDataClient syncClient; + private static MediaStoreDataAsyncClient asyncClient; + private static PutObjectRequest putObjectRequest; + private static DeleteObjectRequest deleteObjectRequest; + private static GetObjectRequest getObjectRequest; + + @BeforeAll + public static void setup() { + RequestCompressionConfiguration compressionConfiguration = + RequestCompressionConfiguration.builder() + .minimumCompressionThresholdInBytes(1) + .requestCompressionEnabled(true) + .build(); + + RequestCompression requestCompressionTrait = RequestCompression.builder() + .encodings("gzip") + .isStreaming(true) + .build(); + + syncClient = MediaStoreDataClient.builder() + .endpointOverride(uri) + .credentialsProvider(credentialsProvider) + .httpClient(ApacheHttpClient.builder().build()) + .overrideConfiguration(o -> o.addExecutionInterceptor(new CaptureTransferEncodingHeaderInterceptor()) + .addExecutionInterceptor(new CaptureContentEncodingHeaderInterceptor()) + .putExecutionAttribute(SdkInternalExecutionAttribute.REQUEST_COMPRESSION, + requestCompressionTrait) + .requestCompressionConfiguration(compressionConfiguration)) + .build(); + + asyncClient = MediaStoreDataAsyncClient.builder() + .endpointOverride(uri) + .credentialsProvider(getCredentialsProvider()) + .httpClient(NettyNioAsyncHttpClient.create()) + .overrideConfiguration(o -> o.addExecutionInterceptor(new CaptureTransferEncodingHeaderInterceptor()) + .addExecutionInterceptor(new CaptureContentEncodingHeaderInterceptor()) + .putExecutionAttribute(SdkInternalExecutionAttribute.REQUEST_COMPRESSION, + requestCompressionTrait) + .requestCompressionConfiguration(compressionConfiguration)) + .build(); + + putObjectRequest = PutObjectRequest.builder() + .contentType("application/octet-stream") + .path("/foo") + .overrideConfiguration( + o -> o.requestCompressionConfiguration( + c -> c.requestCompressionEnabled(true))) + .build(); + deleteObjectRequest = DeleteObjectRequest.builder().path("/foo").build(); + getObjectRequest = GetObjectRequest.builder().path("/foo").build(); + + Compressor compressor = new GzipCompressor(); + byte[] compressedBodyBytes = compressor.compress(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)).asByteArray(); + compressedBody = new String(compressedBodyBytes); + } + + @AfterAll + public static void tearDown() { + syncClient.deleteObject(deleteObjectRequest); + Waiter.run(() -> syncClient.describeObject(r -> r.path("/foo"))) + .untilException(ObjectNotFoundException.class) + .orFailAfter(Duration.ofMinutes(1)); + } + + @AfterEach + public static void reset() { + CaptureContentEncodingHeaderInterceptor.reset(); + } + + @Test + public void putObject_withRequestCompressionSyncStreaming_compressesPayloadAndSendsCorrectly() throws IOException { + TestContentProvider provider = new TestContentProvider(UNCOMPRESSED_BODY.getBytes(StandardCharsets.UTF_8)); + syncClient.putObject(putObjectRequest, RequestBody.fromContentProvider(provider, "binary/octet-stream")); + + assertThat(CaptureTransferEncodingHeaderInterceptor.isChunked).isTrue(); + assertThat(CaptureContentEncodingHeaderInterceptor.isGzip).isTrue(); + + ResponseInputStream response = syncClient.getObject(getObjectRequest); + byte[] buffer = new byte[UNCOMPRESSED_BODY.getBytes(StandardCharsets.UTF_8).length]; + response.read(buffer); + String retrievedContent = new String(buffer); + assertThat(UNCOMPRESSED_BODY).isEqualTo(retrievedContent); + } + + // TODO : uncomment once async streaming compression is implemented + /*@Test + public void nettyClientPutObject_withoutContentLength_sendsSuccessfully() throws IOException { + AsyncRequestBody asyncRequestBody = customAsyncRequestBodyWithoutContentLength(); + asyncClient.putObject(putObjectRequest, asyncRequestBody).join(); + + assertThat(CaptureTransferEncodingHeaderInterceptor.isChunked).isTrue(); + assertThat(CaptureContentEncodingHeaderInterceptor.isGzip).isTrue(); + + // verify stored content is correct + ResponseInputStream response = syncClient.getObject(getObjectRequest); + byte[] buffer = new byte[UNCOMPRESSED_BODY.getBytes(StandardCharsets.UTF_8).length]; + response.read(buffer); + String retrievedContent = new String(buffer); + assertThat(UNCOMPRESSED_BODY).isEqualTo(retrievedContent); + assertThat(CaptureTransferEncodingHeaderInterceptor.isChunked).isTrue(); + }*/ + + private static class CaptureContentEncodingHeaderInterceptor implements ExecutionInterceptor { + public static boolean isGzip; + + public static void reset() { + isGzip = false; + } + + @Override + public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) { + isGzip = context.httpRequest().matchingHeaders("Content-Encoding").contains("gzip"); + } + } +} diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/TransferEncodingChunkedIntegrationTest.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/TransferEncodingChunkedIntegrationTest.java index acab0a8d6723..80fb67dc6fab 100644 --- a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/TransferEncodingChunkedIntegrationTest.java +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/TransferEncodingChunkedIntegrationTest.java @@ -16,70 +16,34 @@ package software.amazon.awssdk.services.mediastoredata; import static org.assertj.core.api.AssertionsForClassTypes.assertThat; -import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely; -import io.reactivex.Flowable; -import java.io.ByteArrayInputStream; -import java.io.FilterInputStream; -import java.io.IOException; -import java.io.InputStream; -import java.net.URI; -import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; -import java.util.List; -import java.util.Optional; import org.apache.commons.lang3.RandomStringUtils; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; -import org.reactivestreams.Subscriber; -import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; -import software.amazon.awssdk.core.async.AsyncRequestBody; -import software.amazon.awssdk.core.interceptor.Context; -import software.amazon.awssdk.core.interceptor.ExecutionAttributes; -import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; import software.amazon.awssdk.core.sync.RequestBody; -import software.amazon.awssdk.http.ContentStreamProvider; import software.amazon.awssdk.http.apache.ApacheHttpClient; import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient; import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient; -import software.amazon.awssdk.services.mediastore.MediaStoreClient; -import software.amazon.awssdk.services.mediastore.model.Container; -import software.amazon.awssdk.services.mediastore.model.ContainerStatus; -import software.amazon.awssdk.services.mediastore.model.DescribeContainerResponse; import software.amazon.awssdk.services.mediastoredata.model.DeleteObjectRequest; import software.amazon.awssdk.services.mediastoredata.model.ObjectNotFoundException; import software.amazon.awssdk.services.mediastoredata.model.PutObjectRequest; import software.amazon.awssdk.testutils.Waiter; -import software.amazon.awssdk.testutils.service.AwsIntegrationTestBase; /** * Integration test to verify Transfer-Encoding:chunked functionalities for all supported HTTP clients. Do not delete. */ -public class TransferEncodingChunkedIntegrationTest extends AwsIntegrationTestBase { - private static final String CONTAINER_NAME = "java-sdk-test-" + Instant.now().toEpochMilli(); - private static MediaStoreClient mediaStoreClient; +public class TransferEncodingChunkedIntegrationTest extends MediaStoreDataIntegrationTestBase { private static MediaStoreDataClient syncClientWithApache; private static MediaStoreDataClient syncClientWithUrlConnection; private static MediaStoreDataAsyncClient asyncClientWithNetty; - private static AwsCredentialsProvider credentialsProvider; - private static Container container; private static PutObjectRequest putObjectRequest; private static DeleteObjectRequest deleteObjectRequest; @BeforeAll public static void setup() { - credentialsProvider = getCredentialsProvider(); - mediaStoreClient = MediaStoreClient.builder() - .credentialsProvider(credentialsProvider) - .httpClient(ApacheHttpClient.builder().build()) - .build(); - container = createContainer(); - URI uri = URI.create(container.endpoint()); - syncClientWithApache = MediaStoreDataClient.builder() .endpointOverride(uri) .credentialsProvider(credentialsProvider) @@ -117,7 +81,7 @@ public static void tearDown() { Waiter.run(() -> syncClientWithApache.describeObject(r -> r.path("/foo"))) .untilException(ObjectNotFoundException.class) .orFailAfter(Duration.ofMinutes(1)); - CaptureTransferEncodingHeaderInterceptor.reset(); + mediaStoreClient.deleteContainer(r -> r.containerName(CONTAINER_NAME)); } @Test @@ -139,86 +103,4 @@ public void nettyClientPutObject_withoutContentLength_sendsSuccessfully() { asyncClientWithNetty.putObject(putObjectRequest, customAsyncRequestBodyWithoutContentLength()).join(); assertThat(CaptureTransferEncodingHeaderInterceptor.isChunked).isTrue(); } - - private static Container createContainer() { - mediaStoreClient.createContainer(r -> r.containerName(CONTAINER_NAME)); - DescribeContainerResponse response = waitContainerToBeActive(); - return response.container(); - } - - private static DescribeContainerResponse waitContainerToBeActive() { - return Waiter.run(() -> mediaStoreClient.describeContainer(r -> r.containerName(CONTAINER_NAME))) - .until(r -> ContainerStatus.ACTIVE.equals(r.container().status())) - .orFailAfter(Duration.ofMinutes(3)); - } - - private static class CaptureTransferEncodingHeaderInterceptor implements ExecutionInterceptor { - private static boolean isChunked; - - public static void reset() { - isChunked = false; - } - - @Override - public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) { - isChunked = context.httpRequest().matchingHeaders("Transfer-Encoding").contains("chunked"); - } - } - - private AsyncRequestBody customAsyncRequestBodyWithoutContentLength() { - return new AsyncRequestBody() { - @Override - public Optional contentLength() { - return Optional.empty(); - } - - @Override - public void subscribe(Subscriber s) { - Flowable.fromPublisher(AsyncRequestBody.fromBytes("Random text".getBytes())) - .subscribe(s); - } - }; - } - - private static class TestContentProvider implements ContentStreamProvider { - private final byte[] content; - private final List createdStreams = new ArrayList<>(); - private CloseTrackingInputStream currentStream; - - private TestContentProvider(byte[] content) { - this.content = content; - } - - @Override - public InputStream newStream() { - if (currentStream != null) { - invokeSafely(currentStream::close); - } - currentStream = new CloseTrackingInputStream(new ByteArrayInputStream(content)); - createdStreams.add(currentStream); - return currentStream; - } - - List getCreatedStreams() { - return createdStreams; - } - } - - private static class CloseTrackingInputStream extends FilterInputStream { - private boolean isClosed = false; - - CloseTrackingInputStream(InputStream in) { - super(in); - } - - @Override - public void close() throws IOException { - super.close(); - isClosed = true; - } - - boolean isClosed() { - return isClosed; - } - } } diff --git a/services/mediastoredata/src/it/resources/log4j2.properties b/services/mediastoredata/src/it/resources/log4j2.properties new file mode 100644 index 000000000000..ea24f17148e6 --- /dev/null +++ b/services/mediastoredata/src/it/resources/log4j2.properties @@ -0,0 +1,38 @@ +# +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# A copy of the License is located at +# +# http://aws.amazon.com/apache2.0 +# +# or in the "license" file accompanying this file. This file is distributed +# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either +# express or implied. See the License for the specific language governing +# permissions and limitations under the License. +# + +status = warn + +appender.console.type = Console +appender.console.name = ConsoleAppender +appender.console.layout.type = PatternLayout +appender.console.layout.pattern = %d{HH:mm:ss.SSS} [%t] %-5level %logger{36} - %msg%n%throwable + +rootLogger.level = info +rootLogger.appenderRef.stdout.ref = ConsoleAppender + +# Uncomment below to enable more specific logging +# +#logger.sdk.name = software.amazon.awssdk +#logger.sdk.level = debug +# +#logger.request.name = software.amazon.awssdk.request +#logger.request.level = debug +# +#logger.apache.name = org.apache.http.wire +#logger.apache.level = debug +# +#logger.netty.name = io.netty.handler.logging +#logger.netty.level = debug \ No newline at end of file From a89b134248342fbfebeed8748311393722808120 Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 21 Jul 2023 13:58:00 -0700 Subject: [PATCH 05/14] Fix integ test --- .../mediastoredata/MediaStoreDataIntegrationTestBase.java | 2 +- .../RequestCompressionStreamingIntegrationTest.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java index bb65e4ff7579..3a0e7006ef8b 100644 --- a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java @@ -67,7 +67,7 @@ public static void init() { } @AfterEach - public static void reset() { + public void reset() { CaptureTransferEncodingHeaderInterceptor.reset(); } diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java index 5e9b29038249..9530f2319b38 100644 --- a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java @@ -116,7 +116,7 @@ public static void tearDown() { } @AfterEach - public static void reset() { + public void cleanUp() { CaptureContentEncodingHeaderInterceptor.reset(); } From 7550c21d00227a6f28bced379d03016ed89d8c4b Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 4 Aug 2023 10:27:13 -0700 Subject: [PATCH 06/14] Async streaming compression --- .../async/ChunkBufferWithUnknownLength.java | 80 ++++++++++ .../async/CompressionAsyncRequestBody.java | 146 ++++++++++++++++++ .../pipeline/stages/CompressRequestStage.java | 9 +- .../CompressionAsyncRequestBodyTest.java | 86 +++++++++++ .../MediaStoreDataIntegrationTestBase.java | 4 +- ...stCompressionStreamingIntegrationTest.java | 28 ++-- ...ransferEncodingChunkedIntegrationTest.java | 5 +- .../services/RequestCompressionTest.java | 55 ++++++- 8 files changed, 391 insertions(+), 22 deletions(-) create mode 100644 core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBufferWithUnknownLength.java create mode 100644 core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java create mode 100644 core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBufferWithUnknownLength.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBufferWithUnknownLength.java new file mode 100644 index 000000000000..5a5cfca6dc9e --- /dev/null +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBufferWithUnknownLength.java @@ -0,0 +1,80 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.core.internal.async; + +import static software.amazon.awssdk.core.HttpChecksumConstant.DEFAULT_ASYNC_CHUNK_SIZE; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.utils.builder.SdkBuilder; + +/** + * Class that will buffer incoming BufferBytes with unknown total length to chunks of bufferSize + */ +@SdkInternalApi +public final class ChunkBufferWithUnknownLength { + private ByteBuffer currentBuffer; + + private ChunkBufferWithUnknownLength(Integer bufferSize) { + int chunkSize = bufferSize != null ? bufferSize : DEFAULT_ASYNC_CHUNK_SIZE; + this.currentBuffer = ByteBuffer.allocate(chunkSize); + } + + public static Builder builder() { + return new DefaultBuilder(); + } + + public synchronized Iterable bufferAndCreateChunks(ByteBuffer buffer) { + List bufferedList = new ArrayList<>(); + while (buffer.hasRemaining()) { + int bytesToCopy = Math.min(buffer.remaining(), currentBuffer.remaining()); + byte[] bytes = new byte[bytesToCopy]; + buffer.get(bytes); + currentBuffer.put(bytes); + + if (!currentBuffer.hasRemaining() || !buffer.hasRemaining()) { + currentBuffer.flip(); + ByteBuffer bufferToSend = ByteBuffer.allocate(currentBuffer.limit()); + bufferToSend.put(currentBuffer); + bufferToSend.flip(); + bufferedList.add(bufferToSend); + currentBuffer.clear(); + } + } + return bufferedList; + } + + public interface Builder extends SdkBuilder { + Builder bufferSize(int bufferSize); + } + + private static final class DefaultBuilder implements Builder { + private Integer bufferSize; + + @Override + public ChunkBufferWithUnknownLength build() { + return new ChunkBufferWithUnknownLength(bufferSize); + } + + @Override + public Builder bufferSize(int bufferSize) { + this.bufferSize = bufferSize; + return this; + } + } +} diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java new file mode 100644 index 000000000000..49e7762c35d2 --- /dev/null +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java @@ -0,0 +1,146 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.core.internal.async; + +import java.nio.ByteBuffer; +import java.util.Optional; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; +import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.core.internal.compression.Compressor; +import software.amazon.awssdk.utils.Validate; +import software.amazon.awssdk.utils.builder.SdkBuilder; + +/** + * Wrapper class to wrap an AsyncRequestBody. + * This will chunk and compress the payload with the provided {@link Compressor}. + */ +@SdkInternalApi +public class CompressionAsyncRequestBody implements AsyncRequestBody { + + private static final int COMPRESSION_CHUNK_SIZE = 128 * 1024; + private final AsyncRequestBody wrapped; + private final Compressor compressor; + + private CompressionAsyncRequestBody(DefaultBuilder builder) { + Validate.notNull(builder.asyncRequestBody, "wrapped AsyncRequestBody cannot be null"); + Validate.notNull(builder.compressor, "compressor cannot be null"); + this.wrapped = builder.asyncRequestBody; + this.compressor = builder.compressor; + } + + /** + * @return Builder instance to construct a {@link CompressionAsyncRequestBody}. + */ + public static Builder builder() { + return new DefaultBuilder(); + } + + public interface Builder extends SdkBuilder { + + /** + * Sets the AsyncRequestBody that will be wrapped. + * @param asyncRequestBody + * @return This builder for method chaining. + */ + Builder asyncRequestBody(AsyncRequestBody asyncRequestBody); + + /** + * Sets the compressor to compress the request. + * @param compressor + * @return This builder for method chaining. + */ + Builder compressor(Compressor compressor); + } + + private static final class DefaultBuilder implements Builder { + + private AsyncRequestBody asyncRequestBody; + private Compressor compressor; + + @Override + public CompressionAsyncRequestBody build() { + return new CompressionAsyncRequestBody(this); + } + + @Override + public Builder asyncRequestBody(AsyncRequestBody asyncRequestBody) { + this.asyncRequestBody = asyncRequestBody; + return this; + } + + @Override + public Builder compressor(Compressor compressor) { + this.compressor = compressor; + return this; + } + } + + @Override + public Optional contentLength() { + return Optional.empty(); + } + + @Override + public String contentType() { + return wrapped.contentType(); + } + + @Override + public void subscribe(Subscriber s) { + Validate.notNull(s, "Subscription MUST NOT be null."); + + ChunkBufferWithUnknownLength chunkBuffer = ChunkBufferWithUnknownLength.builder() + .bufferSize(COMPRESSION_CHUNK_SIZE) + .build(); + + wrapped.flatMapIterable(chunkBuffer::bufferAndCreateChunks) + .subscribe(new CompressionSubscriber(s, compressor)); + } + + private static final class CompressionSubscriber implements Subscriber { + + private final Subscriber wrapped; + private final Compressor compressor; + + CompressionSubscriber(Subscriber wrapped, Compressor compressor) { + this.wrapped = wrapped; + this.compressor = compressor; + } + + @Override + public void onSubscribe(Subscription subscription) { + wrapped.onSubscribe(subscription); + } + + @Override + public void onNext(ByteBuffer byteBuffer) { + ByteBuffer compressedBuffer = compressor.compress(byteBuffer); + wrapped.onNext(compressedBuffer); + } + + @Override + public void onError(Throwable t) { + wrapped.onError(t); + } + + @Override + public void onComplete() { + wrapped.onComplete(); + } + } +} diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java index 1eadb88d32db..6ca358bd312c 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java @@ -30,6 +30,7 @@ import software.amazon.awssdk.core.exception.SdkClientException; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute; +import software.amazon.awssdk.core.internal.async.CompressionAsyncRequestBody; import software.amazon.awssdk.core.internal.compression.Compressor; import software.amazon.awssdk.core.internal.compression.CompressorType; import software.amazon.awssdk.core.internal.http.HttpClientDependencies; @@ -78,10 +79,14 @@ public SdkHttpFullRequest.Builder execute(SdkHttpFullRequest.Builder input, Requ if (context.requestProvider() == null) { // sync streaming input.contentStreamProvider(new CompressionContentStreamProvider(input.contentStreamProvider(), compressor)); + } else { + // async streaming + context.requestProvider(CompressionAsyncRequestBody.builder() + .asyncRequestBody(context.requestProvider()) + .compressor(compressor) + .build()); } - // TODO : streaming - async - updateContentEncodingHeader(input, compressor); return input; } diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java new file mode 100644 index 000000000000..fb1873d3c57b --- /dev/null +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java @@ -0,0 +1,86 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.core.internal.async; + +import static org.assertj.core.api.Assertions.assertThat; + +import io.reactivex.Flowable; +import java.nio.ByteBuffer; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.Test; +import org.reactivestreams.Subscriber; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.core.internal.compression.Compressor; +import software.amazon.awssdk.core.internal.compression.GzipCompressor; +import software.amazon.awssdk.http.async.SimpleSubscriber; + +public final class CompressionAsyncRequestBodyTest { + private static final Compressor compressor = new GzipCompressor();; + private final static String TEST_STRING = + "RequestCompressionTest-RequestCompressionTest-RequestCompressionTest-RequestCompressionTest-RequestCompressionTest"; + private static final String EXPECTED_TEST_STRING = + new String(compressor.compress(TEST_STRING.getBytes())); + + private final AsyncRequestBody provider = CompressionAsyncRequestBody.builder() + .compressor(compressor) + .asyncRequestBody(customAsyncRequestBodyWithoutContentLength()) + .build(); + + @Test + public void hasCorrectContent() throws InterruptedException { + StringBuilder sb = new StringBuilder(); + CountDownLatch done = new CountDownLatch(1); + + Subscriber subscriber = new SimpleSubscriber(buffer -> { + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + sb.append(new String(bytes)); + }) { + @Override + public void onError(Throwable t) { + super.onError(t); + done.countDown(); + } + + @Override + public void onComplete() { + super.onComplete(); + done.countDown(); + } + }; + + provider.subscribe(subscriber); + done.await(10, TimeUnit.SECONDS); + assertThat(sb).hasToString(EXPECTED_TEST_STRING); + } + + protected AsyncRequestBody customAsyncRequestBodyWithoutContentLength() { + return new AsyncRequestBody() { + @Override + public Optional contentLength() { + return Optional.empty(); + } + + @Override + public void subscribe(Subscriber s) { + Flowable.fromPublisher(AsyncRequestBody.fromBytes(TEST_STRING.getBytes())) + .subscribe(s); + } + }; + } +} \ No newline at end of file diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java index 3a0e7006ef8b..20688925bc86 100644 --- a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/MediaStoreDataIntegrationTestBase.java @@ -83,7 +83,7 @@ private static DescribeContainerResponse waitContainerToBeActive() { .orFailAfter(Duration.ofMinutes(3)); } - protected AsyncRequestBody customAsyncRequestBodyWithoutContentLength() { + protected AsyncRequestBody customAsyncRequestBodyWithoutContentLength(byte[] body) { return new AsyncRequestBody() { @Override public Optional contentLength() { @@ -92,7 +92,7 @@ public Optional contentLength() { @Override public void subscribe(Subscriber s) { - Flowable.fromPublisher(AsyncRequestBody.fromBytes("Random text".getBytes())) + Flowable.fromPublisher(AsyncRequestBody.fromBytes(body)) .subscribe(s); } }; diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java index 9530f2319b38..41c5b2fc849d 100644 --- a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java @@ -27,6 +27,7 @@ import software.amazon.awssdk.core.RequestCompressionConfiguration; import software.amazon.awssdk.core.ResponseInputStream; import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.interceptor.Context; import software.amazon.awssdk.core.interceptor.ExecutionAttributes; import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; @@ -83,7 +84,7 @@ public static void setup() { asyncClient = MediaStoreDataAsyncClient.builder() .endpointOverride(uri) - .credentialsProvider(getCredentialsProvider()) + .credentialsProvider(credentialsProvider) .httpClient(NettyNioAsyncHttpClient.create()) .overrideConfiguration(o -> o.addExecutionInterceptor(new CaptureTransferEncodingHeaderInterceptor()) .addExecutionInterceptor(new CaptureContentEncodingHeaderInterceptor()) @@ -108,11 +109,13 @@ public static void setup() { } @AfterAll - public static void tearDown() { + public static void tearDown() throws InterruptedException { syncClient.deleteObject(deleteObjectRequest); Waiter.run(() -> syncClient.describeObject(r -> r.path("/foo"))) .untilException(ObjectNotFoundException.class) .orFailAfter(Duration.ofMinutes(1)); + Thread.sleep(500); + mediaStoreClient.deleteContainer(r -> r.containerName(CONTAINER_NAME)); } @AfterEach @@ -121,7 +124,7 @@ public void cleanUp() { } @Test - public void putObject_withRequestCompressionSyncStreaming_compressesPayloadAndSendsCorrectly() throws IOException { + public void putObject_withSyncStreamingRequestCompression_compressesPayloadAndSendsCorrectly() throws IOException { TestContentProvider provider = new TestContentProvider(UNCOMPRESSED_BODY.getBytes(StandardCharsets.UTF_8)); syncClient.putObject(putObjectRequest, RequestBody.fromContentProvider(provider, "binary/octet-stream")); @@ -129,29 +132,26 @@ public void putObject_withRequestCompressionSyncStreaming_compressesPayloadAndSe assertThat(CaptureContentEncodingHeaderInterceptor.isGzip).isTrue(); ResponseInputStream response = syncClient.getObject(getObjectRequest); - byte[] buffer = new byte[UNCOMPRESSED_BODY.getBytes(StandardCharsets.UTF_8).length]; + byte[] buffer = new byte[UNCOMPRESSED_BODY.getBytes().length]; response.read(buffer); String retrievedContent = new String(buffer); - assertThat(UNCOMPRESSED_BODY).isEqualTo(retrievedContent); + assertThat(retrievedContent).isEqualTo(UNCOMPRESSED_BODY); } - // TODO : uncomment once async streaming compression is implemented - /*@Test - public void nettyClientPutObject_withoutContentLength_sendsSuccessfully() throws IOException { - AsyncRequestBody asyncRequestBody = customAsyncRequestBodyWithoutContentLength(); + @Test + public void putObject_withAsyncStreamingRequestCompression_compressesPayloadAndSendsCorrectly() throws IOException { + AsyncRequestBody asyncRequestBody = customAsyncRequestBodyWithoutContentLength(UNCOMPRESSED_BODY.getBytes()); asyncClient.putObject(putObjectRequest, asyncRequestBody).join(); assertThat(CaptureTransferEncodingHeaderInterceptor.isChunked).isTrue(); assertThat(CaptureContentEncodingHeaderInterceptor.isGzip).isTrue(); - // verify stored content is correct ResponseInputStream response = syncClient.getObject(getObjectRequest); - byte[] buffer = new byte[UNCOMPRESSED_BODY.getBytes(StandardCharsets.UTF_8).length]; + byte[] buffer = new byte[UNCOMPRESSED_BODY.getBytes().length]; response.read(buffer); String retrievedContent = new String(buffer); - assertThat(UNCOMPRESSED_BODY).isEqualTo(retrievedContent); - assertThat(CaptureTransferEncodingHeaderInterceptor.isChunked).isTrue(); - }*/ + assertThat(retrievedContent).isEqualTo(UNCOMPRESSED_BODY); + } private static class CaptureContentEncodingHeaderInterceptor implements ExecutionInterceptor { public static boolean isGzip; diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/TransferEncodingChunkedIntegrationTest.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/TransferEncodingChunkedIntegrationTest.java index 80fb67dc6fab..34522618f759 100644 --- a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/TransferEncodingChunkedIntegrationTest.java +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/TransferEncodingChunkedIntegrationTest.java @@ -76,11 +76,12 @@ public static void setup() { } @AfterAll - public static void tearDown() { + public static void tearDown() throws InterruptedException { syncClientWithApache.deleteObject(deleteObjectRequest); Waiter.run(() -> syncClientWithApache.describeObject(r -> r.path("/foo"))) .untilException(ObjectNotFoundException.class) .orFailAfter(Duration.ofMinutes(1)); + Thread.sleep(500); mediaStoreClient.deleteContainer(r -> r.containerName(CONTAINER_NAME)); } @@ -100,7 +101,7 @@ public void urlConnectionClientPutObject_withoutContentLength_sendsSuccessfully( @Test public void nettyClientPutObject_withoutContentLength_sendsSuccessfully() { - asyncClientWithNetty.putObject(putObjectRequest, customAsyncRequestBodyWithoutContentLength()).join(); + asyncClientWithNetty.putObject(putObjectRequest, customAsyncRequestBodyWithoutContentLength("TestBody".getBytes())).join(); assertThat(CaptureTransferEncodingHeaderInterceptor.isChunked).isTrue(); } } diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java index 29664c5f53f3..9113c131307a 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java @@ -18,19 +18,24 @@ import static org.assertj.core.api.Assertions.assertThat; import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely; +import io.reactivex.Flowable; import java.io.ByteArrayInputStream; import java.io.FilterInputStream; import java.io.IOException; import java.io.InputStream; -import java.nio.charset.StandardCharsets; +import java.nio.ByteBuffer; import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.AfterEach; +import org.reactivestreams.Subscriber; import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.core.async.AsyncResponseTransformer; import software.amazon.awssdk.core.internal.compression.Compressor; import software.amazon.awssdk.core.internal.compression.GzipCompressor; import software.amazon.awssdk.core.sync.RequestBody; @@ -77,7 +82,7 @@ public void setUp() { byte[] compressedBodyBytes = compressor.compress(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)).asByteArray(); compressedLen = compressedBodyBytes.length; compressedBody = new String(compressedBodyBytes); - TestContentProvider provider = new TestContentProvider(UNCOMPRESSED_BODY.getBytes(StandardCharsets.UTF_8)); + TestContentProvider provider = new TestContentProvider(UNCOMPRESSED_BODY.getBytes()); requestBody = RequestBody.fromContentProvider(provider, "binary/octet-stream"); } @@ -150,6 +155,21 @@ public void sync_streaming_compression_compressesCorrectly() { assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); } + @Test + public void async_streaming_compression_hasCorrectHeaders() { + mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + PutOperationWithStreamingRequestCompressionRequest request = + PutOperationWithStreamingRequestCompressionRequest.builder().build(); + asyncClient.putOperationWithStreamingRequestCompression(request, customAsyncRequestBodyWithoutContentLength(), + AsyncResponseTransformer.toBytes()).join(); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + assertThat(loggedRequest.matchingHeaders("Content-Length")).isEmpty(); + assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); + } + @Test public void sync_nonStreaming_compression_withRetry_compressesCorrectly() { mockHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); @@ -216,6 +236,22 @@ public void sync_streaming_compression_withRetry_compressesCorrectly() { assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); } + @Test + public void async_streaming_compression_withRetry_hasCorrectHeaders() { + mockAsyncHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); + mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + PutOperationWithStreamingRequestCompressionRequest request = + PutOperationWithStreamingRequestCompressionRequest.builder().build(); + asyncClient.putOperationWithStreamingRequestCompression(request, customAsyncRequestBodyWithoutContentLength(), + AsyncResponseTransformer.toBytes()).join(); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + assertThat(loggedRequest.matchingHeaders("Content-Length")).isEmpty(); + assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); + } + private HttpExecuteResponse mockResponse() { return HttpExecuteResponse.builder() .response(SdkHttpResponse.builder().statusCode(200).build()) @@ -228,6 +264,21 @@ private HttpExecuteResponse mockErrorResponse() { .build(); } + protected AsyncRequestBody customAsyncRequestBodyWithoutContentLength() { + return new AsyncRequestBody() { + @Override + public Optional contentLength() { + return Optional.empty(); + } + + @Override + public void subscribe(Subscriber s) { + Flowable.fromPublisher(AsyncRequestBody.fromBytes(UNCOMPRESSED_BODY.getBytes())) + .subscribe(s); + } + }; + } + private static final class TestContentProvider implements ContentStreamProvider { private final byte[] content; private final List createdStreams = new ArrayList<>(); From df597143e142a2085f4ea107d1783348839b6fea Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 11 Aug 2023 09:08:19 -0700 Subject: [PATCH 07/14] Address comments --- .../async/CompressionAsyncRequestBody.java | 14 +- .../pipeline/stages/CompressRequestStage.java | 3 - .../CompressionAsyncRequestBodyTest.java | 115 ++++++++-- .../services/AsyncRequestCompressionTest.java | 205 ++++++++++++++++++ .../services/RequestCompressionTest.java | 114 +--------- .../service/http/MockAsyncHttpClient.java | 62 +++++- 6 files changed, 385 insertions(+), 128 deletions(-) create mode 100644 test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/AsyncRequestCompressionTest.java diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java index 49e7762c35d2..3dbe21e07558 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java @@ -114,33 +114,33 @@ public void subscribe(Subscriber s) { private static final class CompressionSubscriber implements Subscriber { - private final Subscriber wrapped; + private final Subscriber subscriber; private final Compressor compressor; - CompressionSubscriber(Subscriber wrapped, Compressor compressor) { - this.wrapped = wrapped; + CompressionSubscriber(Subscriber subscriber, Compressor compressor) { + this.subscriber = subscriber; this.compressor = compressor; } @Override public void onSubscribe(Subscription subscription) { - wrapped.onSubscribe(subscription); + subscriber.onSubscribe(subscription); } @Override public void onNext(ByteBuffer byteBuffer) { ByteBuffer compressedBuffer = compressor.compress(byteBuffer); - wrapped.onNext(compressedBuffer); + subscriber.onNext(compressedBuffer); } @Override public void onError(Throwable t) { - wrapped.onError(t); + subscriber.onError(t); } @Override public void onComplete() { - wrapped.onComplete(); + subscriber.onComplete(); } } } diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java index 6ca358bd312c..e002697c5f15 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/CompressRequestStage.java @@ -64,7 +64,6 @@ public SdkHttpFullRequest.Builder execute(SdkHttpFullRequest.Builder input, Requ Compressor compressor = resolveCompressorType(context.executionAttributes()); - // non-streaming if (!isStreaming(context)) { compressEntirePayload(input, compressor); updateContentEncodingHeader(input, compressor); @@ -77,10 +76,8 @@ public SdkHttpFullRequest.Builder execute(SdkHttpFullRequest.Builder input, Requ } if (context.requestProvider() == null) { - // sync streaming input.contentStreamProvider(new CompressionContentStreamProvider(input.contentStreamProvider(), compressor)); } else { - // async streaming context.requestProvider(CompressionAsyncRequestBody.builder() .asyncRequestBody(context.requestProvider()) .compressor(compressor) diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java index fb1873d3c57b..cf3a9b4e3c8e 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java @@ -18,31 +18,87 @@ import static org.assertj.core.api.Assertions.assertThat; import io.reactivex.Flowable; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; import java.nio.ByteBuffer; +import java.util.Arrays; import java.util.Optional; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.zip.GZIPInputStream; import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; import org.reactivestreams.Subscriber; import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.internal.compression.Compressor; import software.amazon.awssdk.core.internal.compression.GzipCompressor; +import software.amazon.awssdk.core.internal.util.Mimetype; import software.amazon.awssdk.http.async.SimpleSubscriber; +@RunWith(Parameterized.class) public final class CompressionAsyncRequestBodyTest { - private static final Compressor compressor = new GzipCompressor();; - private final static String TEST_STRING = - "RequestCompressionTest-RequestCompressionTest-RequestCompressionTest-RequestCompressionTest-RequestCompressionTest"; - private static final String EXPECTED_TEST_STRING = - new String(compressor.compress(TEST_STRING.getBytes())); + private static final Compressor compressor = new GzipCompressor(); + private final String testString; + private final AsyncRequestBody provider; - private final AsyncRequestBody provider = CompressionAsyncRequestBody.builder() - .compressor(compressor) - .asyncRequestBody(customAsyncRequestBodyWithoutContentLength()) - .build(); + public CompressionAsyncRequestBodyTest(String testString) { + this.testString = testString; + this.provider = CompressionAsyncRequestBody.builder() + .compressor(compressor) + .asyncRequestBody(customAsyncRequestBodyWithoutContentLength(testString.getBytes())) + .build(); + } + + @Parameterized.Parameters + public static String[] data() { + String[] testStrings = { + createCompressibleStringOfGivenSize(1000), + // chunk size = 128 * 1024 + createCompressibleStringOfGivenSize(130 * 1024), + }; + return testStrings; + } @Test - public void hasCorrectContent() throws InterruptedException { + public void hasCorrectContent() throws Exception { + ByteBuffer byteBuffer = ByteBuffer.allocate(testString.length()); + CountDownLatch done = new CountDownLatch(1); + + Subscriber subscriber = new SimpleSubscriber(buffer -> { + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + byteBuffer.put(bytes); + }) { + @Override + public void onError(Throwable t) { + super.onError(t); + done.countDown(); + } + + @Override + public void onComplete() { + super.onComplete(); + done.countDown(); + } + }; + + provider.subscribe(subscriber); + done.await(10, TimeUnit.SECONDS); + + byte[] retrieved = byteBuffer.array(); + byte[] uncompressed = decompress(retrieved); + assertThat(new String(uncompressed)).isEqualTo(testString); + } + + @Test + public void emptyBytesConstructor_hasEmptyContent() throws Exception { + AsyncRequestBody requestBody = CompressionAsyncRequestBody.builder() + .compressor(compressor) + .asyncRequestBody(AsyncRequestBody.empty()) + .build(); + StringBuilder sb = new StringBuilder(); CountDownLatch done = new CountDownLatch(1); @@ -64,12 +120,43 @@ public void onComplete() { } }; - provider.subscribe(subscriber); + requestBody.subscribe(subscriber); done.await(10, TimeUnit.SECONDS); - assertThat(sb).hasToString(EXPECTED_TEST_STRING); + assertThat(sb).hasToString(""); + assertThat(requestBody.contentType()).isEqualTo(Mimetype.MIMETYPE_OCTET_STREAM); + } + + private static String createCompressibleStringOfGivenSize(int size) { + ByteBuffer data = ByteBuffer.allocate(size); + + byte[] a = new byte[size / 4]; + byte[] b = new byte[size / 4]; + Arrays.fill(a, (byte) 'a'); + Arrays.fill(b, (byte) 'b'); + + data.put(a); + data.put(b); + data.put(a); + data.put(b); + + return new String(data.array()); + } + + private static byte[] decompress(byte[] compressedData) throws IOException { + ByteArrayInputStream bais = new ByteArrayInputStream(compressedData); + GZIPInputStream gzipInputStream = new GZIPInputStream(bais); + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + byte[] buffer = new byte[1024]; + int bytesRead; + while ((bytesRead = gzipInputStream.read(buffer)) != -1) { + baos.write(buffer, 0, bytesRead); + } + gzipInputStream.close(); + byte[] decompressedData = baos.toByteArray(); + return decompressedData; } - protected AsyncRequestBody customAsyncRequestBodyWithoutContentLength() { + private static AsyncRequestBody customAsyncRequestBodyWithoutContentLength(byte[] content) { return new AsyncRequestBody() { @Override public Optional contentLength() { @@ -78,7 +165,7 @@ public Optional contentLength() { @Override public void subscribe(Subscriber s) { - Flowable.fromPublisher(AsyncRequestBody.fromBytes(TEST_STRING.getBytes())) + Flowable.fromPublisher(AsyncRequestBody.fromBytes(content)) .subscribe(s); } }; diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/AsyncRequestCompressionTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/AsyncRequestCompressionTest.java new file mode 100644 index 000000000000..5a8f1f50dbc8 --- /dev/null +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/AsyncRequestCompressionTest.java @@ -0,0 +1,205 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.services; + +import static org.assertj.core.api.Assertions.assertThat; + +import io.reactivex.Flowable; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.time.Duration; +import java.util.Optional; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.reactivestreams.Subscriber; +import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.core.async.AsyncResponseTransformer; +import software.amazon.awssdk.core.internal.compression.Compressor; +import software.amazon.awssdk.core.internal.compression.GzipCompressor; +import software.amazon.awssdk.http.HttpExecuteResponse; +import software.amazon.awssdk.http.SdkHttpFullRequest; +import software.amazon.awssdk.http.SdkHttpResponse; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient; +import software.amazon.awssdk.services.protocolrestjson.model.PutOperationWithRequestCompressionRequest; +import software.amazon.awssdk.services.protocolrestjson.model.PutOperationWithStreamingRequestCompressionRequest; +import software.amazon.awssdk.testutils.service.http.MockAsyncHttpClient; + +public class AsyncRequestCompressionTest { + private static final String UNCOMPRESSED_BODY = + "RequestCompressionTest-RequestCompressionTest-RequestCompressionTest-RequestCompressionTest-RequestCompressionTest"; + private String compressedBody; + private int compressedLen; + private MockAsyncHttpClient mockAsyncHttpClient; + private ProtocolRestJsonAsyncClient asyncClient; + private Compressor compressor; + + @BeforeEach + public void setUp() { + mockAsyncHttpClient = new MockAsyncHttpClient(); + asyncClient = ProtocolRestJsonAsyncClient.builder() + .credentialsProvider(AnonymousCredentialsProvider.create()) + .region(Region.US_EAST_1) + .httpClient(mockAsyncHttpClient) + .build(); + compressor = new GzipCompressor(); + byte[] compressedBodyBytes = compressor.compress(UNCOMPRESSED_BODY.getBytes()); + compressedBody = new String(compressedBodyBytes); + compressedLen = compressedBodyBytes.length; + } + + @AfterEach + public void reset() { + mockAsyncHttpClient.reset(); + } + + @Test + public void asyncNonStreamingOperation_compressionEnabledThresholdOverridden_compressesCorrectly() { + mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + PutOperationWithRequestCompressionRequest request = + PutOperationWithRequestCompressionRequest.builder() + .body(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)) + .overrideConfiguration(o -> o.requestCompressionConfiguration( + c -> c.minimumCompressionThresholdInBytes(1))) + .build(); + + asyncClient.putOperationWithRequestCompression(request); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); + InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); + String loggedBody = new String(SdkBytes.fromInputStream(loggedStream).asByteArray()); + int loggedSize = Integer.valueOf(loggedRequest.firstMatchingHeader("Content-Length").get()); + + assertThat(loggedBody).isEqualTo(compressedBody); + assertThat(loggedSize).isEqualTo(compressedLen); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + } + + @Test + public void asyncNonStreamingOperation_payloadSizeLessThanCompressionThreshold_doesNotCompress() { + mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + PutOperationWithRequestCompressionRequest request = + PutOperationWithRequestCompressionRequest.builder() + .body(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)) + .build(); + + asyncClient.putOperationWithRequestCompression(request); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); + InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); + String loggedBody = new String(SdkBytes.fromInputStream(loggedStream).asByteArray()); + int loggedSize = Integer.valueOf(loggedRequest.firstMatchingHeader("Content-Length").get()); + + assertThat(loggedBody).isEqualTo(UNCOMPRESSED_BODY); + assertThat(loggedSize).isEqualTo(UNCOMPRESSED_BODY.length()); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding")).isEmpty(); + } + + @Test + public void asyncStreamingOperation_compressionEnabled_compressesCorrectly() { + mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + mockAsyncHttpClient.setAsyncRequestBodyLength(compressedBody.length()); + PutOperationWithStreamingRequestCompressionRequest request = + PutOperationWithStreamingRequestCompressionRequest.builder().build(); + asyncClient.putOperationWithStreamingRequestCompression(request, customAsyncRequestBodyWithoutContentLength(), + AsyncResponseTransformer.toBytes()).join(); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); + String loggedBody = new String(mockAsyncHttpClient.getStreamingPayload()); + + assertThat(loggedBody).isEqualTo(compressedBody); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + assertThat(loggedRequest.matchingHeaders("Content-Length")).isEmpty(); + assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); + } + + @Test + public void asyncNonStreamingOperation_compressionEnabledThresholdOverriddenWithRetry_compressesCorrectly() { + mockAsyncHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); + mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + PutOperationWithRequestCompressionRequest request = + PutOperationWithRequestCompressionRequest.builder() + .body(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)) + .overrideConfiguration(o -> o.requestCompressionConfiguration( + c -> c.minimumCompressionThresholdInBytes(1))) + .build(); + + asyncClient.putOperationWithRequestCompression(request); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); + InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); + String loggedBody = new String(SdkBytes.fromInputStream(loggedStream).asByteArray()); + int loggedSize = Integer.valueOf(loggedRequest.firstMatchingHeader("Content-Length").get()); + + assertThat(loggedBody).isEqualTo(compressedBody); + assertThat(loggedSize).isEqualTo(compressedLen); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + } + + @Test + public void asyncStreamingOperation_compressionEnabledWithRetry_compressesCorrectly() { + mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + + mockAsyncHttpClient.setAsyncRequestBodyLength(compressedBody.length()); + PutOperationWithStreamingRequestCompressionRequest request = + PutOperationWithStreamingRequestCompressionRequest.builder().build(); + asyncClient.putOperationWithStreamingRequestCompression(request, customAsyncRequestBodyWithoutContentLength(), + AsyncResponseTransformer.toBytes()).join(); + + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); + String loggedBody = new String(mockAsyncHttpClient.getStreamingPayload()); + + assertThat(loggedBody).isEqualTo(compressedBody); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + assertThat(loggedRequest.matchingHeaders("Content-Length")).isEmpty(); + assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); + } + + private HttpExecuteResponse mockResponse() { + return HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(200).build()) + .build(); + } + + private HttpExecuteResponse mockErrorResponse() { + return HttpExecuteResponse.builder() + .response(SdkHttpResponse.builder().statusCode(500).build()) + .build(); + } + + protected AsyncRequestBody customAsyncRequestBodyWithoutContentLength() { + return new AsyncRequestBody() { + @Override + public Optional contentLength() { + return Optional.empty(); + } + + @Override + public void subscribe(Subscriber s) { + Flowable.fromPublisher(AsyncRequestBody.fromBytes(UNCOMPRESSED_BODY.getBytes())) + .subscribe(s); + } + }; + } +} diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java index 9113c131307a..14cb07125037 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/RequestCompressionTest.java @@ -18,24 +18,18 @@ import static org.assertj.core.api.Assertions.assertThat; import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely; -import io.reactivex.Flowable; import java.io.ByteArrayInputStream; import java.io.FilterInputStream; import java.io.IOException; import java.io.InputStream; -import java.nio.ByteBuffer; import java.time.Duration; import java.util.ArrayList; import java.util.List; -import java.util.Optional; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.AfterEach; -import org.reactivestreams.Subscriber; import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; import software.amazon.awssdk.core.SdkBytes; -import software.amazon.awssdk.core.async.AsyncRequestBody; -import software.amazon.awssdk.core.async.AsyncResponseTransformer; import software.amazon.awssdk.core.internal.compression.Compressor; import software.amazon.awssdk.core.internal.compression.GzipCompressor; import software.amazon.awssdk.core.sync.RequestBody; @@ -45,11 +39,9 @@ import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.http.SdkHttpResponse; import software.amazon.awssdk.regions.Region; -import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient; import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient; import software.amazon.awssdk.services.protocolrestjson.model.PutOperationWithRequestCompressionRequest; import software.amazon.awssdk.services.protocolrestjson.model.PutOperationWithStreamingRequestCompressionRequest; -import software.amazon.awssdk.testutils.service.http.MockAsyncHttpClient; import software.amazon.awssdk.testutils.service.http.MockSyncHttpClient; public class RequestCompressionTest { @@ -58,28 +50,20 @@ public class RequestCompressionTest { private String compressedBody; private int compressedLen; private MockSyncHttpClient mockHttpClient; - private MockAsyncHttpClient mockAsyncHttpClient; private ProtocolRestJsonClient syncClient; - private ProtocolRestJsonAsyncClient asyncClient; private Compressor compressor; private RequestBody requestBody; @BeforeEach public void setUp() { mockHttpClient = new MockSyncHttpClient(); - mockAsyncHttpClient = new MockAsyncHttpClient(); syncClient = ProtocolRestJsonClient.builder() .credentialsProvider(AnonymousCredentialsProvider.create()) .region(Region.US_EAST_1) .httpClient(mockHttpClient) .build(); - asyncClient = ProtocolRestJsonAsyncClient.builder() - .credentialsProvider(AnonymousCredentialsProvider.create()) - .region(Region.US_EAST_1) - .httpClient(mockAsyncHttpClient) - .build(); compressor = new GzipCompressor(); - byte[] compressedBodyBytes = compressor.compress(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)).asByteArray(); + byte[] compressedBodyBytes = compressor.compress(UNCOMPRESSED_BODY.getBytes()); compressedLen = compressedBodyBytes.length; compressedBody = new String(compressedBodyBytes); TestContentProvider provider = new TestContentProvider(UNCOMPRESSED_BODY.getBytes()); @@ -89,11 +73,10 @@ public void setUp() { @AfterEach public void reset() { mockHttpClient.reset(); - mockAsyncHttpClient.reset(); } @Test - public void sync_nonStreaming_compression_compressesCorrectly() { + public void syncNonStreamingOperation_compressionEnabledThresholdOverridden_compressesCorrectly() { mockHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); PutOperationWithRequestCompressionRequest request = @@ -115,30 +98,25 @@ public void sync_nonStreaming_compression_compressesCorrectly() { } @Test - public void async_nonStreaming_compression_compressesCorrectly() { - mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); + public void syncNonStreamingOperation_payloadSizeLessThanCompressionThreshold_doesNotCompress() { + mockHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); PutOperationWithRequestCompressionRequest request = PutOperationWithRequestCompressionRequest.builder() .body(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)) - .overrideConfiguration(o -> o.requestCompressionConfiguration( - c -> c.minimumCompressionThresholdInBytes(1))) .build(); + syncClient.putOperationWithRequestCompression(request); - asyncClient.putOperationWithRequestCompression(request); - - SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); + SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockHttpClient.getLastRequest(); InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); String loggedBody = new String(SdkBytes.fromInputStream(loggedStream).asByteArray()); - int loggedSize = Integer.valueOf(loggedRequest.firstMatchingHeader("Content-Length").get()); - assertThat(loggedBody).isEqualTo(compressedBody); - assertThat(loggedSize).isEqualTo(compressedLen); - assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); + assertThat(loggedBody).isEqualTo(UNCOMPRESSED_BODY); + assertThat(loggedRequest.firstMatchingHeader("Content-encoding")).isEmpty(); } @Test - public void sync_streaming_compression_compressesCorrectly() { + public void syncStreamingOperation_compressionEnabled_compressesCorrectly() { mockHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); PutOperationWithStreamingRequestCompressionRequest request = @@ -156,22 +134,7 @@ public void sync_streaming_compression_compressesCorrectly() { } @Test - public void async_streaming_compression_hasCorrectHeaders() { - mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); - - PutOperationWithStreamingRequestCompressionRequest request = - PutOperationWithStreamingRequestCompressionRequest.builder().build(); - asyncClient.putOperationWithStreamingRequestCompression(request, customAsyncRequestBodyWithoutContentLength(), - AsyncResponseTransformer.toBytes()).join(); - - SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); - assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); - assertThat(loggedRequest.matchingHeaders("Content-Length")).isEmpty(); - assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); - } - - @Test - public void sync_nonStreaming_compression_withRetry_compressesCorrectly() { + public void syncNonStreamingOperation_compressionEnabledThresholdOverriddenWithRetry_compressesCorrectly() { mockHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); mockHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); @@ -194,31 +157,7 @@ public void sync_nonStreaming_compression_withRetry_compressesCorrectly() { } @Test - public void async_nonStreaming_compression_withRetry_compressesCorrectly() { - mockAsyncHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); - mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); - - PutOperationWithRequestCompressionRequest request = - PutOperationWithRequestCompressionRequest.builder() - .body(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)) - .overrideConfiguration(o -> o.requestCompressionConfiguration( - c -> c.minimumCompressionThresholdInBytes(1))) - .build(); - - asyncClient.putOperationWithRequestCompression(request); - - SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); - InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); - String loggedBody = new String(SdkBytes.fromInputStream(loggedStream).asByteArray()); - int loggedSize = Integer.valueOf(loggedRequest.firstMatchingHeader("Content-Length").get()); - - assertThat(loggedBody).isEqualTo(compressedBody); - assertThat(loggedSize).isEqualTo(compressedLen); - assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); - } - - @Test - public void sync_streaming_compression_withRetry_compressesCorrectly() { + public void syncStreamingOperation_compressionEnabledWithRetry_compressesCorrectly() { mockHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); mockHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); @@ -236,22 +175,6 @@ public void sync_streaming_compression_withRetry_compressesCorrectly() { assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); } - @Test - public void async_streaming_compression_withRetry_hasCorrectHeaders() { - mockAsyncHttpClient.stubNextResponse(mockErrorResponse(), Duration.ofMillis(500)); - mockAsyncHttpClient.stubNextResponse(mockResponse(), Duration.ofMillis(500)); - - PutOperationWithStreamingRequestCompressionRequest request = - PutOperationWithStreamingRequestCompressionRequest.builder().build(); - asyncClient.putOperationWithStreamingRequestCompression(request, customAsyncRequestBodyWithoutContentLength(), - AsyncResponseTransformer.toBytes()).join(); - - SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); - assertThat(loggedRequest.firstMatchingHeader("Content-encoding").get()).isEqualTo("gzip"); - assertThat(loggedRequest.matchingHeaders("Content-Length")).isEmpty(); - assertThat(loggedRequest.firstMatchingHeader("Transfer-Encoding").get()).isEqualTo("chunked"); - } - private HttpExecuteResponse mockResponse() { return HttpExecuteResponse.builder() .response(SdkHttpResponse.builder().statusCode(200).build()) @@ -264,21 +187,6 @@ private HttpExecuteResponse mockErrorResponse() { .build(); } - protected AsyncRequestBody customAsyncRequestBodyWithoutContentLength() { - return new AsyncRequestBody() { - @Override - public Optional contentLength() { - return Optional.empty(); - } - - @Override - public void subscribe(Subscriber s) { - Flowable.fromPublisher(AsyncRequestBody.fromBytes(UNCOMPRESSED_BODY.getBytes())) - .subscribe(s); - } - }; - } - private static final class TestContentProvider implements ContentStreamProvider { private final byte[] content; private final List createdStreams = new ArrayList<>(); diff --git a/test/service-test-utils/src/main/java/software/amazon/awssdk/testutils/service/http/MockAsyncHttpClient.java b/test/service-test-utils/src/main/java/software/amazon/awssdk/testutils/service/http/MockAsyncHttpClient.java index 8a3f62f7838e..16a7732cb186 100644 --- a/test/service-test-utils/src/main/java/software/amazon/awssdk/testutils/service/http/MockAsyncHttpClient.java +++ b/test/service-test-utils/src/main/java/software/amazon/awssdk/testutils/service/http/MockAsyncHttpClient.java @@ -26,6 +26,7 @@ import java.util.List; import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicInteger; @@ -50,6 +51,8 @@ public final class MockAsyncHttpClient implements SdkAsyncHttpClient, MockHttpCl private final List> responses = new LinkedList<>(); private final AtomicInteger responseIndex = new AtomicInteger(0); private final ExecutorService executor; + private int asyncRequestBodyLength = -1; + private byte[] streamingPayload; public MockAsyncHttpClient() { this.executor = Executors.newFixedThreadPool(3); @@ -66,6 +69,11 @@ public CompletableFuture execute(AsyncExecuteRequest request) { request.responseHandler().onHeaders(nextResponse.httpResponse()); CompletableFuture.runAsync(() -> request.responseHandler().onStream(new ResponsePublisher(content, index)), executor); + + if (asyncRequestBodyLength > 0) { + captureStreamingPayload(request.requestContentPublisher()); + } + return CompletableFuture.completedFuture(null); } @@ -122,7 +130,28 @@ public void stubResponses(HttpExecuteResponse... responses) { this.responseIndex.set(0); } - private class ResponsePublisher implements SdkHttpContentPublisher { + /** + * Enable capturing the streaming payload by setting the length of the AsyncRequestBody. + */ + public void setAsyncRequestBodyLength(int asyncRequestBodyLength) { + this.asyncRequestBodyLength = asyncRequestBodyLength; + } + + private void captureStreamingPayload(SdkHttpContentPublisher publisher) { + ByteBuffer byteBuffer = ByteBuffer.allocate(asyncRequestBodyLength); + Subscriber subscriber = new CapturingSubscriber(byteBuffer); + publisher.subscribe(subscriber); + streamingPayload = byteBuffer.array(); + } + + /** + * Returns the streaming payload byte array, if the asyncRequestBodyLength was set correctly. Otherwise, returns null. + */ + public byte[] getStreamingPayload() { + return streamingPayload.clone(); + } + + private final class ResponsePublisher implements SdkHttpContentPublisher { private final byte[] content; private final int index; @@ -165,4 +194,35 @@ public void cancel() { }); } } + + private static class CapturingSubscriber implements Subscriber { + private ByteBuffer byteBuffer; + private CountDownLatch done = new CountDownLatch(1); + + CapturingSubscriber(ByteBuffer byteBuffer) { + this.byteBuffer = byteBuffer; + } + + @Override + public void onSubscribe(Subscription subscription) { + subscription.request(Long.MAX_VALUE); + } + + @Override + public void onNext(ByteBuffer buffer) { + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + byteBuffer.put(bytes); + } + + @Override + public void onError(Throwable t) { + done.countDown(); + } + + @Override + public void onComplete() { + done.countDown(); + } + } } From 369be3ec8f96a681ed24546997917c2715436412 Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 11 Aug 2023 14:58:42 -0700 Subject: [PATCH 08/14] Refactor ChunkBuffer class --- .../core/internal/async/ChunkBuffer.java | 41 ++++++++-- .../async/ChunkBufferWithUnknownLength.java | 80 ------------------- .../async/CompressionAsyncRequestBody.java | 6 +- .../awssdk/core/async/ChunkBufferTest.java | 5 -- 4 files changed, 37 insertions(+), 95 deletions(-) delete mode 100644 core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBufferWithUnknownLength.java diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java index 93d6d09578a6..3f95c043daf9 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java @@ -23,7 +23,6 @@ import java.util.concurrent.atomic.AtomicLong; import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.utils.BinaryUtils; -import software.amazon.awssdk.utils.Validate; import software.amazon.awssdk.utils.builder.SdkBuilder; /** @@ -31,26 +30,34 @@ */ @SdkInternalApi public final class ChunkBuffer { - private final AtomicLong remainingBytes; + private AtomicLong remainingBytes; private final ByteBuffer currentBuffer; - private final int bufferSize; + private int bufferSize; private ChunkBuffer(Long totalBytes, Integer bufferSize) { - Validate.notNull(totalBytes, "The totalBytes must not be null"); - int chunkSize = bufferSize != null ? bufferSize : DEFAULT_ASYNC_CHUNK_SIZE; this.bufferSize = chunkSize; this.currentBuffer = ByteBuffer.allocate(chunkSize); - this.remainingBytes = new AtomicLong(totalBytes); + + if (totalBytes != null) { + this.remainingBytes = new AtomicLong(totalBytes); + } } public static Builder builder() { return new DefaultBuilder(); } + public synchronized Iterable bufferAndCreateChunks(ByteBuffer buffer) { + if (remainingBytes == null) { + return bufferAndCreateChunksWithUnknownLength(buffer); + } else { + return bufferAndCreateChunksWithKnownLength(buffer); + } + } // currentBuffer and bufferedList can get over written if concurrent Threads calls this method at the same time. - public synchronized Iterable bufferAndCreateChunks(ByteBuffer buffer) { + private synchronized Iterable bufferAndCreateChunksWithKnownLength(ByteBuffer buffer) { int startPosition = 0; List bufferedList = new ArrayList<>(); int currentBytesRead = buffer.remaining(); @@ -97,6 +104,26 @@ public synchronized Iterable bufferAndCreateChunks(ByteBuffer buffer return bufferedList; } + private synchronized Iterable bufferAndCreateChunksWithUnknownLength(ByteBuffer buffer) { + List bufferedList = new ArrayList<>(); + while (buffer.hasRemaining()) { + int bytesToCopy = Math.min(buffer.remaining(), currentBuffer.remaining()); + byte[] bytes = new byte[bytesToCopy]; + buffer.get(bytes); + currentBuffer.put(bytes); + + if (!currentBuffer.hasRemaining() || !buffer.hasRemaining()) { + currentBuffer.flip(); + ByteBuffer bufferToSend = ByteBuffer.allocate(currentBuffer.limit()); + bufferToSend.put(currentBuffer); + bufferToSend.flip(); + bufferedList.add(bufferToSend); + currentBuffer.clear(); + } + } + return bufferedList; + } + public interface Builder extends SdkBuilder { Builder bufferSize(int bufferSize); diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBufferWithUnknownLength.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBufferWithUnknownLength.java deleted file mode 100644 index 5a5cfca6dc9e..000000000000 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBufferWithUnknownLength.java +++ /dev/null @@ -1,80 +0,0 @@ -/* - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"). - * You may not use this file except in compliance with the License. - * A copy of the License is located at - * - * http://aws.amazon.com/apache2.0 - * - * or in the "license" file accompanying this file. This file is distributed - * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either - * express or implied. See the License for the specific language governing - * permissions and limitations under the License. - */ - -package software.amazon.awssdk.core.internal.async; - -import static software.amazon.awssdk.core.HttpChecksumConstant.DEFAULT_ASYNC_CHUNK_SIZE; - -import java.nio.ByteBuffer; -import java.util.ArrayList; -import java.util.List; -import software.amazon.awssdk.annotations.SdkInternalApi; -import software.amazon.awssdk.utils.builder.SdkBuilder; - -/** - * Class that will buffer incoming BufferBytes with unknown total length to chunks of bufferSize - */ -@SdkInternalApi -public final class ChunkBufferWithUnknownLength { - private ByteBuffer currentBuffer; - - private ChunkBufferWithUnknownLength(Integer bufferSize) { - int chunkSize = bufferSize != null ? bufferSize : DEFAULT_ASYNC_CHUNK_SIZE; - this.currentBuffer = ByteBuffer.allocate(chunkSize); - } - - public static Builder builder() { - return new DefaultBuilder(); - } - - public synchronized Iterable bufferAndCreateChunks(ByteBuffer buffer) { - List bufferedList = new ArrayList<>(); - while (buffer.hasRemaining()) { - int bytesToCopy = Math.min(buffer.remaining(), currentBuffer.remaining()); - byte[] bytes = new byte[bytesToCopy]; - buffer.get(bytes); - currentBuffer.put(bytes); - - if (!currentBuffer.hasRemaining() || !buffer.hasRemaining()) { - currentBuffer.flip(); - ByteBuffer bufferToSend = ByteBuffer.allocate(currentBuffer.limit()); - bufferToSend.put(currentBuffer); - bufferToSend.flip(); - bufferedList.add(bufferToSend); - currentBuffer.clear(); - } - } - return bufferedList; - } - - public interface Builder extends SdkBuilder { - Builder bufferSize(int bufferSize); - } - - private static final class DefaultBuilder implements Builder { - private Integer bufferSize; - - @Override - public ChunkBufferWithUnknownLength build() { - return new ChunkBufferWithUnknownLength(bufferSize); - } - - @Override - public Builder bufferSize(int bufferSize) { - this.bufferSize = bufferSize; - return this; - } - } -} diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java index 3dbe21e07558..f2036889f944 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java @@ -104,9 +104,9 @@ public String contentType() { public void subscribe(Subscriber s) { Validate.notNull(s, "Subscription MUST NOT be null."); - ChunkBufferWithUnknownLength chunkBuffer = ChunkBufferWithUnknownLength.builder() - .bufferSize(COMPRESSION_CHUNK_SIZE) - .build(); + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(COMPRESSION_CHUNK_SIZE) + .build(); wrapped.flatMapIterable(chunkBuffer::bufferAndCreateChunks) .subscribe(new CompressionSubscriber(s, compressor)); diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java index 8b73402dc468..822b05f19cbc 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java @@ -35,11 +35,6 @@ class ChunkBufferTest { - @Test - void builderWithNoTotalSize() { - assertThatThrownBy(() -> ChunkBuffer.builder().build()).isInstanceOf(NullPointerException.class); - } - @Test void numberOfChunkMultipleOfTotalBytes() { String inputString = StringUtils.repeat("*", 25); From 280ffc0d30e3adf387812c948a550316dd269428 Mon Sep 17 00:00:00 2001 From: hdavidh Date: Mon, 21 Aug 2023 12:15:13 -0700 Subject: [PATCH 09/14] Address comments --- .../core/internal/async/ChunkBuffer.java | 34 +++--- .../async/CompressionAsyncRequestBody.java | 63 ++++------- .../awssdk/core/async/ChunkBufferTest.java | 107 ++++++++++++++++-- .../CompressionAsyncRequestBodyTckTest.java | 89 +++++++++++++++ .../CompressionAsyncRequestBodyTest.java | 56 ++++----- 5 files changed, 255 insertions(+), 94 deletions(-) create mode 100644 core/sdk-core/src/test/java/software/amazon/awssdk/core/async/CompressionAsyncRequestBodyTckTest.java diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java index 5b6f15f81e56..86a87c6e38b8 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java @@ -35,7 +35,7 @@ public final class ChunkBuffer { private AtomicLong transferredBytes; private final ByteBuffer currentBuffer; private final int chunkSize; - private long totalBytes; + private Long totalBytes; private ChunkBuffer(Long totalBytes, Integer bufferSize) { int chunkSize = bufferSize != null ? bufferSize : DEFAULT_ASYNC_CHUNK_SIZE; @@ -57,11 +57,18 @@ public static Builder builder() { * be buffered. */ public synchronized Iterable split(ByteBuffer inputByteBuffer) { - if (!inputByteBuffer.hasRemaining()) { return Collections.singletonList(inputByteBuffer); } + if (totalBytes != null) { + return splitWithKnownLength(inputByteBuffer); + } else { + return splitWithUnknownLength(inputByteBuffer); + } + } + + private synchronized Iterable splitWithKnownLength(ByteBuffer inputByteBuffer) { List byteBuffers = new ArrayList<>(); // If current buffer is not empty, fill the buffer first. @@ -86,19 +93,17 @@ public synchronized Iterable split(ByteBuffer inputByteBuffer) { return byteBuffers; } - /** - * Split the input {@link ByteBuffer} into multiple smaller {@link ByteBuffer}s, each of which contains {@link #chunkSize} - * worth of bytes, when the total length is unknown. - */ - public synchronized Iterable splitWithUnknownLength(ByteBuffer buffer) { + private synchronized Iterable splitWithUnknownLength(ByteBuffer inputByteBuffer) { + boolean isLastChunk = inputByteBuffer.remaining() != chunkSize; List bufferedList = new ArrayList<>(); - while (buffer.hasRemaining()) { - int bytesToCopy = Math.min(buffer.remaining(), currentBuffer.remaining()); + + while (inputByteBuffer.hasRemaining()) { + int bytesToCopy = Math.min(inputByteBuffer.remaining(), currentBuffer.remaining()); byte[] bytes = new byte[bytesToCopy]; - buffer.get(bytes); + inputByteBuffer.get(bytes); currentBuffer.put(bytes); - if (!currentBuffer.hasRemaining() || !buffer.hasRemaining()) { + if (!currentBuffer.hasRemaining()) { currentBuffer.flip(); ByteBuffer bufferToSend = ByteBuffer.allocate(currentBuffer.limit()); bufferToSend.put(currentBuffer); @@ -107,6 +112,11 @@ public synchronized Iterable splitWithUnknownLength(ByteBuffer buffe currentBuffer.clear(); } } + + if (isLastChunk && currentBuffer.position() != 0) { + bufferedList.add((ByteBuffer) currentBuffer.flip()); + } + return bufferedList; } @@ -173,8 +183,6 @@ public interface Builder extends SdkBuilder { Builder bufferSize(int bufferSize); Builder totalBytes(long totalBytes); - - } private static final class DefaultBuilder implements Builder { diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java index 5445c9c1b4f1..0f520a38a04a 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java @@ -18,7 +18,6 @@ import java.nio.ByteBuffer; import java.util.Optional; import org.reactivestreams.Subscriber; -import org.reactivestreams.Subscription; import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.internal.compression.Compressor; @@ -35,12 +34,12 @@ public class CompressionAsyncRequestBody implements AsyncRequestBody { private static final int COMPRESSION_CHUNK_SIZE = 128 * 1024; private final AsyncRequestBody wrapped; private final Compressor compressor; + private final int chunkSize; private CompressionAsyncRequestBody(DefaultBuilder builder) { - Validate.notNull(builder.asyncRequestBody, "wrapped AsyncRequestBody cannot be null"); - Validate.notNull(builder.compressor, "compressor cannot be null"); - this.wrapped = builder.asyncRequestBody; - this.compressor = builder.compressor; + this.wrapped = Validate.paramNotNull(builder.asyncRequestBody, "asyncRequestBody"); + this.compressor = Validate.paramNotNull(builder.compressor, "compressor"); + this.chunkSize = builder.chunkSize != null ? builder.chunkSize : COMPRESSION_CHUNK_SIZE; } /** @@ -65,12 +64,20 @@ public interface Builder extends SdkBuilder contentLength() { - return Optional.empty(); + return wrapped.contentLength(); } @Override @@ -105,42 +118,10 @@ public void subscribe(Subscriber s) { Validate.notNull(s, "Subscription MUST NOT be null."); ChunkBuffer chunkBuffer = ChunkBuffer.builder() - .bufferSize(COMPRESSION_CHUNK_SIZE) + .bufferSize(chunkSize) .build(); - wrapped.flatMapIterable(chunkBuffer::splitWithUnknownLength) - .subscribe(new CompressionSubscriber(s, compressor)); - } - - private static final class CompressionSubscriber implements Subscriber { - - private final Subscriber subscriber; - private final Compressor compressor; - - CompressionSubscriber(Subscriber subscriber, Compressor compressor) { - this.subscriber = subscriber; - this.compressor = compressor; - } - - @Override - public void onSubscribe(Subscription subscription) { - subscriber.onSubscribe(subscription); - } - - @Override - public void onNext(ByteBuffer byteBuffer) { - ByteBuffer compressedBuffer = compressor.compress(byteBuffer); - subscriber.onNext(compressedBuffer); - } - - @Override - public void onError(Throwable t) { - subscriber.onError(t); - } - - @Override - public void onComplete() { - subscriber.onComplete(); - } + wrapped.flatMapIterable(chunkBuffer::split) + .map(compressor::compress).subscribe(s); } } diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java index 61b4fd05e70f..18aaddcb17ba 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java @@ -42,12 +42,51 @@ class ChunkBufferTest { @ParameterizedTest @ValueSource(ints = {1, 6, 10, 23, 25}) - void numberOfChunk_Not_MultipleOfTotalBytes(int totalBytes) { + void numberOfChunk_Not_MultipleOfTotalBytes_KnownLength(int totalBytes) { int bufferSize = 5; String inputString = RandomStringUtils.randomAscii(totalBytes); - ChunkBuffer chunkBuffer = - ChunkBuffer.builder().bufferSize(bufferSize).totalBytes(inputString.getBytes(StandardCharsets.UTF_8).length).build(); + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(bufferSize) + .totalBytes(inputString.getBytes(StandardCharsets.UTF_8).length) + .build(); + Iterable byteBuffers = + chunkBuffer.split(ByteBuffer.wrap(inputString.getBytes(StandardCharsets.UTF_8))); + + AtomicInteger index = new AtomicInteger(0); + int count = (int) Math.ceil(totalBytes / (double) bufferSize); + int remainder = totalBytes % bufferSize; + + byteBuffers.forEach(r -> { + int i = index.get(); + + try (ByteArrayInputStream inputStream = new ByteArrayInputStream(inputString.getBytes(StandardCharsets.UTF_8))) { + byte[] expected; + if (i == count - 1 && remainder != 0) { + expected = new byte[remainder]; + } else { + expected = new byte[bufferSize]; + } + inputStream.skip(i * bufferSize); + inputStream.read(expected); + byte[] actualBytes = BinaryUtils.copyBytesFrom(r); + assertThat(actualBytes).isEqualTo(expected); + index.incrementAndGet(); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + } + + @ParameterizedTest + @ValueSource(ints = {1, 6, 10, 23, 25}) + void numberOfChunk_Not_MultipleOfTotalBytes_UnknownLength(int totalBytes) { + int bufferSize = 5; + + String inputString = RandomStringUtils.randomAscii(totalBytes); + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(bufferSize) + .build(); Iterable byteBuffers = chunkBuffer.split(ByteBuffer.wrap(inputString.getBytes(StandardCharsets.UTF_8))); @@ -77,10 +116,12 @@ void numberOfChunk_Not_MultipleOfTotalBytes(int totalBytes) { } @Test - void zeroTotalBytesAsInput_returnsZeroByte() { + void zeroTotalBytesAsInput_returnsZeroByte_KnownLength() { byte[] zeroByte = new byte[0]; - ChunkBuffer chunkBuffer = - ChunkBuffer.builder().bufferSize(5).totalBytes(zeroByte.length).build(); + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(5) + .totalBytes(zeroByte.length) + .build(); Iterable byteBuffers = chunkBuffer.split(ByteBuffer.wrap(zeroByte)); @@ -92,13 +133,54 @@ void zeroTotalBytesAsInput_returnsZeroByte() { } @Test - void emptyAllocatedBytes_returnSameNumberOfEmptyBytes() { + void zeroTotalBytesAsInput_returnsZeroByte_UnknownLength() { + byte[] zeroByte = new byte[0]; + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(5) + .build(); + Iterable byteBuffers = + chunkBuffer.split(ByteBuffer.wrap(zeroByte)); + + AtomicInteger iteratedCounts = new AtomicInteger(); + byteBuffers.forEach(r -> { + iteratedCounts.getAndIncrement(); + }); + assertThat(iteratedCounts.get()).isEqualTo(1); + } + + @Test + void emptyAllocatedBytes_returnSameNumberOfEmptyBytes_knownLength() { + int totalBytes = 17; + int bufferSize = 5; + ByteBuffer wrap = ByteBuffer.allocate(totalBytes); + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(bufferSize) + .totalBytes(wrap.remaining()) + .build(); + Iterable byteBuffers = + chunkBuffer.split(wrap); + + AtomicInteger iteratedCounts = new AtomicInteger(); + byteBuffers.forEach(r -> { + iteratedCounts.getAndIncrement(); + if (iteratedCounts.get() * bufferSize < totalBytes) { + // array of empty bytes + assertThat(BinaryUtils.copyBytesFrom(r)).isEqualTo(ByteBuffer.allocate(bufferSize).array()); + } else { + assertThat(BinaryUtils.copyBytesFrom(r)).isEqualTo(ByteBuffer.allocate(totalBytes % bufferSize).array()); + } + }); + assertThat(iteratedCounts.get()).isEqualTo(4); + } + @Test + void emptyAllocatedBytes_returnSameNumberOfEmptyBytes_unknownLength() { int totalBytes = 17; int bufferSize = 5; ByteBuffer wrap = ByteBuffer.allocate(totalBytes); - ChunkBuffer chunkBuffer = - ChunkBuffer.builder().bufferSize(bufferSize).totalBytes(wrap.remaining()).build(); + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(bufferSize) + .build(); Iterable byteBuffers = chunkBuffer.split(wrap); @@ -152,8 +234,10 @@ void concurrentTreads_calling_bufferAndCreateChunks() throws ExecutionException, int threads = 8; ByteBuffer wrap = ByteBuffer.allocate(totalBytes); - ChunkBuffer chunkBuffer = - ChunkBuffer.builder().bufferSize(bufferSize).totalBytes(wrap.remaining() * threads).build(); + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(bufferSize) + .totalBytes(wrap.remaining() * threads) + .build(); ExecutorService service = Executors.newFixedThreadPool(threads); @@ -192,5 +276,4 @@ void concurrentTreads_calling_bufferAndCreateChunks() throws ExecutionException, assertThat(remainderBytesBuffers.get()).isOne(); assertThat(otherSizeBuffers.get()).isZero(); } - } diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/CompressionAsyncRequestBodyTckTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/CompressionAsyncRequestBodyTckTest.java new file mode 100644 index 000000000000..dd0cc312ad1f --- /dev/null +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/CompressionAsyncRequestBodyTckTest.java @@ -0,0 +1,89 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.core.async; + +import io.reactivex.Flowable; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Optional; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.tck.PublisherVerification; +import org.reactivestreams.tck.TestEnvironment; +import software.amazon.awssdk.core.internal.async.CompressionAsyncRequestBody; +import software.amazon.awssdk.core.internal.compression.Compressor; +import software.amazon.awssdk.core.internal.compression.GzipCompressor; + +public class CompressionAsyncRequestBodyTckTest extends PublisherVerification { + + private static final int MAX_ELEMENTS = 1000; + private static final int CHUNK_SIZE = 128 * 1024; + private static final Compressor compressor = new GzipCompressor(); + + public CompressionAsyncRequestBodyTckTest() { + super(new TestEnvironment()); + } + + @Override + public long maxElementsFromPublisher() { + return MAX_ELEMENTS; + } + + @Override + public Publisher createPublisher(long n) { + return CompressionAsyncRequestBody.builder() + .asyncRequestBody(customAsyncRequestBodyWithoutContentLength(createCompressibleStringOfNChunks(n).getBytes())) + .compressor(compressor) + .build(); + } + + @Override + public Publisher createFailedPublisher() { + return null; + } + + private static String createCompressibleStringOfNChunks(long nChunks) { + int size = Math.toIntExact(nChunks * CHUNK_SIZE); + ByteBuffer data = ByteBuffer.allocate(size); + + byte[] a = new byte[size / 4]; + byte[] b = new byte[size / 4]; + Arrays.fill(a, (byte) 'a'); + Arrays.fill(b, (byte) 'b'); + + data.put(a); + data.put(b); + data.put(a); + data.put(b); + + return new String(data.array()); + } + + private static AsyncRequestBody customAsyncRequestBodyWithoutContentLength(byte[] content) { + return new AsyncRequestBody() { + @Override + public Optional contentLength() { + return Optional.empty(); + } + + @Override + public void subscribe(Subscriber s) { + Flowable.fromPublisher(AsyncRequestBody.fromBytes(content)) + .subscribe(s); + } + }; + } +} diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java index cf3a9b4e3c8e..ffb15e282a13 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBodyTest.java @@ -26,10 +26,11 @@ import java.util.Optional; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.zip.GZIPInputStream; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.reactivestreams.Subscriber; import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.internal.compression.Compressor; @@ -37,39 +38,37 @@ import software.amazon.awssdk.core.internal.util.Mimetype; import software.amazon.awssdk.http.async.SimpleSubscriber; -@RunWith(Parameterized.class) public final class CompressionAsyncRequestBodyTest { private static final Compressor compressor = new GzipCompressor(); - private final String testString; - private final AsyncRequestBody provider; - - public CompressionAsyncRequestBodyTest(String testString) { - this.testString = testString; - this.provider = CompressionAsyncRequestBody.builder() - .compressor(compressor) - .asyncRequestBody(customAsyncRequestBodyWithoutContentLength(testString.getBytes())) - .build(); - } - @Parameterized.Parameters - public static String[] data() { - String[] testStrings = { - createCompressibleStringOfGivenSize(1000), - // chunk size = 128 * 1024 - createCompressibleStringOfGivenSize(130 * 1024), - }; - return testStrings; - } + @ParameterizedTest + @ValueSource(ints = {80, 1000}) + public void hasCorrectContent(int bodySize) throws Exception { + String testString = createCompressibleStringOfGivenSize(bodySize); + byte[] testBytes = testString.getBytes(); + int chunkSize = 133; + AsyncRequestBody provider = CompressionAsyncRequestBody.builder() + .compressor(compressor) + .asyncRequestBody(customAsyncRequestBodyWithoutContentLength(testBytes)) + .chunkSize(chunkSize) + .build(); - @Test - public void hasCorrectContent() throws Exception { ByteBuffer byteBuffer = ByteBuffer.allocate(testString.length()); CountDownLatch done = new CountDownLatch(1); + AtomicInteger pos = new AtomicInteger(); Subscriber subscriber = new SimpleSubscriber(buffer -> { byte[] bytes = new byte[buffer.remaining()]; buffer.get(bytes); byteBuffer.put(bytes); + + // verify each chunk + byte[] chunkToVerify = new byte[chunkSize]; + System.arraycopy(testBytes, pos.get(), chunkToVerify, 0, chunkSize); + chunkToVerify = compressor.compress(chunkToVerify); + + assertThat(bytes).isEqualTo(chunkToVerify); + pos.addAndGet(chunkSize); }) { @Override public void onError(Throwable t) { @@ -99,13 +98,13 @@ public void emptyBytesConstructor_hasEmptyContent() throws Exception { .asyncRequestBody(AsyncRequestBody.empty()) .build(); - StringBuilder sb = new StringBuilder(); + ByteBuffer byteBuffer = ByteBuffer.allocate(0); CountDownLatch done = new CountDownLatch(1); Subscriber subscriber = new SimpleSubscriber(buffer -> { byte[] bytes = new byte[buffer.remaining()]; buffer.get(bytes); - sb.append(new String(bytes)); + byteBuffer.put(bytes); }) { @Override public void onError(Throwable t) { @@ -122,7 +121,8 @@ public void onComplete() { requestBody.subscribe(subscriber); done.await(10, TimeUnit.SECONDS); - assertThat(sb).hasToString(""); + assertThat(byteBuffer.array()).isEmpty(); + assertThat(byteBuffer.array()).isEqualTo(new byte[0]); assertThat(requestBody.contentType()).isEqualTo(Mimetype.MIMETYPE_OCTET_STREAM); } From 933e2596cca48520ae30c4dde12122a0731974a8 Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 25 Aug 2023 16:28:46 -0700 Subject: [PATCH 10/14] Address comments --- .../core/internal/async/ChunkBuffer.java | 83 ++++++++----------- .../async/CompressionAsyncRequestBody.java | 72 +++++++++++----- .../awssdk/core/async/ChunkBufferTest.java | 52 +++++++++++- .../CompressionAsyncRequestBodyTckTest.java | 58 +++++++++---- 4 files changed, 178 insertions(+), 87 deletions(-) diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java index 86a87c6e38b8..f14f73b2466a 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java @@ -21,13 +21,16 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Optional; import java.util.concurrent.atomic.AtomicLong; import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.utils.Logger; import software.amazon.awssdk.utils.builder.SdkBuilder; /** - * Class that will buffer incoming BufferBytes of totalBytes length to chunks of bufferSize* + * Class that will buffer incoming BufferBytes to chunks of bufferSize. + * If totalBytes is not provided, i.e. content-length is unknown, {@code getBufferedData()} should be used in the Subscriber's + * {@code onComplete} to check for a final chunk that is smaller than the chunk size, and send if present. */ @SdkInternalApi public final class ChunkBuffer { @@ -36,14 +39,15 @@ public final class ChunkBuffer { private final ByteBuffer currentBuffer; private final int chunkSize; private Long totalBytes; + private boolean isCompleted; private ChunkBuffer(Long totalBytes, Integer bufferSize) { int chunkSize = bufferSize != null ? bufferSize : DEFAULT_ASYNC_CHUNK_SIZE; this.chunkSize = chunkSize; this.currentBuffer = ByteBuffer.allocate(chunkSize); + this.transferredBytes = new AtomicLong(0); if (totalBytes != null) { this.totalBytes = totalBytes; - this.transferredBytes = new AtomicLong(0); } } @@ -61,14 +65,6 @@ public synchronized Iterable split(ByteBuffer inputByteBuffer) { return Collections.singletonList(inputByteBuffer); } - if (totalBytes != null) { - return splitWithKnownLength(inputByteBuffer); - } else { - return splitWithUnknownLength(inputByteBuffer); - } - } - - private synchronized Iterable splitWithKnownLength(ByteBuffer inputByteBuffer) { List byteBuffers = new ArrayList<>(); // If current buffer is not empty, fill the buffer first. @@ -76,7 +72,7 @@ private synchronized Iterable splitWithKnownLength(ByteBuffer inputB fillCurrentBuffer(inputByteBuffer); if (isCurrentBufferFull()) { - addCurrentBufferToIterable(byteBuffers, chunkSize); + addCurrentBufferToIterable(byteBuffers); } } @@ -87,39 +83,11 @@ private synchronized Iterable splitWithKnownLength(ByteBuffer inputB // If this is the last chunk, add data buffered to the iterable if (isLastChunk()) { - int remainingBytesInBuffer = currentBuffer.position(); - addCurrentBufferToIterable(byteBuffers, remainingBytesInBuffer); + addCurrentBufferToIterable(byteBuffers); } return byteBuffers; } - private synchronized Iterable splitWithUnknownLength(ByteBuffer inputByteBuffer) { - boolean isLastChunk = inputByteBuffer.remaining() != chunkSize; - List bufferedList = new ArrayList<>(); - - while (inputByteBuffer.hasRemaining()) { - int bytesToCopy = Math.min(inputByteBuffer.remaining(), currentBuffer.remaining()); - byte[] bytes = new byte[bytesToCopy]; - inputByteBuffer.get(bytes); - currentBuffer.put(bytes); - - if (!currentBuffer.hasRemaining()) { - currentBuffer.flip(); - ByteBuffer bufferToSend = ByteBuffer.allocate(currentBuffer.limit()); - bufferToSend.put(currentBuffer); - bufferToSend.flip(); - bufferedList.add(bufferToSend); - currentBuffer.clear(); - } - } - - if (isLastChunk && currentBuffer.position() != 0) { - bufferedList.add((ByteBuffer) currentBuffer.flip()); - } - - return bufferedList; - } - private boolean isCurrentBufferFull() { return currentBuffer.position() == chunkSize; } @@ -143,19 +111,38 @@ private void splitRemainingInputByteBuffer(ByteBuffer inputByteBuffer, List getBufferedData() { + int remainingBytesInBuffer = currentBuffer.position(); + + if (remainingBytesInBuffer == 0) { + return Optional.empty(); + } + + ByteBuffer bufferedChunk = ByteBuffer.allocate(remainingBytesInBuffer); + currentBuffer.flip(); + bufferedChunk.put(currentBuffer); + bufferedChunk.flip(); + return Optional.of(bufferedChunk); + } + private boolean isLastChunk() { + if (totalBytes == null) { + return false; + } long remainingBytes = totalBytes - transferredBytes.get(); return remainingBytes != 0 && remainingBytes == currentBuffer.position(); } - private void addCurrentBufferToIterable(List byteBuffers, int capacity) { - ByteBuffer bufferedChunk = ByteBuffer.allocate(capacity); - currentBuffer.flip(); - bufferedChunk.put(currentBuffer); - bufferedChunk.flip(); - byteBuffers.add(bufferedChunk); - transferredBytes.addAndGet(bufferedChunk.remaining()); - currentBuffer.clear(); + private void addCurrentBufferToIterable(List byteBuffers) { + Optional bufferedChunk = getBufferedData(); + if (bufferedChunk.isPresent()) { + byteBuffers.add(bufferedChunk.get()); + transferredBytes.addAndGet(bufferedChunk.get().remaining()); + currentBuffer.clear(); + } } private void fillCurrentBuffer(ByteBuffer inputByteBuffer) { diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java index 0f520a38a04a..d031aaaee883 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java @@ -15,13 +15,19 @@ package software.amazon.awssdk.core.internal.async; +import static software.amazon.awssdk.core.internal.io.AwsChunkedInputStream.DEFAULT_CHUNK_SIZE; + import java.nio.ByteBuffer; +import java.util.Collections; import java.util.Optional; import org.reactivestreams.Subscriber; import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.core.async.SdkPublisher; import software.amazon.awssdk.core.internal.compression.Compressor; import software.amazon.awssdk.utils.Validate; +import software.amazon.awssdk.utils.async.DelegatingSubscriber; +import software.amazon.awssdk.utils.async.FlatteningSubscriber; import software.amazon.awssdk.utils.builder.SdkBuilder; /** @@ -31,7 +37,6 @@ @SdkInternalApi public class CompressionAsyncRequestBody implements AsyncRequestBody { - private static final int COMPRESSION_CHUNK_SIZE = 128 * 1024; private final AsyncRequestBody wrapped; private final Compressor compressor; private final int chunkSize; @@ -39,7 +44,34 @@ public class CompressionAsyncRequestBody implements AsyncRequestBody { private CompressionAsyncRequestBody(DefaultBuilder builder) { this.wrapped = Validate.paramNotNull(builder.asyncRequestBody, "asyncRequestBody"); this.compressor = Validate.paramNotNull(builder.compressor, "compressor"); - this.chunkSize = builder.chunkSize != null ? builder.chunkSize : COMPRESSION_CHUNK_SIZE; + this.chunkSize = builder.chunkSize != null ? builder.chunkSize : DEFAULT_CHUNK_SIZE; + } + + @Override + public void subscribe(Subscriber s) { + Validate.notNull(s, "Subscription MUST NOT be null."); + + SdkPublisher> split = split(wrapped); + SdkPublisher flattening = flattening(split); + flattening.map(compressor::compress).subscribe(s); + } + + @Override + public Optional contentLength() { + return wrapped.contentLength(); + } + + @Override + public String contentType() { + return wrapped.contentType(); + } + + private SdkPublisher> split(SdkPublisher source) { + return subscriber -> source.subscribe(new SplittingSubscriber(subscriber, chunkSize)); + } + + private SdkPublisher flattening(SdkPublisher> source) { + return subscriber -> source.subscribe(new FlatteningSubscriber<>(subscriber)); } /** @@ -103,25 +135,27 @@ public Builder chunkSize(Integer chunkSize) { } } - @Override - public Optional contentLength() { - return wrapped.contentLength(); - } + private static final class SplittingSubscriber extends DelegatingSubscriber> { + private final ChunkBuffer chunkBuffer; - @Override - public String contentType() { - return wrapped.contentType(); - } - - @Override - public void subscribe(Subscriber s) { - Validate.notNull(s, "Subscription MUST NOT be null."); + protected SplittingSubscriber(Subscriber> subscriber, int chunkSize) { + super(subscriber); + this.chunkBuffer = ChunkBuffer.builder() + .bufferSize(chunkSize) + .build(); + } - ChunkBuffer chunkBuffer = ChunkBuffer.builder() - .bufferSize(chunkSize) - .build(); + @Override + public void onNext(ByteBuffer byteBuffer) { + Iterable buffers = chunkBuffer.split(byteBuffer); + subscriber.onNext(buffers); + } - wrapped.flatMapIterable(chunkBuffer::split) - .map(compressor::compress).subscribe(s); + @Override + public void onComplete() { + Optional byteBuffer = chunkBuffer.getBufferedData(); + byteBuffer.ifPresent(buffer -> subscriber.onNext(Collections.singletonList(buffer))); + super.onComplete(); + } } } diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java index 18aaddcb17ba..f21a54df1675 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java @@ -23,6 +23,7 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Collection; +import java.util.Optional; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -194,7 +195,10 @@ void emptyAllocatedBytes_returnSameNumberOfEmptyBytes_unknownLength() { assertThat(BinaryUtils.copyBytesFrom(r)).isEqualTo(ByteBuffer.allocate(totalBytes % bufferSize).array()); } }); - assertThat(iteratedCounts.get()).isEqualTo(4); + assertThat(iteratedCounts.get()).isEqualTo(3); + + Optional lastBuffer = chunkBuffer.getBufferedData(); + assertThat(lastBuffer.isPresent()); } @@ -228,7 +232,7 @@ void emptyAllocatedBytes_returnSameNumberOfEmptyBytes_unknownLength() { * 111 is given as output since we consumed all the total bytes* */ @Test - void concurrentTreads_calling_bufferAndCreateChunks() throws ExecutionException, InterruptedException { + void concurrentTreads_calling_bufferAndCreateChunks_knownLength() throws ExecutionException, InterruptedException { int totalBytes = 17; int bufferSize = 5; int threads = 8; @@ -276,4 +280,48 @@ void concurrentTreads_calling_bufferAndCreateChunks() throws ExecutionException, assertThat(remainderBytesBuffers.get()).isOne(); assertThat(otherSizeBuffers.get()).isZero(); } + + @Test + void concurrentTreads_calling_bufferAndCreateChunks_unknownLength() throws ExecutionException, InterruptedException { + int totalBytes = 17; + int bufferSize = 5; + int threads = 8; + + ChunkBuffer chunkBuffer = ChunkBuffer.builder() + .bufferSize(bufferSize) + .build(); + + ExecutorService service = Executors.newFixedThreadPool(threads); + + Collection> futures; + + AtomicInteger counter = new AtomicInteger(0); + + futures = IntStream.range(0, threads).>mapToObj(t -> service.submit(() -> { + String inputString = StringUtils.repeat(Integer.toString(counter.incrementAndGet()), totalBytes); + return chunkBuffer.split(ByteBuffer.wrap(inputString.getBytes(StandardCharsets.UTF_8))); + })).collect(Collectors.toCollection(() -> new ArrayList<>(threads))); + + AtomicInteger filledBuffers = new AtomicInteger(0); + AtomicInteger otherSizeBuffers = new AtomicInteger(0); + + for (Future bufferedFuture : futures) { + Iterable buffers = bufferedFuture.get(); + buffers.forEach(b -> { + System.out.println(b.remaining()); + if (b.remaining() == bufferSize) { + filledBuffers.incrementAndGet(); + } else { + otherSizeBuffers.incrementAndGet(); + } + + }); + } + + assertThat(filledBuffers.get()).isEqualTo((totalBytes * threads) / bufferSize); + assertThat(otherSizeBuffers.get()).isZero(); + + ByteBuffer lastBuffer = chunkBuffer.getBufferedData().get(); + assertThat(lastBuffer.remaining()).isOne(); + } } diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/CompressionAsyncRequestBodyTckTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/CompressionAsyncRequestBodyTckTest.java index dd0cc312ad1f..54c74e1e97e9 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/CompressionAsyncRequestBodyTckTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/CompressionAsyncRequestBodyTckTest.java @@ -15,8 +15,16 @@ package software.amazon.awssdk.core.async; +import com.google.common.jimfs.Configuration; +import com.google.common.jimfs.Jimfs; import io.reactivex.Flowable; +import java.io.IOException; +import java.io.OutputStream; +import java.io.UncheckedIOException; import java.nio.ByteBuffer; +import java.nio.file.FileSystem; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Arrays; import java.util.Optional; import org.reactivestreams.Publisher; @@ -29,6 +37,8 @@ public class CompressionAsyncRequestBodyTckTest extends PublisherVerification { + private static final FileSystem fs = Jimfs.newFileSystem(Configuration.unix()); + private static final Path rootDir = fs.getRootDirectories().iterator().next(); private static final int MAX_ELEMENTS = 1000; private static final int CHUNK_SIZE = 128 * 1024; private static final Compressor compressor = new GzipCompressor(); @@ -45,7 +55,7 @@ public long maxElementsFromPublisher() { @Override public Publisher createPublisher(long n) { return CompressionAsyncRequestBody.builder() - .asyncRequestBody(customAsyncRequestBodyWithoutContentLength(createCompressibleStringOfNChunks(n).getBytes())) + .asyncRequestBody(customAsyncRequestBodyFromFileWithoutContentLength(n)) .compressor(compressor) .build(); } @@ -55,7 +65,34 @@ public Publisher createFailedPublisher() { return null; } - private static String createCompressibleStringOfNChunks(long nChunks) { + private static AsyncRequestBody customAsyncRequestBodyFromFileWithoutContentLength(long nChunks) { + return new AsyncRequestBody() { + @Override + public Optional contentLength() { + return Optional.empty(); + } + + @Override + public void subscribe(Subscriber s) { + Flowable.fromPublisher(AsyncRequestBody.fromFile(fileOfNChunks(nChunks))).subscribe(s); + } + }; + } + + private static Path fileOfNChunks(long nChunks) { + String name = String.format("%d-chunks-file.dat", nChunks); + Path p = rootDir.resolve(name); + if (!Files.exists(p)) { + try (OutputStream os = Files.newOutputStream(p)) { + os.write(createCompressibleArrayOfNChunks(nChunks)); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + return p; + } + + private static byte[] createCompressibleArrayOfNChunks(long nChunks) { int size = Math.toIntExact(nChunks * CHUNK_SIZE); ByteBuffer data = ByteBuffer.allocate(size); @@ -69,21 +106,6 @@ private static String createCompressibleStringOfNChunks(long nChunks) { data.put(a); data.put(b); - return new String(data.array()); - } - - private static AsyncRequestBody customAsyncRequestBodyWithoutContentLength(byte[] content) { - return new AsyncRequestBody() { - @Override - public Optional contentLength() { - return Optional.empty(); - } - - @Override - public void subscribe(Subscriber s) { - Flowable.fromPublisher(AsyncRequestBody.fromBytes(content)) - .subscribe(s); - } - }; + return data.array(); } } From 14f79fbad7c2b08e30c6c6e8ab547d17b482904a Mon Sep 17 00:00:00 2001 From: hdavidh Date: Fri, 25 Aug 2023 16:29:47 -0700 Subject: [PATCH 11/14] Remove unused field --- .../software/amazon/awssdk/core/internal/async/ChunkBuffer.java | 1 - 1 file changed, 1 deletion(-) diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java index f14f73b2466a..94b84c7a7c3b 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java @@ -39,7 +39,6 @@ public final class ChunkBuffer { private final ByteBuffer currentBuffer; private final int chunkSize; private Long totalBytes; - private boolean isCompleted; private ChunkBuffer(Long totalBytes, Integer bufferSize) { int chunkSize = bufferSize != null ? bufferSize : DEFAULT_ASYNC_CHUNK_SIZE; From 5212d3e963bb2563fa06716cead195978b4037ca Mon Sep 17 00:00:00 2001 From: hdavidh Date: Tue, 29 Aug 2023 09:33:30 -0700 Subject: [PATCH 12/14] Handle demand in Subscriber --- .../core/internal/async/ChunkBuffer.java | 4 +- .../async/CompressionAsyncRequestBody.java | 57 ++++++++++++++++++- ...stCompressionStreamingIntegrationTest.java | 2 +- 3 files changed, 57 insertions(+), 6 deletions(-) diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java index 94b84c7a7c3b..1981c432a5a6 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java @@ -29,8 +29,8 @@ /** * Class that will buffer incoming BufferBytes to chunks of bufferSize. - * If totalBytes is not provided, i.e. content-length is unknown, {@code getBufferedData()} should be used in the Subscriber's - * {@code onComplete} to check for a final chunk that is smaller than the chunk size, and send if present. + * If totalBytes is not provided, i.e. content-length is unknown, {@link #getBufferedData()} should be used in the Subscriber's + * {@code onComplete()} to check for a final chunk that is smaller than the chunk size, and send if present. */ @SdkInternalApi public final class ChunkBuffer { diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java index d031aaaee883..82da601f0acc 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/CompressionAsyncRequestBody.java @@ -20,7 +20,10 @@ import java.nio.ByteBuffer; import java.util.Collections; import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.async.SdkPublisher; @@ -137,6 +140,10 @@ public Builder chunkSize(Integer chunkSize) { private static final class SplittingSubscriber extends DelegatingSubscriber> { private final ChunkBuffer chunkBuffer; + private final AtomicBoolean upstreamDone = new AtomicBoolean(false); + private final AtomicLong downstreamDemand = new AtomicLong(); + private final Object lock = new Object(); + private volatile boolean sentFinalChunk = false; protected SplittingSubscriber(Subscriber> subscriber, int chunkSize) { super(subscriber); @@ -145,17 +152,61 @@ protected SplittingSubscriber(Subscriber> subscribe .build(); } + @Override + public void onSubscribe(Subscription s) { + subscriber.onSubscribe(new Subscription() { + @Override + public void request(long n) { + if (n <= 0) { + throw new IllegalArgumentException("n > 0 required but it was " + n); + } + + downstreamDemand.getAndAdd(n); + + if (upstreamDone.get()) { + sendFinalChunk(); + } else { + s.request(n); + } + } + + @Override + public void cancel() { + s.cancel(); + } + }); + } + @Override public void onNext(ByteBuffer byteBuffer) { + downstreamDemand.decrementAndGet(); Iterable buffers = chunkBuffer.split(byteBuffer); subscriber.onNext(buffers); } @Override public void onComplete() { - Optional byteBuffer = chunkBuffer.getBufferedData(); - byteBuffer.ifPresent(buffer -> subscriber.onNext(Collections.singletonList(buffer))); - super.onComplete(); + upstreamDone.compareAndSet(false, true); + if (downstreamDemand.get() > 0) { + sendFinalChunk(); + } + } + + @Override + public void onError(Throwable t) { + upstreamDone.compareAndSet(false, true); + super.onError(t); + } + + private void sendFinalChunk() { + synchronized (lock) { + if (!sentFinalChunk) { + sentFinalChunk = true; + Optional byteBuffer = chunkBuffer.getBufferedData(); + byteBuffer.ifPresent(buffer -> subscriber.onNext(Collections.singletonList(buffer))); + subscriber.onComplete(); + } + } } } } diff --git a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java index 41c5b2fc849d..bb4a2a9bf0c3 100644 --- a/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java +++ b/services/mediastoredata/src/it/java/software/amazon/awssdk/services/mediastoredata/RequestCompressionStreamingIntegrationTest.java @@ -114,7 +114,7 @@ public static void tearDown() throws InterruptedException { Waiter.run(() -> syncClient.describeObject(r -> r.path("/foo"))) .untilException(ObjectNotFoundException.class) .orFailAfter(Duration.ofMinutes(1)); - Thread.sleep(500); + Thread.sleep(1000); mediaStoreClient.deleteContainer(r -> r.containerName(CONTAINER_NAME)); } From 563b7810b837371ef90fe4cd81e71fb49257facf Mon Sep 17 00:00:00 2001 From: hdavidh Date: Tue, 29 Aug 2023 15:36:46 -0700 Subject: [PATCH 13/14] Address comments --- .../core/internal/async/ChunkBuffer.java | 6 +-- .../awssdk/core/async/ChunkBufferTest.java | 45 +------------------ 2 files changed, 3 insertions(+), 48 deletions(-) diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java index 1981c432a5a6..9599bb75655f 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java @@ -35,7 +35,7 @@ @SdkInternalApi public final class ChunkBuffer { private static final Logger log = Logger.loggerFor(ChunkBuffer.class); - private AtomicLong transferredBytes; + private final AtomicLong transferredBytes; private final ByteBuffer currentBuffer; private final int chunkSize; private Long totalBytes; @@ -45,9 +45,7 @@ private ChunkBuffer(Long totalBytes, Integer bufferSize) { this.chunkSize = chunkSize; this.currentBuffer = ByteBuffer.allocate(chunkSize); this.transferredBytes = new AtomicLong(0); - if (totalBytes != null) { - this.totalBytes = totalBytes; - } + this.totalBytes = totalBytes; } public static Builder builder() { diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java index f21a54df1675..41250225664a 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/ChunkBufferTest.java @@ -199,6 +199,7 @@ void emptyAllocatedBytes_returnSameNumberOfEmptyBytes_unknownLength() { Optional lastBuffer = chunkBuffer.getBufferedData(); assertThat(lastBuffer.isPresent()); + assertThat(lastBuffer.get().remaining()).isEqualTo(2); } @@ -280,48 +281,4 @@ void concurrentTreads_calling_bufferAndCreateChunks_knownLength() throws Executi assertThat(remainderBytesBuffers.get()).isOne(); assertThat(otherSizeBuffers.get()).isZero(); } - - @Test - void concurrentTreads_calling_bufferAndCreateChunks_unknownLength() throws ExecutionException, InterruptedException { - int totalBytes = 17; - int bufferSize = 5; - int threads = 8; - - ChunkBuffer chunkBuffer = ChunkBuffer.builder() - .bufferSize(bufferSize) - .build(); - - ExecutorService service = Executors.newFixedThreadPool(threads); - - Collection> futures; - - AtomicInteger counter = new AtomicInteger(0); - - futures = IntStream.range(0, threads).>mapToObj(t -> service.submit(() -> { - String inputString = StringUtils.repeat(Integer.toString(counter.incrementAndGet()), totalBytes); - return chunkBuffer.split(ByteBuffer.wrap(inputString.getBytes(StandardCharsets.UTF_8))); - })).collect(Collectors.toCollection(() -> new ArrayList<>(threads))); - - AtomicInteger filledBuffers = new AtomicInteger(0); - AtomicInteger otherSizeBuffers = new AtomicInteger(0); - - for (Future bufferedFuture : futures) { - Iterable buffers = bufferedFuture.get(); - buffers.forEach(b -> { - System.out.println(b.remaining()); - if (b.remaining() == bufferSize) { - filledBuffers.incrementAndGet(); - } else { - otherSizeBuffers.incrementAndGet(); - } - - }); - } - - assertThat(filledBuffers.get()).isEqualTo((totalBytes * threads) / bufferSize); - assertThat(otherSizeBuffers.get()).isZero(); - - ByteBuffer lastBuffer = chunkBuffer.getBufferedData().get(); - assertThat(lastBuffer.remaining()).isOne(); - } } From 2d1b6feefab654612b6a55f1f5d9b6a8512cd457 Mon Sep 17 00:00:00 2001 From: hdavidh Date: Tue, 29 Aug 2023 16:10:32 -0700 Subject: [PATCH 14/14] Add back final modifier --- .../software/amazon/awssdk/core/internal/async/ChunkBuffer.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java index 9599bb75655f..bdf84d549b8f 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/ChunkBuffer.java @@ -38,7 +38,7 @@ public final class ChunkBuffer { private final AtomicLong transferredBytes; private final ByteBuffer currentBuffer; private final int chunkSize; - private Long totalBytes; + private final Long totalBytes; private ChunkBuffer(Long totalBytes, Integer bufferSize) { int chunkSize = bufferSize != null ? bufferSize : DEFAULT_ASYNC_CHUNK_SIZE;