Add a timeout option

Author: nex3

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 * Add `SyncMessagePort.receiveMessageIfAvailable()`.
 
+* Add `timeout` and `timeoutValue` options to
+  `SyncMessagePort.receiveMessage()`.
+
 ## 1.0.0
 
 * Initial release

lib/index.test.ts

@@ -6,7 +6,7 @@ import * as fs from 'fs';
 import * as p from 'path';
 import {MessagePort, Worker} from 'worker_threads';
 
-import {SyncMessagePort} from './index';
+import {SyncMessagePort, TimeoutException} from './index';
 
 describe('SyncMessagePort', () => {
   describe('sends a message', () => {
@@ -127,6 +127,70 @@ describe('SyncMessagePort', () => {
     });
   });
 
+  describe('timeout', () => {
+    it("returns a value if it's already available", () => {
+      const channel = SyncMessagePort.createChannel();
+      const port1 = new SyncMessagePort(channel.port1);
+      const port2 = new SyncMessagePort(channel.port2);
+      port1.postMessage('message');
+      expect(port2.receiveMessage({timeout: 0})).toBe('message');
+    });
+
+    it('returns a value if it becomes available before the timeout', () => {
+      const channel = SyncMessagePort.createChannel();
+      const port = new SyncMessagePort(channel.port1);
+
+      spawnWorker(
+        `
+          port.postMessage('ready');
+          setTimeout(() => {
+            port.postMessage('message');
+            port.close();
+          }, 100);
+        `,
+        channel.port2,
+      );
+
+      expect(port.receiveMessage()).toEqual('ready');
+      expect(port.receiveMessage({timeout: 200})).toEqual('message');
+    });
+
+    it('throws an error if it times out before a value is available', () => {
+      const channel = SyncMessagePort.createChannel();
+      const port = new SyncMessagePort(channel.port1);
+      expect(() => port.receiveMessage({timeout: 0})).toThrow(TimeoutException);
+    });
+
+    it('returns timeoutValue if it times out before a value is available', () => {
+      const channel = SyncMessagePort.createChannel();
+      const port = new SyncMessagePort(channel.port1);
+      expect(port.receiveMessage({timeout: 0, timeoutValue: 'timed out'})).toBe(
+        'timed out',
+      );
+    });
+
+    it('throws an error if the channel closes before the request times out', () => {
+      const channel = SyncMessagePort.createChannel();
+      const port = new SyncMessagePort(channel.port1);
+
+      spawnWorker(
+        `
+          port.postMessage('ready');
+          setTimeout(() => {
+            port.close();
+          }, 100);
+        `,
+        channel.port2,
+      );
+
+      expect(port.receiveMessage()).toEqual('ready');
+      // timeoutValue shouldn't take precedence over this error
+      expect(() =>
+        port.receiveMessage({timeout: 10000, timeoutValue: 'timed out'}),
+      ).toThrow();
+    });
+  });
+
   describe('with an asynchronous listener', () => {
     it('receives a message sent before listening', async () => {
       const channel = SyncMessagePort.createChannel();

lib/index.ts

@@ -8,6 +8,7 @@ import {
   MessageChannel,
   MessagePort,
   TransferListItem,
+  Worker,
   receiveMessageOnPort,
 } from 'worker_threads';
 
@@ -36,6 +37,33 @@ enum BufferState {
   Closed = 0b10,
 }
 
+/**
+ * Options that can be passed to {@link SyncMessagePort.receiveMessage}.
+ */
+export interface ReceiveMessageOptions {
+  /**
+   * The time (in milliseconds) to wait for a message before returning {@link
+   * timeoutValue} (if set) or throwing a [TimeoutException] otherwise.
+   */
+  timeout?: number;
+
+  /**
+   * If a message isn't received within {@link timeout} milliseconds, this value
+   * is returned. Ignored if {@link timeout} is not set.
+   */
+  timeoutValue?: unknown;
+}
+
+/**
+ * An exception thrown by {@link SyncMessagePort.receiveMessage} if a message
+ * isn't received within {@link ReceivedMessageOptions.timeout} milliseconds.
+ */
+export class TimeoutException extends Error {
+  constructor(message: string) {
+    super(message);
+  }
+}
+
 /**
  * A communication port that can receive messages synchronously from another
  * `SyncMessagePort`.
@@ -132,7 +160,6 @@ export class SyncMessagePort extends EventEmitter {
   }
 
   // TODO(nex3):
-  // * Add a timeout option to `receiveMessage()`
   // * Add an option to `receiveMessage()` to return a special value if the
   //   channel is closed.
 
@@ -143,7 +170,7 @@ export class SyncMessagePort extends EventEmitter {
    * Throws an error if the channel is closed, including if it closes while this
    * is waiting for a message.
    */
-  receiveMessage(): unknown {
+  receiveMessage(options?: ReceiveMessageOptions): unknown {
     if (this.listenerCount('message')) {
       throw new Error(
         'SyncMessageChannel.receiveMessage() may not be called while there ' +
@@ -156,14 +183,13 @@ export class SyncMessagePort extends EventEmitter {
     // `receiveMessageOnPort` and the call to `Atomics.wait()`, we won't
     // overwrite it. Use `Atomics.compareExchange` so that we don't overwrite
     // the "closed" state.
-    if (
-      Atomics.compareExchange(
-        this.buffer,
-        0,
-        BufferState.MessageSent,
-        BufferState.AwaitingMessage,
-      ) === BufferState.Closed
-    ) {
+    const previousState = Atomics.compareExchange(
+      this.buffer,
+      0,
+      BufferState.MessageSent,
+      BufferState.AwaitingMessage,
+    );
+    if (previousState === BufferState.Closed) {
       throw new Error("The SyncMessagePort's channel is closed.");
     }
 
@@ -173,10 +199,20 @@ export class SyncMessagePort extends EventEmitter {
     // If there's no new message, wait for the other port to flip the "new
     // message" indicator to 1. If it's been set to 1 since we stored 0, this
     // will terminate immediately.
-    Atomics.wait(this.buffer, 0, BufferState.AwaitingMessage);
+    const result = Atomics.wait(
+      this.buffer,
+      0,
+      BufferState.AwaitingMessage,
+      options?.timeout,
+    );
     message = receiveMessageOnPort(this.port);
     if (message) return message.message;
 
+    if (result === 'timed-out') {
+      if ('timeoutValue' in options!) return options.timeoutValue;
+      throw new TimeoutException('SyncMessagePort.receiveMessage() timed out.');
+    }
+
     // Update the state to 0b10 after the last message is consumed.
     const oldState = Atomics.and(this.buffer, 0, BufferState.Closed);
     // Assert the old state was either 0b10 or 0b11.
@@ -187,6 +223,7 @@ export class SyncMessagePort extends EventEmitter {
   /** See `MessagePort.close()`. */
   close(): void {
     Atomics.or(this.buffer, 0, BufferState.Closed);
+    Atomics.notify(this.buffer, 0);
     this.port.close();
   }
 }