Merge pull request #1364 from carterkozak/ckozak/UNDERTOW-2142-batch-chunk-proposal

[UNDERTOW-2638] Proposal for a multi-buffer ChunkedStreamSinkConduit.write implementation

Author: fl4via

core/src/main/java/io/undertow/conduits/ChunkedStreamSinkConduit.java

@@ -35,6 +35,7 @@
 import io.undertow.util.HeaderValues;
 import io.undertow.util.Headers;
 import io.undertow.util.ImmediatePooledByteBuffer;
+import org.xnio.Buffers;
 import org.xnio.IoUtils;
 import io.undertow.connector.ByteBufferPool;
 import io.undertow.connector.PooledByteBuffer;
@@ -128,6 +129,109 @@ public int write(final ByteBuffer src) throws IOException {
         return doWrite(src);
     }
 
+    long doWrite(final ByteBuffer[] srcs, int offset, int length) throws IOException {
+        if (anyAreSet(state, FLAG_WRITES_SHUTDOWN)) {
+            throw new ClosedChannelException();
+        }
+        // Write as many buffers as possible without a chunk-size overflowing an integer.
+        long totalRemaining = 0;
+        for (int i = 0; i < length; i++) {
+            ByteBuffer buf = srcs[i + offset];
+            int remaining = buf.remaining();
+            if (totalRemaining + remaining > Integer.MAX_VALUE) {
+                // Avoid producing chunks too large for clients by reducing the number of buffers
+                // until total remaining fits within a 32-bit signed integer value. This is safe
+                // because a single java ByteBuffer has a capacity represented by an integer.
+                length = i;
+                break;
+            }
+            totalRemaining += remaining;
+        }
+        if(totalRemaining == 0) {
+            return 0;
+        }
+        int remaining = (int) totalRemaining;
+        this.state |= FLAG_FIRST_DATA_WRITTEN;
+        int oldLimit = srcs[length - 1].limit();
+        boolean dataRemaining = false; //set to true if there is data in src that still needs to be written out
+        if (chunkleft == 0 && !chunkingSepBuffer.hasRemaining()) {
+            chunkingBuffer.clear();
+            putIntAsHexString(chunkingBuffer, remaining);
+            chunkingBuffer.put(CRLF);
+            chunkingBuffer.flip();
+            chunkingSepBuffer.clear();
+            chunkingSepBuffer.put(CRLF);
+            chunkingSepBuffer.flip();
+            state |= FLAG_WRITTEN_FIRST_CHUNK;
+            chunkleft = remaining;
+        } else {
+            int maxRemaining = chunkleft;
+            for (int i = 0; i < length; i++) {
+                ByteBuffer buf = srcs[offset + i];
+                int bufRemaining = buf.remaining();
+                if (bufRemaining >= maxRemaining) {
+                    length = i + 1;
+                    oldLimit = buf.limit();
+                    dataRemaining = true;
+                    buf.limit(buf.position() + maxRemaining);
+                    break;
+                }
+                maxRemaining -= bufRemaining;
+            }
+        }
+        try {
+            int chunkingSize = chunkingBuffer.remaining();
+            int chunkingSepSize = chunkingSepBuffer.remaining();
+            if (chunkingSize > 0 || chunkingSepSize > 0 || lastChunkBuffer != null) {
+                int originalRemaining = (int) Buffers.remaining(srcs, offset, length);
+                long result;
+                if (lastChunkBuffer == null || dataRemaining) {
+                    // chunkingBuffer
+                    // srcs (taking into account offset+length)
+                    // chunkingSepBuffer
+                    final ByteBuffer[] buf = new ByteBuffer[2 + length];
+                    buf[0] = chunkingBuffer;
+                    System.arraycopy(srcs, offset , buf, 1, length);
+                    buf[length + 1] = chunkingSepBuffer;
+                    result = next.write(buf, 0, buf.length);
+                } else {
+                    // chunkingBuffer
+                    // srcs (taking into account offset+length)
+                    // lastChunkBuffer
+                    final ByteBuffer[] buf = new ByteBuffer[2 + length];
+                    buf[0] = chunkingBuffer;
+                    System.arraycopy(srcs, offset , buf, 1, length);
+                    buf[length + 1] = lastChunkBuffer.getBuffer();
+                    if (anyAreSet(state, CONF_FLAG_PASS_CLOSE)) {
+                        result = next.writeFinal(buf, 0, buf.length);
+                    } else {
+                        result = next.write(buf, 0, buf.length);
+                    }
+                    if (Buffers.remaining(srcs, offset, length) == 0) {
+                        state |= FLAG_WRITES_SHUTDOWN;
+                    }
+                    if (!lastChunkBuffer.getBuffer().hasRemaining()) {
+                        state |= FLAG_NEXT_SHUTDOWN;
+                        lastChunkBuffer.close();
+                    }
+                }
+                int srcWritten = originalRemaining - (int) Buffers.remaining(srcs, offset, length);
+                chunkleft -= srcWritten;
+                if (result < chunkingSize) {
+                    return 0;
+                } else {
+                    return srcWritten;
+                }
+            } else {
+                long result = next.write(srcs, offset, length);
+                chunkleft -= result;
+                return result;
+
+            }
+        } finally {
+            srcs[length - 1].limit(oldLimit);
+        }
+    }
 
     int doWrite(final ByteBuffer src) throws IOException {
         if (anyAreSet(state, FLAG_WRITES_SHUTDOWN)) {
@@ -195,7 +299,6 @@ int doWrite(final ByteBuffer src) throws IOException {
         } finally {
             src.limit(oldLimit);
         }
-
     }
 
     @Override
@@ -217,13 +320,7 @@ public void truncateWrites() throws IOException {
 
     @Override
     public long write(final ByteBuffer[] srcs, final int offset, final int length) throws IOException {
-        for (int i = 0; i < length; i++) {
-            ByteBuffer srcBuffer = srcs[offset + i];
-            if (srcBuffer.hasRemaining()) {
-                return write(srcBuffer);
-            }
-        }
-        return 0;
+        return doWrite(srcs, offset, length);
     }
 
     @Override
@@ -382,7 +479,7 @@ private void createLastChunk(final boolean writeFinal) throws UnsupportedEncodin
             lastChunkBuffer.put(CRLF);
         }
         //horrible hack
-        //there is a situation where we can get a buffer leak here if the connection is terminated abnormaly
+        //there is a situation where we can get a buffer leak here if the connection is terminated abnormally
         //this should be fixed once this channel has its lifecycle tied to the connection, same as fixed length
         lastChunkBuffer.flip();
         ByteBuffer data = ByteBuffer.allocate(lastChunkBuffer.remaining());