Combine chunking + encoding

Combine these two steps to improve performance by reducing memory
allocations.

Author: dagnir

core/http-auth-aws/src/main/java/software/amazon/awssdk/http/auth/aws/internal/signer/chunkedencoding/ChunkedEncodedPublisher.java

@@ -61,6 +61,7 @@
  */
 @SdkInternalApi
 public class ChunkedEncodedPublisher implements Publisher<ByteBuffer> {
+    private final ByteBuffer EMPTY = ByteBuffer.allocate(0);
     private static final byte[] CRLF = {'\r', '\n'};
     private static final byte SEMICOLON = ';';
     private static final byte EQUALS = '=';
@@ -83,6 +84,7 @@ public ChunkedEncodedPublisher(Builder b) {
         this.extensions.addAll(b.extensions);
         this.trailers.addAll(b.trailers);
         this.addEmptyTrailingChunk = b.addEmptyTrailingChunk;
+        this.chunkBuffer = ByteBuffer.allocate(chunkSize);
     }
 
     @Override
@@ -93,9 +95,8 @@ public void subscribe(Subscriber<? super ByteBuffer> subscriber) {
         Publisher<Iterable<ByteBuffer>> chunked = chunk(lengthEnforced);
         Publisher<Iterable<ByteBuffer>> trailingAdded = addTrailingChunks(chunked);
         Publisher<ByteBuffer> flattened = flatten(trailingAdded);
-        Publisher<ByteBuffer> encoded = map(flattened, this::encodeChunk);
 
-        encoded.subscribe(subscriber);
+        flattened.subscribe(subscriber);
     }
 
     public static Builder builder() {
@@ -105,7 +106,7 @@ public static Builder builder() {
     private void resetState() {
         extensions.forEach(Resettable::reset);
         trailers.forEach(Resettable::reset);
-        chunkBuffer = null;
+        chunkBuffer = ByteBuffer.allocate(chunkSize);
     }
 
     private Iterable<Iterable<ByteBuffer>> getTrailingChunks() {
@@ -114,12 +115,13 @@ private Iterable<Iterable<ByteBuffer>> getTrailingChunks() {
         if (chunkBuffer != null) {
             chunkBuffer.flip();
             if (chunkBuffer.hasRemaining()) {
-                trailing.add(chunkBuffer);
+                trailing.add(encodeChunk(chunkBuffer));
+                chunkBuffer = null;
             }
         }
 
         if (addEmptyTrailingChunk) {
-            trailing.add(ByteBuffer.allocate(0));
+            trailing.add(encodeChunk(EMPTY.duplicate()));
         }
 
         return Collections.singletonList(trailing);
@@ -174,6 +176,7 @@ private ByteBuffer encodeChunk(ByteBuffer byteBuffer) {
                                     .mapToInt(t -> t.remaining() + CRLF.length)
                                     .sum();
 
+
         int encodedLen = chunkSizeHex.length + extensionsLength + CRLF.length + contentLen + trailerLen + CRLF.length;
 
         if (isTrailerChunk) {
@@ -272,37 +275,57 @@ protected ChunkingSubscriber(Subscriber<? super Iterable<ByteBuffer>> subscriber
         }
 
         @Override
-        public void onNext(ByteBuffer byteBuffer) {
-            if (chunkBuffer == null) {
-                chunkBuffer = ByteBuffer.allocate(chunkSize);
-            }
-
-            long totalBufferedBytes = (long) chunkBuffer.position() + byteBuffer.remaining();
+        public void onNext(ByteBuffer inputBuffer) {
+            long totalBufferedBytes = (long) chunkBuffer.position() + inputBuffer.remaining();
+            // compute the number full chunks we have currently
             int nBufferedChunks = (int) (totalBufferedBytes / chunkSize);
 
             List<ByteBuffer> chunks = new ArrayList<>(nBufferedChunks);
 
             if (nBufferedChunks > 0) {
-                for (int i = 0; i < nBufferedChunks; i++) {
-                    ByteBuffer slice = byteBuffer.slice();
-                    int maxBytesToCopy = Math.min(chunkBuffer.remaining(), slice.remaining());
-                    slice.limit(maxBytesToCopy);
+                // We have some data from the previous inputBuffer
+                if (chunkBuffer.position() > 0) {
+                    int bytesToFill = chunkBuffer.remaining();
 
+                    ByteBuffer slice = inputBuffer.slice();
+
+                    slice.limit(slice.position() + bytesToFill);
+                    inputBuffer.position(inputBuffer.position() + bytesToFill);
+
+                    // At this point, we know chunkBuffer is full since inputBuffer has at least enough bytes to make up a full
+                    // chunk along with the data already in chunkBuffer
                     chunkBuffer.put(slice);
-                    if (!chunkBuffer.hasRemaining()) {
-                        chunkBuffer.flip();
-                        chunks.add(chunkBuffer);
-                        chunkBuffer = ByteBuffer.allocate(chunkSize);
-                    }
+                    chunkBuffer.flip();
+                    chunks.add(encodeChunk(chunkBuffer));
+
+                    chunkBuffer.flip();
 
-                    byteBuffer.position(byteBuffer.position() + maxBytesToCopy);
+                    nBufferedChunks--;
+                }
+
+                // Now encode all the remaining full chunks from inputBuffer.
+                // At this point chunkBuffer has no data in it; slice off chunks from inputBuffer and encode directly
+                for (int i = 0; i < nBufferedChunks; i++) {
+                    ByteBuffer slice = inputBuffer.slice();
+
+                    int sliceLimit = Math.min(slice.limit(), chunkSize);
+                    slice.limit(sliceLimit);
+
+                    inputBuffer.position(inputBuffer.position() + slice.remaining());
+
+                    if (slice.remaining() >= chunkSize) {
+                        slice.limit(slice.position() + chunkSize);
+                        chunks.add(encodeChunk(slice));
+                    } else {
+                        chunkBuffer.put(slice);
+                    }
                 }
 
-                if (byteBuffer.hasRemaining()) {
-                    chunkBuffer.put(byteBuffer);
+                if (inputBuffer.hasRemaining()) {
+                    chunkBuffer.put(inputBuffer);
                 }
             } else {
-                chunkBuffer.put(byteBuffer);
+                chunkBuffer.put(inputBuffer);
             }
 
             subscriber.onNext(chunks);

core/http-auth-aws/src/test/java/software/amazon/awssdk/http/auth/aws/internal/signer/chunkedencoding/ChunkedEncodedPublisherTest.java

@@ -264,6 +264,7 @@ void subscribe_randomElementSizes_chunksHaveExtensions_dataChunkedCorrectly() {
                                                                               .build();
 
             List<ByteBuffer> chunks = getAllElements(chunkedPublisher);
+            assertThat(chunks.size()).isEqualTo(24);
 
             chunks.forEach(c -> {
                 String header = StandardCharsets.UTF_8.decode(getHeader(c)).toString();

services/s3/pom.xml

@@ -235,5 +235,9 @@
             <artifactId>rxjava</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-slf4j-impl</artifactId>
+        </dependency>
     </dependencies>
 </project>

services/s3/src/it/resources/log4j2.properties

@@ -20,7 +20,7 @@ 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.level = debug
 rootLogger.appenderRef.stdout.ref = ConsoleAppender
 
 # Uncomment below to enable more specific logging
@@ -34,5 +34,5 @@ rootLogger.appenderRef.stdout.ref = ConsoleAppender
 #logger.apache.name = org.apache.http.wire
 #logger.apache.level = debug
 #
-#logger.netty.name = io.netty.handler.logging
-#logger.netty.level = debug
+logger.netty.name = io.netty.handler.logging
+logger.netty.level = debug

test/sdk-benchmarks/pom.xml

@@ -82,6 +82,12 @@
             <version>${awsjavasdk.version}</version>
         </dependency>
 
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>s3</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
+
         <dependency>
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-java-sdk-ec2</artifactId>
@@ -94,6 +100,12 @@
             <version>${awsjavasdk.version}</version>
         </dependency>
 
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>s3</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
+
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>aws-query-protocol</artifactId>