feat(core): add signal option to waitFor (#5042)

* feat(core): add signal option to waitFor

This PR adds a `signal?: AbortSignal` option to `WaitForOptions`, so that a call to `waitFor()` can be externally aborted.

* chore(core): update changeset

.changeset/pretty-chicken-lie.md

@@ -0,0 +1,5 @@
+---
+'xstate': minor
+---
+
+`waitFor()` now accepts a `{signal: AbortSignal}` in `WaitForOptions`

packages/core/src/waitFor.ts

@@ -9,6 +9,9 @@ interface WaitForOptions {
    * @defaultValue Infinity
    */
   timeout: number;
+
+  /** A signal which stops waiting when aborted. */
+  signal?: AbortSignal;
 }
 
 const defaultWaitForOptions: WaitForOptions = {
@@ -46,6 +49,11 @@ export function waitFor<TActorRef extends AnyActorRef>(
     ...options
   };
   return new Promise((res, rej) => {
+    const { signal } = resolvedOptions;
+    if (signal?.aborted) {
+      rej(signal.reason);
+      return;
+    }
     let done = false;
     if (isDevelopment && resolvedOptions.timeout < 0) {
       console.error(
@@ -56,14 +64,17 @@ export function waitFor<TActorRef extends AnyActorRef>(
       resolvedOptions.timeout === Infinity
         ? undefined
         : setTimeout(() => {
-            sub!.unsubscribe();
+            dispose();
             rej(new Error(`Timeout of ${resolvedOptions.timeout} ms exceeded`));
           }, resolvedOptions.timeout);
 
     const dispose = () => {
       clearTimeout(handle!);
       done = true;
       sub?.unsubscribe();
+      if (abortListener) {
+        signal!.removeEventListener('abort', abortListener);
+      }
     };
 
     function checkEmitted(emitted: SnapshotFrom<TActorRef>) {
@@ -73,6 +84,11 @@ export function waitFor<TActorRef extends AnyActorRef>(
       }
     }
 
+    /**
+     * If the `signal` option is provided, this will be the listener for its
+     * `abort` event
+     */
+    let abortListener: () => void | undefined;
     let sub: Subscription | undefined; // avoid TDZ when disposing synchronously
 
     // See if the current snapshot already matches the predicate
@@ -81,6 +97,16 @@ export function waitFor<TActorRef extends AnyActorRef>(
       return;
     }
 
+    // only define the `abortListener` if the `signal` option is provided
+    if (signal) {
+      abortListener = () => {
+        dispose();
+        // XState does not "own" the signal, so we should reject with its reason (if any)
+        rej(signal.reason);
+      };
+      signal.addEventListener('abort', abortListener);
+    }
+
     sub = actorRef.subscribe({
       next: checkEmitted,
       error: (err) => {

packages/core/test/waitFor.test.ts

@@ -196,4 +196,206 @@ describe('waitFor', () => {
       `[Error: Actor terminated without satisfying predicate]`
     );
   });
+
+  it('should not subscribe to the actor when it receives an aborted signal', async () => {
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          on: {
+            NEXT: 'b'
+          }
+        },
+        b: {
+          type: 'final'
+        }
+      }
+    });
+
+    const service = createActor(machine).start();
+    service.send({ type: 'NEXT' });
+
+    const controller = new AbortController();
+    const { signal } = controller;
+    controller.abort(new Error('Aborted!'));
+    const spy = jest.fn();
+    service.subscribe = spy;
+    try {
+      await waitFor(service, (state) => state.matches('b'), { signal });
+      fail('should have rejected');
+    } catch {
+      expect(spy).not.toHaveBeenCalled();
+    }
+  });
+
+  it('should not listen for the "abort" event when it receives an aborted signal', async () => {
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          on: {
+            NEXT: 'b'
+          }
+        },
+        b: {
+          type: 'final'
+        }
+      }
+    });
+
+    const service = createActor(machine).start();
+    service.send({ type: 'NEXT' });
+
+    const controller = new AbortController();
+    const { signal } = controller;
+    controller.abort(new Error('Aborted!'));
+
+    const spy = jest.fn();
+    signal.addEventListener = spy;
+
+    try {
+      await waitFor(service, (state) => state.matches('b'), { signal });
+      fail('should have rejected');
+    } catch {
+      expect(spy).not.toHaveBeenCalled();
+    }
+  });
+
+  it('should not listen for the "abort" event for actor in its final state that matches the predicate', async () => {
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          on: {
+            NEXT: 'b'
+          }
+        },
+        b: {
+          type: 'final'
+        }
+      }
+    });
+
+    const service = createActor(machine).start();
+    service.send({ type: 'NEXT' });
+
+    const controller = new AbortController();
+    const { signal } = controller;
+
+    const spy = jest.fn();
+    signal.addEventListener = spy;
+
+    await waitFor(service, (state) => state.matches('b'), { signal });
+    expect(spy).not.toHaveBeenCalled();
+  });
+
+  it('should immediately reject when it receives an aborted signal', async () => {
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          on: {
+            NEXT: 'b'
+          }
+        },
+        b: {
+          type: 'final'
+        }
+      }
+    });
+
+    const service = createActor(machine).start();
+    service.send({ type: 'NEXT' });
+
+    const controller = new AbortController();
+    const { signal } = controller;
+    controller.abort(new Error('Aborted!'));
+
+    await expect(
+      waitFor(service, (state) => state.matches('b'), { signal })
+    ).rejects.toMatchInlineSnapshot(`[Error: Aborted!]`);
+  });
+
+  it('should reject when the signal is aborted while waiting', async () => {
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          on: { NEXT: 'b' }
+        },
+        b: {}
+      }
+    });
+
+    const service = createActor(machine).start();
+    const controller = new AbortController();
+    const { signal } = controller;
+    setTimeout(() => controller.abort(new Error('Aborted!')), 10);
+
+    await expect(
+      waitFor(service, (state) => state.matches('b'), { signal })
+    ).rejects.toMatchInlineSnapshot(`[Error: Aborted!]`);
+  });
+
+  it('should stop listening for the "abort" event upon successful completion', async () => {
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          on: {
+            NEXT: 'b'
+          }
+        },
+        b: {
+          type: 'final'
+        }
+      }
+    });
+
+    const service = createActor(machine).start();
+    setTimeout(() => {
+      service.send({ type: 'NEXT' });
+    }, 10);
+
+    const controller = new AbortController();
+    const { signal } = controller;
+    const spy = jest.fn();
+    signal.removeEventListener = spy;
+
+    await waitFor(service, (state) => state.matches('b'), { signal });
+
+    expect(spy).toHaveBeenCalledTimes(1);
+  });
+
+  it('should stop listening for the "abort" event upon failure', async () => {
+    const machine = createMachine({
+      initial: 'a',
+      states: {
+        a: {
+          on: { NEXT: 'b' }
+        },
+        b: {
+          type: 'final'
+        }
+      }
+    });
+
+    const service = createActor(machine).start();
+
+    setTimeout(() => {
+      service.send({ type: 'NEXT' });
+    }, 10);
+
+    const controller = new AbortController();
+    const { signal } = controller;
+    const spy = jest.fn();
+    signal.removeEventListener = spy;
+
+    try {
+      await waitFor(service, (state) => state.matches('never'), { signal });
+      fail('should have rejected');
+    } catch {
+      expect(spy).toHaveBeenCalledTimes(1);
+    }
+  });
 });