Optimize how SSE data is written

common/src/main/java/de/bluecolored/bluemap/common/web/MapRequestHandler.java

@@ -32,7 +32,6 @@
 
 import com.flowpowered.math.vector.Vector2i;
 
-import java.io.IOException;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
@@ -75,11 +74,7 @@ public MapRequestHandler(
                 // attempt to turn off buffering in upstream proxy
                 response.addHeader("X-Accel-Buffering", "no");
 
-                try {
-                    response.setBody(sseConnections.openConnection());
-                } catch (IOException e) {
-                    return new HttpResponse(HttpStatusCode.INTERNAL_SERVER_ERROR);
-                }
+                response.setBody(sseConnections::handleConnection);
                 return response;
             });
         }

common/src/main/java/de/bluecolored/bluemap/common/web/SseConnection.java

@@ -26,55 +26,30 @@
 
 import java.io.Closeable;
 import java.io.IOException;
-import java.io.InputStream;
-import java.io.PipedInputStream;
-import java.io.PipedOutputStream;
+import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
-import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
 
-import de.bluecolored.bluemap.core.util.stream.OnCloseInputStream;
 import lombok.SneakyThrows;
 
 /**
  * Represents a single Server-Sent Events (SSE) connection.
  * <p>
- * Read the events from the {@link PipedInputStream} returned from {@link #getInputStream()}.
- * Reading from the stream will block until a new event is delivered to it.
- * <p>
- * Events are queued via {@link #enqueue(String, String)} and delivered via a virtual thread
- * owned by this connection so a slow client only blocks its own delivery.
+ * Events can be queued via {@link #enqueue(String, String)} without blocking.
+ * Call {@link #run(OutputStream)} on the thread that owns the connection's output-stream (e.g.
+ * the HTTP connection's thread) to deliver queued events to it. This will block the calling thread
+ * until the connection is closed.
  */
 public class SseConnection implements Closeable {
 
-    private static final int PIPE_BUFFER_SIZE = 1024;
+    // how many messages can be queued up for sending before being dropped
+    private static final int QUEUE_CAPACITY = 64;
 
-    // how many messages can be queued up for sending (in addition to the above buffer)
-    // before being dropped
-    private static final int QUEUE_CAPACITY = 16;
-
-    private final PipedOutputStream pipeOut;
-    private final InputStream pipeIn;
     private final BlockingQueue<String[]> queue = new LinkedBlockingQueue<>(QUEUE_CAPACITY);
-    private final Thread sendThread;
     private volatile boolean closed = false;
     private volatile Runnable onClose;
-
-    public SseConnection() throws IOException {
-        // add a hook to the pipe to close the conneciton if the stream is closed
-        this.pipeOut = new PipedOutputStream();
-        this.pipeIn = new OnCloseInputStream(new PipedInputStream(pipeOut, PIPE_BUFFER_SIZE), SseConnection.this);
-
-        this.sendThread = Thread.ofVirtual().name("BlueMap-SSE-send").start(this::sendLoop);
-    }
-
-    /**
-     * Returns an {@link InputStream} to read events from.
-     * Closing it also closes this connection.
-     */
-    public InputStream getInputStream() {
-        return pipeIn;
-    }
+    private volatile Thread runningThread;
 
     public boolean isClosed() {
         return closed;
@@ -104,44 +79,51 @@ public void enqueue(String eventType, String data) {
         }
     }
 
-    private void sendLoop() {
+    /**
+     * Delivers queued events directly to {@code out}, blocking the calling thread until this
+     * connection is closed either explicitly via {@link #close()}, or because writing to
+     * {@code out} fails (happens if the client disconnects).
+     */
+    public void run(OutputStream out) throws IOException {
+        runningThread = Thread.currentThread();
+        String[] event;
         try {
             while (!closed) {
-                String[] event = queue.take();
-                send(event[0], event[1]);
+                try {
+                    event = queue.take();
+                } catch (InterruptedException _) {
+                    runningThread.interrupt();
+                    break;
+                }
+                send(out, event[0], event[1]);
             }
-        } catch (InterruptedException | IOException ignored) {}
+        } finally {
+            close();
+        }
     }
 
     @SneakyThrows(IOException.class)  // allows using this function in the forEach below
-    private void writeLine(String line){
-        pipeOut.write((line + "\n").getBytes(StandardCharsets.UTF_8));
+    private void writeLine(OutputStream out, String line) {
+        out.write((line + "\n").getBytes(StandardCharsets.UTF_8));
     }
 
     /**
      * Write one SSE event with optional data to the stream and flush it.
      *
-     * @throws IOException if the connection is closed or the client has disconnected
+     * @throws IOException if the client has disconnected
      */
-    private synchronized void send(String eventType, String data) throws IOException {
-        if (closed) throw new IOException("SSE connection is closed");
-        try {
-            writeLine("event: " + eventType);
-            data.lines().forEach(l -> writeLine("data: " + l));
-            pipeOut.write('\n');
-            pipeOut.flush();
-        } catch (IOException e) {
-            close();
-            throw e;
-        }
+    private void send(OutputStream out, String eventType, String data) throws IOException {
+        writeLine(out, "event: " + eventType);
+        data.lines().forEach(l -> writeLine(out, "data: " + l));
+        out.write('\n');
+        out.flush();
     }
 
     @Override
     public synchronized void close() {
         if (closed) return;
         closed = true;
-        sendThread.interrupt();
-        try { pipeOut.close(); } catch (IOException ignored) {}
+        if (runningThread != null) runningThread.interrupt();
         if (onClose != null) onClose.run();
     }
 

common/src/main/java/de/bluecolored/bluemap/common/web/SseConnectionManager.java

@@ -26,7 +26,7 @@
 
 import java.io.Closeable;
 import java.io.IOException;
-import java.io.InputStream;
+import java.io.OutputStream;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
@@ -54,15 +54,15 @@ public void removeHasConnectionsListener(Consumer<Boolean> listener) {
     }
 
     /**
-     * Creates a new {@link SseConnection}, registers it, and returns an {@link InputStream} suitable
-     * for use as an HTTP response body. When the stream is closed (either because the client
-     * disconnected or the server closed the connection), the connection is automatically removed
-     * from this manager.
+     * Creates a new {@link SseConnection}, registers it, and delivers events to {@code out} until
+     * the connection closes (either because the client disconnected or the server closed the
+     * connection), blocking the calling thread for that whole time. The connection is
+     * automatically removed from this manager once it closes.
      */
-    public InputStream openConnection() throws IOException {
+    public void handleConnection(OutputStream out) throws IOException {
         SseConnection connection = new SseConnection();
         add(connection);
-        return connection.getInputStream();
+        connection.run(out);
     }
 
     public void add(SseConnection connection) {

common/src/main/java/de/bluecolored/bluemap/common/web/http/ChunkedOutputStream.java

@@ -0,0 +1,136 @@
+/*
+ * This file is part of BlueMap, licensed under the MIT License (MIT).
+ *
+ * Copyright (c) Blue (Lukas Rieger) <https://bluecolored.de>
+ * Copyright (c) contributors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+package de.bluecolored.bluemap.common.web.http;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
+
+/**
+ * Wraps an {@link OutputStream}, buffering writes and framing them as HTTP/1.1 chunks.
+ * <p>
+ * {@link #endChunk()} ends the current chunk without flushing the wrapped stream.
+ * {@link #flush()} ends the current chunk *and* flushes the wrapped stream.
+ * <p>
+ * Closing this stream ends the current chunk and writes the terminating zero-length chunk, but
+ * doesn't close the wrapped stream since it's expected to outlive an individual chunked response.
+ * <p>
+ * Any write made after this stream has been closed throws an {@link IOException} to avoid bytes
+ * being written to the wrapped stream outside the required chunk framing.
+ */
+public class ChunkedOutputStream extends OutputStream {
+
+    private static final int AUTO_FLUSH_CHUNK_SIZE = 1024;
+    private static final byte[] CRLF = "\r\n".getBytes(StandardCharsets.UTF_8);
+
+    private final OutputStream out;
+    private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+    private boolean closed = false;
+
+    public ChunkedOutputStream(OutputStream out) {
+        this.out = out;
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+        ensureOpen();
+        buffer.write(b);
+        if (buffer.size() >= AUTO_FLUSH_CHUNK_SIZE) endChunk();
+    }
+
+    @Override
+    public void write(byte[] b, int off, int len) throws IOException {
+        ensureOpen();
+        buffer.write(b, off, len);
+        if (buffer.size() >= AUTO_FLUSH_CHUNK_SIZE) endChunk();
+    }
+
+    /**
+     * Writes out any currently buffered bytes as one HTTP chunk.
+     */
+    public void endChunk() throws IOException {
+        ensureOpen();
+        if (buffer.size() > 0) {
+            writeChunkHeader(buffer.size());
+            buffer.writeTo(out);
+            out.write(CRLF);
+            buffer.reset();
+        }
+    }
+
+    /**
+     * Writes {@code len} bytes from {@code b} starting at {@code off} as single chunk.
+     * This avoids the buffering overhead incurred by using {@code write(...)}.
+     * <p>
+     * Any currently buffered bytes are written with {@link #endChunk()} first.
+     */
+    public void writeChunk(byte[] b, int off, int len) throws IOException {
+        endChunk();
+        if (len > 0) {
+            writeChunkHeader(len);
+            out.write(b, off, len);
+            out.write(CRLF);
+        }
+    }
+
+    private void writeChunkHeader(int len) throws IOException {
+        out.write(Integer.toHexString(len).getBytes(StandardCharsets.UTF_8));
+        out.write(CRLF);
+    }
+
+    /**
+     * Ends the current chunk and flushes the wrapped stream to push all the buffered
+     * data to the client.
+     */
+    @Override
+    public void flush() throws IOException {
+        endChunk();
+        out.flush();
+    }
+
+    /**
+     * Ends the current chunk, writes the terminating zero-length chunk, and flushes the
+     * wrapped stream (without closing it).
+     */
+    @Override
+    public void close() throws IOException {
+        if (closed) return;
+        endChunk();
+        closed = true;
+        out.write('0');
+        out.write(CRLF);
+        out.write(CRLF);
+        out.flush();
+    }
+
+    /**
+     * @throws IOException if this stream has already been closed.
+     */
+    private void ensureOpen() throws IOException {
+        if (closed) throw new IOException("stream closed");
+    }
+
+}

common/src/main/java/de/bluecolored/bluemap/common/web/http/HttpResponse.java

@@ -40,33 +40,52 @@ public class HttpResponse implements Closeable, HttpHeaderCarrier {
     private @NonNull String version = "HTTP/1.1";
     private @NonNull HttpStatusCode statusCode;
     private @NonNull @Singular Map<String, HttpHeader> headers = new LinkedHashMap<>();
-    private @Nullable InputStream body;
+    private @Nullable HttpResponseStreamWriter body;
 
     public void setBody(@Nullable InputStream body) {
-        this.body = body;
+        this.body = body == null ? null : asStreamWriter(body);
     }
 
     public void setBody(byte[] data) {
-        if (data == null) {
-            this.body = null;
-            return;
-        }
-
-        setBody(new ByteArrayInputStream(data));
+        setBody(data == null ? null : new ByteArrayInputStream(data));
     }
 
     public void setBody(String data) {
-        if (data == null) {
-            this.body = null;
-            return;
-        }
+        setBody(data == null ? null : data.getBytes(StandardCharsets.UTF_8));
+    }
 
-        setBody(data.getBytes(StandardCharsets.UTF_8));
+    public void setBody(@Nullable HttpResponseStreamWriter streamWriter) {
+        this.body = streamWriter;
     }
 
     @Override
     public void close() throws IOException {
         if (body != null) body.close();
     }
 
+    /**
+     * Adapt an {@link InputStream} body into a {@link HttpResponseStreamWriter}
+     * that writes the data to the client in chunks.
+     */
+    private static HttpResponseStreamWriter asStreamWriter(InputStream body) {
+        return new HttpResponseStreamWriter() {
+            private final byte[] byteBuffer = new byte[1024];
+
+            @Override
+            public void write(ChunkedOutputStream out) throws IOException {
+                while (true) {
+                    int read = body.read(byteBuffer);
+                    if (read == -1) break;
+                    if (read == 0) continue;
+                    out.writeChunk(byteBuffer, 0, read);
+                }
+            }
+
+            @Override
+            public void close() throws IOException {
+                body.close();
+            }
+        };
+    }
+
 }