custom non blocking reader (#322)

The previous import workflow streaming approach used PipedReader, which
would block the web socket when receive many chunks of a single message
(like say a large import). Replaced PipedReader with a custom reader
that stores all the received text chunks in a queue w/o blocking the web
socket thread

Author: devhawk

transact/src/main/java/dev/dbos/transact/conductor/Conductor.java

@@ -44,8 +44,6 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.io.PipedReader;
-import java.io.PipedWriter;
 import java.io.Reader;
 import java.net.InetAddress;
 import java.net.URI;
@@ -62,6 +60,7 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executors;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -165,16 +164,12 @@ public void checkServerTrusted(X509Certificate[] certs, String authType) {}
 
   private class WebSocketListener implements WebSocket.Listener {
     private final StringBuilder messageBuffer = new StringBuilder();
-    private PipedWriter importPipeWriter = null;
+    private ImportFrameQueue importFrameQueue = null;
 
     private void closePipe() {
-      if (importPipeWriter != null) {
-        try {
-          importPipeWriter.close();
-        } catch (IOException e) {
-          logger.debug("Error closing import pipe writer", e);
-        }
-        importPipeWriter = null;
+      if (importFrameQueue != null) {
+        importFrameQueue.finish();
+        importFrameQueue = null;
       }
     }
 
@@ -188,12 +183,14 @@ public void onOpen(WebSocket ws) {
 
     @Override
     public CompletableFuture<?> onText(WebSocket ws, CharSequence data, boolean last) {
-      // Streaming import path: pipe each frame directly to the worker thread
-      if (importPipeWriter != null) {
+      logger.debug("onText data size {} last {}", data.length(), last);
+
+      // Streaming import path: queue each frame for the worker thread (non-blocking)
+      if (importFrameQueue != null) {
         try {
-          importPipeWriter.append(data);
-        } catch (IOException e) {
-          logger.error("Error writing frame to import pipe; streaming import aborted", e);
+          importFrameQueue.addFrame(data);
+        } catch (Exception e) {
+          logger.error("Error queuing frame for import; streaming import aborted", e);
           closePipe();
         }
         if (last) {
@@ -205,21 +202,18 @@ public CompletableFuture<?> onText(WebSocket ws, CharSequence data, boolean last
 
       // Detect import_workflow on the first frame of a new message
       if (messageBuffer.length() == 0 && isImportMessage(data)) {
-        try {
-          PipedReader pipeReader = new PipedReader(8 * 1024);
-          importPipeWriter = new PipedWriter(pipeReader);
-          importPipeWriter.append(data);
-          if (last) {
-            closePipe();
-          }
-          streamImportAsync(Conductor.this, ws, pipeReader);
-          ws.request(1);
-          return null;
-        } catch (IOException e) {
-          logger.error("Failed to start streaming import; falling back to buffered path", e);
+        logger.debug("import message detected");
+
+        importFrameQueue = new ImportFrameQueue();
+        importFrameQueue.addFrame(data);
+        streamImportAsync(Conductor.this, ws, importFrameQueue);
+        logger.debug("streamImportAsync started");
+        if (last) {
           closePipe();
-          // fall through to buffered path
         }
+
+        ws.request(1);
+        return null;
       }
 
       messageBuffer.append(data);
@@ -343,6 +337,62 @@ public void onError(WebSocket ws, Throwable error) {
     }
   }
 
+  /**
+   * A Reader backed by an unbounded queue of CharSequence frames. The producer side (onText) calls
+   * addFrame/finish which never block. The consumer side (streamImportAsync) reads data normally,
+   * blocking only when waiting for more frames to arrive. This avoids blocking the WebSocket I/O
+   * thread, which would prevent the TCP receive buffer from draining and cause the conductor's pong
+   * writes to time out.
+   */
+  private static class ImportFrameQueue extends Reader {
+    private final LinkedBlockingQueue<CharSequence> frames = new LinkedBlockingQueue<>();
+    private final AtomicBoolean done = new AtomicBoolean(false);
+    private CharSequence current = null;
+    private int pos = 0;
+
+    void addFrame(CharSequence data) {
+      frames.add(data);
+    }
+
+    void finish() {
+      if (done.compareAndSet(false, true)) {
+        frames.add(""); // wake up a blocked reader
+      }
+    }
+
+    @Override
+    public int read(char[] cbuf, int off, int len) throws IOException {
+      while (true) {
+        if (current != null && pos < current.length()) {
+          int n = Math.min(len, current.length() - pos);
+          for (int i = 0; i < n; i++) {
+            cbuf[off + i] = current.charAt(pos + i);
+          }
+          pos += n;
+          return n;
+        }
+        try {
+          CharSequence next = frames.poll(100, TimeUnit.MILLISECONDS);
+          if (next == null) {
+            if (done.get() && frames.isEmpty()) return -1;
+            continue;
+          }
+          if (next.length() == 0 && done.get()) return -1;
+          current = next;
+          pos = 0;
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          throw new IOException("Interrupted reading import stream", e);
+        }
+      }
+    }
+
+    @Override
+    public void close() {
+      finish();
+    }
+  }
+
   private static void writeFragmentedResponse(WebSocket ws, BaseResponse response)
       throws Exception {
     int fragmentSize = 128 * 1024; // 128k
@@ -372,7 +422,7 @@ private static class FragmentingOutputStream extends OutputStream {
       this.ws = ws;
       this.fragmentSize = fragmentSize;
       this.buffer = new byte[fragmentSize];
-      logger.debug("Created JdkFragmentingOutputStream with fragment size: {}", fragmentSize);
+      logger.debug("Created FragmentingOutputStream with fragment size: {}", fragmentSize);
     }
 
     @Override

transact/src/test/java/dev/dbos/transact/conductor/ConductorTest.java

@@ -11,6 +11,7 @@
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -1565,6 +1566,80 @@ public void canImportLargePayload() throws Exception {
     }
   }
 
+  // Regression test for: when a large import is being processed, the WebSocket I/O thread must
+  // remain free to deliver pong frames. The old PipedWriter implementation blocked onText() when
+  // the 8KB pipe buffer filled up, stalling the I/O thread and preventing pong delivery, which
+  // caused the server's pong write to time out and the connection to be reset.
+  //
+  // This test verifies the connection stays alive during a long-running import by requiring a
+  // ping/pong cycle to complete while importWorkflow() is blocked. If the I/O thread were stuck,
+  // the pong would not be delivered, the ping would time out, the connection would reset, and
+  // messageLatch would never fire.
+  @RetryingTest(3)
+  public void pingsSucceedDuringLargeImport() throws Exception {
+    class Listener extends MessageListener {
+      final CountDownLatch pingLatch = new CountDownLatch(1);
+      volatile boolean connectionReset = false;
+
+      @Override
+      public void onPing(WebSocket conn, Framedata frame) {
+        super.onPing(conn, frame); // sends pong back to the conductor
+        pingLatch.countDown();
+      }
+
+      @Override
+      public void onClose(WebSocket conn, int code, String reason, boolean remote) {
+        connectionReset = true;
+      }
+    }
+
+    // Block importWorkflow until the test confirms a ping was received (and pong sent back).
+    // At least one full ping/pong cycle must complete before the import finishes.
+    CountDownLatch importMayProceed = new CountDownLatch(1);
+    doAnswer(
+            invocation -> {
+              assertTrue(importMayProceed.await(10, TimeUnit.SECONDS), "import was not released");
+              return null;
+            })
+        .when(mockDB)
+        .importWorkflow(any());
+
+    var workflows = createLargeTestExportedWorkflows();
+    var serialized = Conductor.serializeExportedWorkflows(workflows);
+
+    Listener listener = new Listener();
+    testServer.setListener(listener);
+
+    // Ping fires frequently; timeout is long enough for normal test overhead but short enough
+    // that a missed pong would reset the connection before importMayProceed is released.
+    builder.pingPeriodMs(300).pingTimeoutMs(2000);
+
+    try (Conductor conductor = builder.build()) {
+      conductor.start();
+      assertTrue(listener.openLatch.await(5, TimeUnit.SECONDS), "open latch timed out");
+
+      listener.send(
+          MessageType.IMPORT_WORKFLOW,
+          "ping-during-import",
+          Map.of("serialized_workflow", serialized));
+
+      // Wait for a ping to arrive at the server (and pong to be sent) while import is blocked.
+      assertTrue(listener.pingLatch.await(5, TimeUnit.SECONDS), "no ping received during import");
+      assertFalse(listener.connectionReset, "connection was reset during import");
+
+      // Release importWorkflow and verify the import completes cleanly.
+      importMayProceed.countDown();
+      assertTrue(listener.messageLatch.await(15, TimeUnit.SECONDS), "import did not complete");
+      assertFalse(listener.connectionReset, "connection was reset after import");
+
+      JsonNode jsonNode = mapper.readTree(listener.message);
+      assertEquals("import_workflow", jsonNode.get("type").asText());
+      assertEquals("ping-during-import", jsonNode.get("request_id").asText());
+      assertNull(jsonNode.get("error_message"));
+      assertTrue(jsonNode.get("success").asBoolean());
+    }
+  }
+
   @RetryingTest(3)
   public void canExportLargePayload() throws Exception {
     MessageListener listener = new MessageListener();