Added flush mechanism

gregw · Sep 26, 2023 · Sep 21, 2023 · Sep 21, 2023 · Sep 21, 2023 · Sep 22, 2023
commit 676799cd9f4c2cbb83836b81fcabb4822cb7ec18
diff --git a/jetty-core/jetty-io/src/main/java/org/eclipse/jetty/io/ByteBufferAggregator.java b/jetty-core/jetty-io/src/main/java/org/eclipse/jetty/io/ByteBufferAggregator.java
@@ -58,6 +58,15 @@ public ByteBufferAggregator(ByteBufferPool bufferPool, boolean direct, int start
         _currentSize = startSize;
     }
 
+    /**
+     * Get the currently aggregated length.
+     * @return The current total aggregated bytes.
+     */
+    public int length()
+    {
+        return _aggregatedSize;
+    }
+
     /**
      * Aggregates the given ByteBuffer. This copies bytes up to the specified maximum size, at which
      * time this method returns {@code true} and {@link #takeRetainableByteBuffer()} must be called

diff --git a/jetty-core/jetty-io/src/main/java/org/eclipse/jetty/io/content/BufferedContentSink.java b/jetty-core/jetty-io/src/main/java/org/eclipse/jetty/io/content/BufferedContentSink.java
@@ -35,6 +35,19 @@
  */
 public class BufferedContentSink implements Content.Sink
 {
+    /**
+     * An empty byte array, which if {@link #write(boolean, ByteBuffer, Callback) written}
+     * as a {@link ByteBuffer#wrap(byte[]) wrapped ByteBuffer}
+     * will invoke a {@link #flush(Callback)} operation.
+     */
+    public static final byte[] FLUSH_BYTES = new byte[0];
+
+    /**
+     * An empty {@link ByteBuffer}, which if {@link #write(boolean, ByteBuffer, Callback) written}
+     * will invoke a {@link #flush(Callback)} operation.
+     */
+    public static final ByteBuffer FLUSH_BUFFER = ByteBuffer.wrap(FLUSH_BYTES);
+
     private static final Logger LOG = LoggerFactory.getLogger(BufferedContentSink.class);
 
     private static final int START_BUFFER_SIZE = 1024;
@@ -103,6 +116,15 @@ public void write(boolean last, ByteBuffer byteBuffer, Callback callback)
         }
     }
 
+    /**
+     * Flush the buffered content.
+     * @param callback Callback completed when the flush is complete
+     */
+    public void flush(Callback callback)
+    {
+        flush(false, FLUSH_BUFFER, callback);
+    }
+
     /**
      * Flushes the aggregated buffer if something was aggregated, then flushes the
      * given buffer, bypassing the aggregator.
@@ -119,7 +141,7 @@ private void flush(boolean last, ByteBuffer currentBuffer, Callback callback)
                 LOG.debug("nothing aggregated, flushing current buffer {}", currentBuffer);
             _flusher.offer(last, currentBuffer, callback);
         }
-        else
+        else if (BufferUtil.hasContent(currentBuffer))
         {
             if (LOG.isDebugEnabled())
                 LOG.debug("flushing aggregated buffer {}", aggregatedBuffer);
@@ -144,6 +166,10 @@ public void failed(Throwable x)
                 }
             });
         }
+        else
+        {
+            _flusher.offer(false, aggregatedBuffer.getByteBuffer(), Callback.from(aggregatedBuffer::release, callback));
+        }
     }
 
     /**
@@ -152,7 +178,11 @@ public void failed(Throwable x)
     private void aggregateAndFlush(boolean last, ByteBuffer currentBuffer, Callback callback)
     {
         boolean full = _aggregator.aggregate(currentBuffer);
-        boolean complete = last && !currentBuffer.hasRemaining();
+        boolean empty = !currentBuffer.hasRemaining();
+        boolean flush = full ||
+            currentBuffer == FLUSH_BUFFER ||
+            empty && currentBuffer.hasArray() && currentBuffer.array() == FLUSH_BYTES;
+        boolean complete = last && empty;
         if (LOG.isDebugEnabled())
             LOG.debug("aggregated current buffer, full={}, complete={}, bytes left={}, aggregator={}", full, complete, currentBuffer.remaining(), _aggregator);
         if (complete)
@@ -171,34 +201,42 @@ private void aggregateAndFlush(boolean last, ByteBuffer currentBuffer, Callback
                 _flusher.offer(true, BufferUtil.EMPTY_BUFFER, callback);
             }
         }
-        else if (full)
+        else if (flush)
         {
             RetainableByteBuffer aggregatedBuffer = _aggregator.takeRetainableByteBuffer();
             if (LOG.isDebugEnabled())
-                LOG.debug("writing aggregated buffer: {} bytes", aggregatedBuffer.remaining());
-            _flusher.offer(false, aggregatedBuffer.getByteBuffer(), new Callback.Nested(Callback.from(aggregatedBuffer::release))
+                LOG.debug("writing aggregated buffer: {} bytes, then {}", aggregatedBuffer.remaining(), currentBuffer.remaining());
+
+            if (BufferUtil.hasContent(currentBuffer))
             {
-                @Override
-                public void succeeded()
+                _flusher.offer(false, aggregatedBuffer.getByteBuffer(), new Callback.Nested(Callback.from(aggregatedBuffer::release))
                 {
-                    super.succeeded();
-                    if (LOG.isDebugEnabled())
-                        LOG.debug("written aggregated buffer, writing remaining of current: {} bytes{}", currentBuffer.remaining(), (last ? " (last write)" : ""));
-                    if (last)
-                        _flusher.offer(true, currentBuffer, callback);
-                    else
-                        aggregateAndFlush(false, currentBuffer, callback);
-                }
+                    @Override
+                    public void succeeded()
+                    {
+                        super.succeeded();
+                        if (LOG.isDebugEnabled())
+                            LOG.debug("written aggregated buffer, writing remaining of current: {} bytes{}", currentBuffer.remaining(), (last ? " (last write)" : ""));
+                        if (last)
+                            _flusher.offer(true, currentBuffer, callback);
+                        else
+                            aggregateAndFlush(false, currentBuffer, callback);
+                    }
 
-                @Override
-                public void failed(Throwable x)
-                {
-                    if (LOG.isDebugEnabled())
-                        LOG.debug("failure writing aggregated buffer", x);
-                    super.failed(x);
-                    callback.failed(x);
-                }
-            });
+                    @Override
+                    public void failed(Throwable x)
+                    {
+                        if (LOG.isDebugEnabled())
+                            LOG.debug("failure writing aggregated buffer", x);
+                        super.failed(x);
+                        callback.failed(x);
+                    }
+                });
+            }
+            else
+            {
+                _flusher.offer(false, aggregatedBuffer.getByteBuffer(), Callback.from(aggregatedBuffer::release, callback));
+            }
         }
         else
         {

diff --git a/jetty-core/jetty-io/src/main/java/org/eclipse/jetty/io/content/ContentSinkOutputStream.java b/jetty-core/jetty-io/src/main/java/org/eclipse/jetty/io/content/ContentSinkOutputStream.java
@@ -64,7 +64,7 @@ public void flush() throws IOException
     {
         try (Blocker.Callback callback = _blocking.callback())
         {
-            sink.write(false, null, callback);
+            sink.write(false, BufferedContentSink.FLUSH_BUFFER, callback);
             callback.block();
         }
         catch (Throwable x)
@@ -78,12 +78,17 @@ public void close() throws IOException
     {
         try (Blocker.Callback callback = _blocking.callback())
         {
-            sink.write(true, null, callback);
+            close(callback);
             callback.block();
         }
         catch (Throwable x)
         {
             throw IO.rethrow(x);
         }
     }
+
+    public void close(Callback callback) throws IOException
+    {
+        sink.write(true, null, callback);
+    }
 }
diff --git a/jetty-core/jetty-io/src/test/java/org/eclipse/jetty/io/BufferedContentSinkTest.java b/jetty-core/jetty-io/src/test/java/org/eclipse/jetty/io/BufferedContentSinkTest.java
@@ -21,6 +21,8 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
+import java.util.stream.Stream;
 
 import org.eclipse.jetty.io.content.AsyncContent;
 import org.eclipse.jetty.io.content.BufferedContentSink;
@@ -29,6 +31,8 @@
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.awaitility.Awaitility.await;
@@ -39,6 +43,7 @@
 import static org.hamcrest.Matchers.nullValue;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -229,6 +234,46 @@ public void testSmallBuffer()
         }
     }
 
+    public static Stream<BiConsumer<BufferedContentSink, Callback>> flushers()
+    {
+        return Stream.of(
+            BufferedContentSink::flush,
+            (b, callback) -> b.write(false, BufferedContentSink.FLUSH_BUFFER, callback),
+            (b, callback) -> b.write(false, ByteBuffer.wrap(BufferedContentSink.FLUSH_BYTES), callback)
+        );
+    }
+
+    @ParameterizedTest
+    @MethodSource("flushers")
+    public void testFlush(BiConsumer<BufferedContentSink, Callback> flusher) throws Exception
+    {
+        ByteBuffer accumulatingBuffer = BufferUtil.allocate(4096);
+        BufferUtil.flipToFill(accumulatingBuffer);
+
+        try (AsyncContent async = new AsyncContent(); )
+        {
+            BufferedContentSink buffered = new BufferedContentSink(async, _bufferPool, false, 8192, 8192);
+
+            Callback.Completable callback = new Callback.Completable();
+            buffered.write(false, BufferUtil.toBuffer("Hello "), callback);
+            callback.get(5, TimeUnit.SECONDS);
+            assertNull(async.read());
+
+            callback = new Callback.Completable();
+            buffered.write(false, BufferUtil.toBuffer("World!"), callback);
+            callback.get(5, TimeUnit.SECONDS);
+            assertNull(async.read());
+
+            callback = new Callback.Completable();
+            flusher.accept(buffered, callback);
+            Content.Chunk chunk = async.read();
+            assertThat(chunk.isLast(), is(false));
+            assertThat(BufferUtil.toString(chunk.getByteBuffer()), is("Hello World!"));
+            chunk.release();
+            callback.get(5, TimeUnit.SECONDS);
+        }
+    }
+
     @Test
     public void testMaxAggregationSizeExceeded()
     {