WriteStream implementation (#132)

__tests__/streams.test.ts

@@ -1,5 +1,9 @@
 import { describe, it, expect, vi } from 'vitest';
-import { InterruptedStreamError, ReadStreamImpl } from '../router/streams';
+import {
+  InterruptedStreamError,
+  ReadStreamImpl,
+  WriteStreamImpl,
+} from '../router/streams';
 
 const noopCb = () => undefined;
 
@@ -334,3 +338,50 @@ describe('ReadStream unit', () => {
     });
   });
 });
+
+describe('WriteStream unit', () => {
+  it('should write', async () => {
+    const writeCb = vi.fn();
+    const stream = new WriteStreamImpl<number>(writeCb, noopCb);
+    stream.write(1);
+    stream.write(2);
+
+    expect(writeCb).toHaveBeenNthCalledWith(1, 1);
+    expect(writeCb).toHaveBeenNthCalledWith(2, 2);
+  });
+
+  it('should close the stream', async () => {
+    const closeCb = vi.fn();
+    const stream = new WriteStreamImpl<number>(noopCb, closeCb);
+
+    expect(stream.isClosed()).toBeFalsy();
+
+    stream.close();
+    expect(closeCb).toHaveBeenCalled();
+    expect(stream.isClosed()).toBeTruthy();
+  });
+
+  it('should throw when writing after close', async () => {
+    const stream = new WriteStreamImpl<number>(noopCb, noopCb);
+    stream.close();
+    expect(() => stream.write(1)).toThrowError(Error);
+  });
+
+  it('should handle close requests', async () => {
+    const stream = new WriteStreamImpl<number>(noopCb, noopCb);
+
+    expect(stream.isCloseRequested()).toBeFalsy();
+
+    const closeRequestP = stream.waitForCloseRequest();
+    expect(
+      await Promise.race([
+        new Promise((resolve) => setTimeout(() => resolve('timeout'), 10)),
+        closeRequestP,
+      ]),
+    ).toEqual('timeout');
+
+    stream.triggerCloseRequest();
+    expect(stream.isCloseRequested()).toBeTruthy();
+    await expect(closeRequestP).resolves.toEqual(undefined);
+  });
+});

router/streams.ts

@@ -103,11 +103,7 @@ export interface WriteStream<T> {
    * `waitForCloseRequest` returns a promise that resolves when the reader requests
    * to close the stream.
    */
-  waitForCloseRequest(): Promise<void>;
-  /**
-   * `waitForClose` returns a promise that resolves when the stream is closed
-   */
-  waitForClose(): Promise<void>;
+  waitForCloseRequest(): Promise<undefined>;
   /**
    * `isClosed` returns true if the stream was closed by the writer.
    */
@@ -130,7 +126,9 @@ export interface ReadWriteStream<TRead, TWrite> {
 
 /**
  * Internal implementation of a `ReadStream`.
- * Has internal methods to pushed data to the stream and close it.
+ * This won't be exposed as an interface to river
+ * consumers directly, it has internal river methods
+ * to pushed data to the stream and close it.
  */
 export class ReadStreamImpl<T> implements ReadStream<T> {
   /**
@@ -180,7 +178,7 @@ export class ReadStreamImpl<T> implements ReadStream<T> {
   /**
    * Resolves nextPromise
    */
-  private resolveNext: null | (() => void) = null;
+  private resolveNextPromise: null | (() => void) = null;
 
   constructor(closeRequestCallback: () => void) {
     this.closeRequestCallback = closeRequestCallback;
@@ -223,13 +221,13 @@ export class ReadStreamImpl<T> implements ReadStream<T> {
 
           if (!this.nextPromise) {
             this.nextPromise = new Promise<void>((resolve) => {
-              this.resolveNext = resolve;
+              this.resolveNextPromise = resolve;
             });
           }
 
           await this.nextPromise;
           this.nextPromise = null;
-          this.resolveNext = null;
+          this.resolveNextPromise = null;
         }
 
         // Unfortunately we have to use non-null assertion here, because T can be undefined
@@ -266,7 +264,7 @@ export class ReadStreamImpl<T> implements ReadStream<T> {
     this.didDrainDisposeValues = this.queue.length > 0;
     this.queue.length = 0;
 
-    this.resolveNext?.();
+    this.resolveNextPromise?.();
   }
 
   public isClosed(): boolean {
@@ -310,7 +308,7 @@ export class ReadStreamImpl<T> implements ReadStream<T> {
     }
 
     this.queue.push(value);
-    this.resolveNext?.();
+    this.resolveNextPromise?.();
   }
 
   /**
@@ -325,7 +323,7 @@ export class ReadStreamImpl<T> implements ReadStream<T> {
     }
 
     this.closed = true;
-    this.resolveNext?.();
+    this.resolveNextPromise?.();
     this.resolveClosePromise();
   }
 
@@ -336,3 +334,93 @@ export class ReadStreamImpl<T> implements ReadStream<T> {
     return this.queue.length > 0;
   }
 }
+
+/**
+ * Internal implementation of a `WriteStream`.
+ * This won't be exposed as an interface to river
+ * consumers directly, it has internal river methods
+ * to trigger a close request, a way to pass on close
+ * signals, and a way to push data to the stream.
+ */
+export class WriteStreamImpl<T> implements WriteStream<T> {
+  /**
+   * Passed via constructor to pass on write requests
+   */
+  private writeCb: (value: T) => void;
+  /**
+   * Passed via constructor to pass on close requests
+   */
+  private closeCb: () => void;
+  /**
+   * Whether the stream is closed.
+   */
+  private closed = false;
+  /**
+   * Whether the reader has requested to close the stream.
+   */
+  private closeRequested = false;
+  /**
+   * A promise that resolves when we recive a close request.
+   */
+  private closeRequestPromise: Promise<undefined>;
+  /**
+   * Resolves closePromise
+   */
+  private resolveRequestClosePromise: () => void = () => undefined;
+
+  constructor(writeCb: (value: T) => void, closeCb: () => void) {
+    this.writeCb = writeCb;
+    this.closeCb = closeCb;
+
+    this.closeRequestPromise = new Promise((resolve) => {
+      this.resolveRequestClosePromise = () => resolve(undefined);
+    });
+  }
+
+  public write(value: T): undefined {
+    if (this.closed) {
+      throw new Error('Cannot write to closed stream');
+    }
+
+    this.writeCb(value);
+  }
+
+  public close(): undefined {
+    if (this.closed) {
+      return;
+    }
+
+    this.closed = true;
+    this.closeCb();
+  }
+
+  public isCloseRequested(): boolean {
+    return this.closeRequested;
+  }
+
+  public waitForCloseRequest(): Promise<undefined> {
+    return this.closeRequestPromise;
+  }
+
+  public isClosed(): boolean {
+    return this.closed;
+  }
+
+  /**
+   * @internal meant for use within river, not exposed as a public API
+   *
+   * Triggers a close request.
+   */
+  public triggerCloseRequest(): undefined {
+    if (this.closeRequested) {
+      throw new Error('Cannot trigger close request multiple times');
+    }
+
+    if (this.closed) {
+      throw new Error('Cannot trigger close request on closed stream');
+    }
+
+    this.closeRequested = true;
+    this.resolveRequestClosePromise();
+  }
+}