Issue #6566 - use the demand API for the websocket MessageSinks

jetty-io/src/main/java/org/eclipse/jetty/io/BufferCallbackAccumulator.java

@@ -0,0 +1,99 @@
+//
+// ========================================================================
+// Copyright (c) 1995-2021 Mort Bay Consulting Pty Ltd and others.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// https://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
+// which is available at https://www.apache.org/licenses/LICENSE-2.0.
+//
+// SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
+// ========================================================================
+//
+
+package org.eclipse.jetty.io;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+import org.eclipse.jetty.util.BufferUtil;
+import org.eclipse.jetty.util.Callback;
+
+public class BufferCallbackAccumulator
+{
+    private final List<Entry> _entries = new ArrayList<>();
+    private int _length;
+
+    private static class Entry
+    {
+        private final ByteBuffer buffer;
+        private final Callback callback;
+
+        Entry(ByteBuffer buffer, Callback callback)
+        {
+            this.buffer = buffer;
+            this.callback = callback;
+        }
+    }
+
+    public void addEntry(ByteBuffer buffer, Callback callback)
+    {
+        _entries.add(new Entry(buffer, callback));
+        _length = Math.addExact(_length, buffer.remaining());
+    }
+
+    /**
+     * Get the amount of bytes which have been accumulated.
+     * This will add up the remaining of each buffer in the accumulator.
+     * @return the total length of the content in the accumulator.
+     */
+    public int getLength()
+    {
+        return _length;
+    }
+
+    /**
+     * @return a newly allocated byte array containing all content written into the accumulator.
+     */
+    public byte[] takeByteArray()
+    {
+        int length = getLength();
+        if (length == 0)
+            return new byte[0];
+
+        byte[] bytes = new byte[length];
+        ByteBuffer buffer = BufferUtil.toBuffer(bytes);
+        BufferUtil.clearToFill(buffer);
+        writeTo(buffer);
+        return bytes;
+    }
+
+    public void writeTo(ByteBuffer buffer)
+    {
+        for (Iterator<Entry> iterator = _entries.iterator(); iterator.hasNext();)
+        {
+            Entry entry = iterator.next();
+            _length -= entry.buffer.remaining();
+            buffer.put(entry.buffer);
+            iterator.remove();
+            entry.callback.succeeded();
+        }
+
+        if (!_entries.isEmpty())
+            throw new IllegalStateException("remaining entries: " + _entries.size());
+        if (_length != 0)
+            throw new IllegalStateException("non-zero length: " + _length);
+    }
+
+    public void fail(Throwable t)
+    {
+        for (Entry entry : _entries)
+        {
+            entry.callback.failed(t);
+        }
+        _entries.clear();
+        _length = 0;
+    }
+}

jetty-util/src/main/java/org/eclipse/jetty/util/Callback.java

@@ -133,7 +133,7 @@ public void failed(Throwable x)
     }
 
     /**
-     * Creaste a callback that runs completed when it succeeds or fails
+     * Create a callback that runs completed when it succeeds or fails
      *
      * @param completed The completion to run on success or failure
      * @return a new callback
@@ -396,6 +396,32 @@ public InvocationType getInvocationType()
      */
     class Completable extends CompletableFuture<Void> implements Callback
     {
+        /**
+         * Create a completable future given a callback.
+         *
+         * @param callback The nested callback.
+         * @return a new Completable which will succeed this callback when completed.
+         */
+        public static Completable from(Callback callback)
+        {
+            return new Completable()
+            {
+                @Override
+                public void succeeded()
+                {
+                    callback.succeeded();
+                    super.succeeded();
+                }
+
+                @Override
+                public void failed(Throwable x)
+                {
+                    callback.failed(x);
+                    super.failed(x);
+                }
+            };
+        }
+
         private final InvocationType invocation;
 
         public Completable()

jetty-websocket/websocket-core-common/src/main/java/org/eclipse/jetty/websocket/core/CoreSession.java

@@ -110,6 +110,11 @@ public interface CoreSession extends OutgoingFrames, Configuration
      */
     SocketAddress getRemoteAddress();
 
+    /**
+     * @return True if the websocket is open inbound
+     */
+    boolean isInputOpen();
+
     /**
      * @return True if the websocket is open outbound
      */
@@ -253,6 +258,12 @@ public SocketAddress getRemoteAddress()
             return null;
         }
 
+        @Override
+        public boolean isInputOpen()
+        {
+            return true;
+        }
+
         @Override
         public boolean isOutputOpen()
         {

jetty-websocket/websocket-core-common/src/main/java/org/eclipse/jetty/websocket/core/internal/WebSocketCoreSession.java

@@ -183,6 +183,12 @@ public SocketAddress getRemoteAddress()
         return getConnection().getEndPoint().getRemoteSocketAddress();
     }
 
+    @Override
+    public boolean isInputOpen()
+    {
+        return sessionState.isInputOpen();
+    }
+
     @Override
     public boolean isOutputOpen()
     {
@@ -416,8 +422,10 @@ public void demand(long n)
     {
         if (!demanding)
             throw new IllegalStateException("FrameHandler is not demanding: " + this);
+
         if (!sessionState.isInputOpen())
             throw new IllegalStateException("FrameHandler input not open: " + this);
+
         connection.demand(n);
     }
 

jetty-websocket/websocket-core-common/src/main/java/org/eclipse/jetty/websocket/core/internal/messages/ByteArrayMessageSink.java

@@ -13,12 +13,11 @@
 
 package org.eclipse.jetty.websocket.core.internal.messages;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodType;
 import java.nio.ByteBuffer;
 
+import org.eclipse.jetty.io.BufferCallbackAccumulator;
 import org.eclipse.jetty.util.BufferUtil;
 import org.eclipse.jetty.util.Callback;
 import org.eclipse.jetty.websocket.core.CoreSession;
@@ -29,9 +28,7 @@
 public class ByteArrayMessageSink extends AbstractMessageSink
 {
     private static final byte[] EMPTY_BUFFER = new byte[0];
-    private static final int BUFFER_SIZE = 65535;
-    private ByteArrayOutputStream out;
-    private int size;
+    private BufferCallbackAccumulator out;
 
     public ByteArrayMessageSink(CoreSession session, MethodHandle methodHandle)
     {
@@ -51,12 +48,12 @@ public void accept(Frame frame, Callback callback)
     {
         try
         {
-            size += frame.getPayloadLength();
+            long size = (out == null ? 0 : out.getLength()) + frame.getPayloadLength();
             long maxBinaryMessageSize = session.getMaxBinaryMessageSize();
             if (maxBinaryMessageSize > 0 && size > maxBinaryMessageSize)
             {
-                throw new MessageTooLargeException(String.format("Binary message too large: (actual) %,d > (configured max binary message size) %,d",
-                    size, maxBinaryMessageSize));
+                throw new MessageTooLargeException(
+                    String.format("Binary message too large: (actual) %,d > (configured max binary message size) %,d", size, maxBinaryMessageSize));
             }
 
             // If we are fin and no OutputStream has been created we don't need to aggregate.
@@ -71,19 +68,33 @@ public void accept(Frame frame, Callback callback)
                     methodHandle.invoke(EMPTY_BUFFER, 0, 0);
 
                 callback.succeeded();
+                session.demand(1);
                 return;
             }
 
-            aggregatePayload(frame);
+            // Aggregate the frame payload.
+            if (frame.hasPayload())
+            {
+                ByteBuffer payload = frame.getPayload();
+                if (out == null)
+                    out = new BufferCallbackAccumulator();
+                out.addEntry(payload, callback);
+            }
+
+            // If the methodHandle throws we don't want to fail callback twice.
+            callback = Callback.NOOP;
             if (frame.isFin())
             {
-                byte[] buf = out.toByteArray();
+                byte[] buf = out.takeByteArray();
                 methodHandle.invoke(buf, 0, buf.length);
             }
-            callback.succeeded();
+
+            session.demand(1);
         }
         catch (Throwable t)
         {
+            if (out != null)
+                out.fail(t);
             callback.failed(t);
         }
         finally
@@ -92,19 +103,7 @@ public void accept(Frame frame, Callback callback)
             {
                 // reset
                 out = null;
-                size = 0;
             }
         }
     }
-
-    private void aggregatePayload(Frame frame) throws IOException
-    {
-        if (frame.hasPayload())
-        {
-            ByteBuffer payload = frame.getPayload();
-            if (out == null)
-                out = new ByteArrayOutputStream(BUFFER_SIZE);
-            BufferUtil.writeTo(payload, out);
-        }
-    }
 }

jetty-websocket/websocket-core-common/src/main/java/org/eclipse/jetty/websocket/core/internal/messages/ByteBufferMessageSink.java

@@ -13,13 +13,13 @@
 
 package org.eclipse.jetty.websocket.core.internal.messages;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodType;
 import java.nio.ByteBuffer;
 import java.util.Objects;
 
+import org.eclipse.jetty.io.BufferCallbackAccumulator;
+import org.eclipse.jetty.io.ByteBufferPool;
 import org.eclipse.jetty.util.BufferUtil;
 import org.eclipse.jetty.util.Callback;
 import org.eclipse.jetty.websocket.core.CoreSession;
@@ -29,9 +29,7 @@
 
 public class ByteBufferMessageSink extends AbstractMessageSink
 {
-    private static final int BUFFER_SIZE = 65535;
-    private ByteArrayOutputStream out;
-    private int size;
+    private BufferCallbackAccumulator out;
 
     public ByteBufferMessageSink(CoreSession session, MethodHandle methodHandle)
     {
@@ -51,7 +49,7 @@ public void accept(Frame frame, Callback callback)
     {
         try
         {
-            size += frame.getPayloadLength();
+            long size = (out == null ? 0 : out.getLength()) + frame.getPayloadLength();
             long maxBinaryMessageSize = session.getMaxBinaryMessageSize();
             if (maxBinaryMessageSize > 0 && size > maxBinaryMessageSize)
             {
@@ -68,41 +66,53 @@ public void accept(Frame frame, Callback callback)
                     methodHandle.invoke(BufferUtil.EMPTY_BUFFER);
 
                 callback.succeeded();
+                session.demand(1);
                 return;
             }
 
-            aggregatePayload(frame);
+            // Aggregate the frame payload.
+            if (frame.hasPayload())
+            {
+                ByteBuffer payload = frame.getPayload();
+                if (out == null)
+                    out = new BufferCallbackAccumulator();
+                out.addEntry(payload, callback);
+            }
+
+            // If the methodHandle throws we don't want to fail callback twice.
+            callback = Callback.NOOP;
             if (frame.isFin())
-                methodHandle.invoke(ByteBuffer.wrap(out.toByteArray()));
+            {
+                ByteBufferPool bufferPool = session.getByteBufferPool();
+                ByteBuffer buffer = bufferPool.acquire(out.getLength(), false);
+                BufferUtil.clearToFill(buffer);
+                out.writeTo(buffer);
+                BufferUtil.flipToFlush(buffer, 0);
 
-            callback.succeeded();
+                try
+                {
+                    methodHandle.invoke(buffer);
+                }
+                finally
+                {
+                    bufferPool.release(buffer);
+                }
+            }
+
+            session.demand(1);
         }
         catch (Throwable t)
         {
+            if (out != null)
+                out.fail(t);
             callback.failed(t);
         }
         finally
         {
             if (frame.isFin())
             {
-                // reset
                 out = null;
-                size = 0;
             }
         }
     }
-
-    private void aggregatePayload(Frame frame) throws IOException
-    {
-        if (frame.hasPayload())
-        {
-            ByteBuffer payload = frame.getPayload();
-
-            if (out == null)
-                out = new ByteArrayOutputStream(BUFFER_SIZE);
-
-            BufferUtil.writeTo(payload, out);
-            payload.position(payload.limit()); // consume buffer
-        }
-    }
 }