Issue #4475 - proof of concept

Signed-off-by: Lachlan Roberts <[email protected]>

jetty-websocket/javax-websocket-client-impl/src/main/java/org/eclipse/jetty/websocket/jsr356/endpoints/JsrEndpointEventDriver.java

@@ -23,6 +23,7 @@
 import java.io.Reader;
 import java.nio.ByteBuffer;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
 import javax.websocket.CloseReason;
 import javax.websocket.Endpoint;
 import javax.websocket.MessageHandler;
@@ -190,7 +191,23 @@ public void onTextFrame(ByteBuffer buffer, boolean fin) throws IOException
             }
             else if (wrapper.wantsStreams())
             {
-                final MessageReader stream = new MessageReader(new MessageInputStream());
+                final CountDownLatch completed = new CountDownLatch(1);
+                final MessageReader stream = new MessageReader(new MessageInputStream())
+                {
+                    @Override
+                    public void messageComplete()
+                    {
+                        super.messageComplete();
+                        try
+                        {
+                            completed.await();
+                        }
+                        catch (Exception e)
+                        {
+                            throw new RuntimeException(e);
+                        }
+                    }
+                };
                 activeMessage = stream;
 
                 dispatch(new Runnable()
@@ -201,6 +218,7 @@ public void run()
                     {
                         MessageHandler.Whole<Reader> handler = (Whole<Reader>)wrapper.getHandler();
                         handler.onMessage(stream);
+                        completed.countDown();
                     }
                 });
             }

jetty-websocket/javax-websocket-server-impl/src/test/java/org/eclipse/jetty/websocket/jsr356/server/TextStreamTest.java

@@ -23,10 +23,14 @@
 import java.io.Writer;
 import java.net.URI;
 import java.util.Random;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import javax.websocket.ClientEndpoint;
 import javax.websocket.ContainerProvider;
+import javax.websocket.Endpoint;
+import javax.websocket.EndpointConfig;
+import javax.websocket.MessageHandler;
 import javax.websocket.OnMessage;
 import javax.websocket.Session;
 import javax.websocket.WebSocketContainer;
@@ -36,11 +40,15 @@
 import org.eclipse.jetty.server.Server;
 import org.eclipse.jetty.server.ServerConnector;
 import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.eclipse.jetty.util.BlockingArrayQueue;
+import org.eclipse.jetty.util.IO;
 import org.eclipse.jetty.websocket.jsr356.server.deploy.WebSocketServerContainerInitializer;
+import org.hamcrest.Matchers;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -49,6 +57,8 @@ public class TextStreamTest
     private static final String PATH = "/echo";
     private static final String CHARS = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
 
+    private static CompletableFuture<QueuedTextStreamer> queuedTextStreamerFuture = new CompletableFuture<>();
+
     private Server server;
     private ServerConnector connector;
     private WebSocketContainer wsClient;
@@ -62,8 +72,8 @@ public void prepare() throws Exception
 
         ServletContextHandler context = new ServletContextHandler(server, "/", true, false);
         ServerContainer container = WebSocketServerContainerInitializer.configureContext(context);
-        ServerEndpointConfig config = ServerEndpointConfig.Builder.create(ServerTextStreamer.class, PATH).build();
-        container.addEndpoint(config);
+        container.addEndpoint(ServerEndpointConfig.Builder.create(ServerTextStreamer.class, PATH).build());
+        container.addEndpoint(ServerEndpointConfig.Builder.create(QueuedTextStreamer.class, "/test").build());
 
         server.start();
 
@@ -125,6 +135,26 @@ public void testMoreThanLargestMessageOneByteAtATime() throws Exception
         assertArrayEquals(data, client.getEcho());
     }
 
+    @Test
+    public void test() throws Exception
+    {
+        URI uri = URI.create("ws://localhost:" + connector.getLocalPort() + "/test");
+        ClientTextStreamer client = new ClientTextStreamer();
+        Session session = wsClient.connectToServer(client, uri);
+
+        final int numLoops = 20;
+        for (int i = 0; i < numLoops; i++)
+            session.getBasicRemote().sendText(Integer.toString(i));
+        session.close();
+
+        QueuedTextStreamer queuedTextStreamer = queuedTextStreamerFuture.get(5, TimeUnit.SECONDS);
+        for (int i = 0; i < numLoops; i++)
+        {
+            String msg = queuedTextStreamer.messages.poll(5, TimeUnit.SECONDS);
+            assertThat(msg, Matchers.is(Integer.toString(i)));
+        }
+    }
+
     private char[] randomChars(int size)
     {
         char[] data = new char[size];
@@ -183,4 +213,34 @@ public void echo(Session session, Reader input) throws IOException
             }
         }
     }
+
+    public static class QueuedTextStreamer extends Endpoint implements MessageHandler.Whole<Reader>
+    {
+        private BlockingArrayQueue<String> messages = new BlockingArrayQueue<>();
+
+        public QueuedTextStreamer()
+        {
+            queuedTextStreamerFuture.complete(this);
+        }
+
+        @Override
+        public void onOpen(Session session, EndpointConfig config)
+        {
+            session.addMessageHandler(this);
+        }
+
+        @Override
+        public void onMessage(Reader input)
+        {
+            try
+            {
+                Thread.sleep(Math.abs(new Random().nextLong() % 200));
+                messages.add(IO.toString(input));
+            }
+            catch (Exception e)
+            {
+                e.printStackTrace();
+            }
+        }
+    }
 }