Fixed PipeableSSRExecutorManager streaming

Author: smikhalevski

src/main/ssr/node/PipeableSSRExecutorManager.ts

@@ -1,51 +1,44 @@
-import { Writable } from 'stream';
-import { SSRExecutorManager, SSRExecutorManagerOptions } from '../SSRExecutorManager.js';
+import { SSRExecutorManager } from '../SSRExecutorManager.js';
+import { AbortError } from '../../utils.js';
 
 /**
  * Streaming executor manager for NodeJS environment.
  */
 export class PipeableSSRExecutorManager extends SSRExecutorManager {
-  /**
-   * The stream that includes both React rendering chunks and executor hydration chunks.
-   */
-  readonly stream: NodeJS.WritableStream;
+  private _isPiped = false;
+  private _pendingChunk = '';
 
   /**
-   * Creates a new {@link PipeableSSRExecutorManager} instance.
+   * Outputs executor hydration chunks into the provided
+   * [Writable Node.js Stream](https://nodejs.org/api/stream.html#writable-streams).
    *
-   * @param stream The output stream to which both React chunks and executor hydration chunks are written.
-   * @param options Additional options.
+   * @param stream The stream to write to.
+   * @template Stream The stream to write to.
    */
-  constructor(stream: NodeJS.WritableStream, options?: SSRExecutorManagerOptions) {
-    super(options);
-
-    this.stream = new Writable({
-      write: (chunk, encoding, callback) => {
-        stream.write(chunk, encoding, error => {
-          if (error) {
-            callback(error);
-            return;
-          }
-
-          if (!chunk.toString().endsWith('</script>')) {
-            callback();
-            return;
-          }
-
-          const hydrationChunk = this.nextHydrationChunk();
-
-          if (hydrationChunk !== '') {
-            stream.write(hydrationChunk, callback);
-            return;
-          }
-
-          callback();
-        });
-      },
-
-      final: callback => {
-        stream.end(callback);
-      },
-    });
+  pipe<Stream extends NodeJS.WritableStream>(stream: Stream): Stream {
+    if (this._isPiped) {
+      throw new Error('React Executor currently only supports piping to one stream');
+    }
+
+    this._isPiped = true;
+
+    const writeChunk = () => {
+      if ((this._pendingChunk ||= this.nextHydrationChunk()) && stream.write(this._pendingChunk)) {
+        console.log('EXECUTOR ' + this._pendingChunk + '\n\n');
+        this._pendingChunk = '';
+      }
+    };
+
+    this.subscribe(writeChunk);
+
+    stream.on('drain', writeChunk);
+
+    stream.on('error', () => this.abort(AbortError('The stream errored while writing data')));
+
+    stream.on('cancel', () => this.abort(AbortError('The stream closed early')));
+
+    writeChunk();
+
+    return stream;
   }
 }

src/test/ssr/node/PipeableSSRExecutorManager.test.ts

@@ -8,18 +8,20 @@ Date.now = () => 50;
 test('sends hydration chunk after the content chunk', async () => {
   const writeMock = vi.fn();
 
-  const outputStream = new Writable({
+  const writable = new Writable({
     write(chunk, _encoding, callback) {
       writeMock(chunk.toString());
       callback();
     },
   });
 
-  const manager = new PipeableSSRExecutorManager(outputStream);
+  const manager = new PipeableSSRExecutorManager();
 
   manager.getOrCreate('xxx', 111);
 
-  manager.stream.write('aaa</script>');
+  writable.write('aaa</script>');
+
+  manager.pipe(writable);
 
   await delay(200);
 
@@ -34,20 +36,21 @@ test('sends hydration chunk after the content chunk', async () => {
 test('does not send hydration chunk if nothing has changed', async () => {
   const writeMock = vi.fn();
 
-  const outputStream = new Writable({
+  const writable = new Writable({
     write(chunk, _encoding, callback) {
       writeMock(chunk.toString());
       callback();
     },
   });
 
-  const manager = new PipeableSSRExecutorManager(outputStream);
+  const manager = new PipeableSSRExecutorManager();
 
   manager.getOrCreate('xxx', 111);
 
-  manager.stream.write('aaa</script>');
-  manager.stream.write('bbb');
-  manager.stream.write('ccc</script>');
+  writable.write('aaa</script>');
+  manager.pipe(writable);
+  writable.write('bbb');
+  writable.write('ccc</script>');
 
   await delay(200);