Add support for abort signals when calling on and once

src/main.spec.ts

@@ -35,17 +35,171 @@ test('calls global handler when any event is emitted', () => {
   expect(handler).toBeCalledTimes(2)
 })
 
+describe('when on is used', () => {
+  test('calls the handler each time', () => {
+    const eventHandler = vi.fn()
+    const globalHandler = vi.fn()
+    const emitter = createEmitter<{ hello: void }>()
+
+    emitter.on('hello', eventHandler)
+    emitter.on(globalHandler)
+
+    expect(eventHandler).toHaveBeenCalledTimes(0)
+    expect(globalHandler).toHaveBeenCalledTimes(0)
+
+    emitter.emit('hello')
+
+    expect(eventHandler).toHaveBeenCalledTimes(1)
+    expect(globalHandler).toHaveBeenCalledTimes(1)
+
+    emitter.emit('hello')
+
+    expect(eventHandler).toHaveBeenCalledTimes(2)
+    expect(globalHandler).toHaveBeenCalledTimes(2)
+  })
+
+  test('does not call the handler if it is removed', () => {
+    const eventHandler = vi.fn()
+    const globalHandler = vi.fn()
+    const emitter = createEmitter<{ hello: void }>()
+
+    const offEvent = emitter.on('hello', eventHandler)
+    const offGlobal = emitter.on(globalHandler)
+
+    offEvent()
+    offGlobal()
+
+    emitter.emit('hello')
+
+    expect(eventHandler).toHaveBeenCalledTimes(0)
+    expect(globalHandler).toHaveBeenCalledTimes(0)
+  })
+
+  test('does not call the handler after it is removed', () => {
+    const eventHandler = vi.fn()
+    const globalHandler = vi.fn()
+    const emitter = createEmitter<{ hello: void }>()
+
+    const offEvent = emitter.on('hello', eventHandler)
+    const offGlobal = emitter.on(globalHandler)
+
+    emitter.emit('hello')
+
+    offEvent()
+    offGlobal()
+
+    emitter.emit('hello')
+    emitter.emit('hello')
+
+    expect(eventHandler).toHaveBeenCalledTimes(1)
+    expect(globalHandler).toHaveBeenCalledTimes(1)
+  })
+
+  test('does not call the handler if signal is aborted', () => {
+    const eventHandler = vi.fn()
+    const globalHandler = vi.fn()
+    const controller = new AbortController()
+    const emitter = createEmitter<{ hello: void }>()
+
+    emitter.on('hello', eventHandler, { signal: controller.signal })
+    emitter.on(globalHandler, { signal: controller.signal })
+
+    controller.abort()
+
+    emitter.emit('hello')
+
+    expect(eventHandler).toHaveBeenCalledTimes(0)
+    expect(globalHandler).toHaveBeenCalledTimes(0)
+  })
+
+  test('does not call the handler after signal is aborted', () => {
+    const eventHandler = vi.fn()
+    const globalHandler = vi.fn()
+    const controller = new AbortController()
+    const emitter = createEmitter<{ hello: void }>()
+
+    emitter.on('hello', eventHandler, { signal: controller.signal })
+    emitter.on(globalHandler, { signal: controller.signal })
+
+    emitter.emit('hello')
+
+    controller.abort()
+
+    emitter.emit('hello')
+    emitter.emit('hello')
+
+    expect(eventHandler).toHaveBeenCalledTimes(1)
+    expect(globalHandler).toHaveBeenCalledTimes(1)
+  })
+})
+
 describe('when once is used', () => {
   test('calls the handler one time', () => {
+    const eventHandler = vi.fn()
+    const globalHandler = vi.fn()
+    const emitter = createEmitter<{ hello: void }>()
+
+    emitter.once('hello', eventHandler)
+    emitter.once(globalHandler)
+
+    emitter.emit('hello')
+    emitter.emit('hello')
+
+    expect(eventHandler).toHaveBeenCalledOnce()
+    expect(globalHandler).toHaveBeenCalledOnce()
+  })
+
+  test('does not call the handler if it is removed', () => {
     const handler = vi.fn()
     const emitter = createEmitter<{ hello: void }>()
 
-    emitter.once('hello', handler)
+    const off = emitter.once('hello', handler)
+
+    off()
 
     emitter.emit('hello')
+
+    expect(handler).not.toHaveBeenCalled()
+  })
+
+  test('does not call the handler if signal is aborted', () => {
+    const handler = vi.fn()
+    const controller = new AbortController()
+    const emitter = createEmitter<{ hello: void }>()
+
+    emitter.once('hello', handler, { signal: controller.signal })
+
+    controller.abort()
+
     emitter.emit('hello')
 
-    expect(handler).toHaveBeenCalledOnce()
+    expect(handler).not.toHaveBeenCalled()
+  })
+})
+
+describe('when next is used', () => {
+  test('returns the event payload', async () => {
+    const emitter = createEmitter<{ hello: string }>()
+
+    const event = emitter.next('hello')
+
+    emitter.emit('hello', 'world')
+
+    await expect(event).resolves.toEqual('world')
+  })
+
+  test('when timeout is used, rejects if the event is not emitted', async () => {
+    vi.useFakeTimers()
+    const emitter = createEmitter()
+
+    await expect(() => {
+      const payload = emitter.next({ timeout: 100 })
+      vi.advanceTimersByTime(100)
+
+      return payload
+    }).rejects.toThrowError('Timeout waiting for global event after 100ms')
+
+    vi.useRealTimers()
   })
 })
 
@@ -218,26 +372,4 @@ test('next without event returns the global event payload', async () => {
     kind: 'hello',
     payload: 'world',
   })
-})
-
-test('next with event returns the event payload', async () => {
-  const emitter = createEmitter<{ hello: string }>()
-
-  const event = emitter.next('hello')
-
-  emitter.emit('hello', 'world')
-
-  await expect(event).resolves.toEqual('world')
-})
-
-test('next with timeout rejects if the event is not emitted', async () => {
-  vi.useFakeTimers()
-  const emitter = createEmitter()
-
-  await expect(() => {
-    const payload = emitter.next({ timeout: 100 })
-    vi.advanceTimersByTime(100)
-
-    return payload
-  }).rejects.toThrowError('Timeout waiting for global event after 100ms')
-})
+})

src/main.ts

@@ -19,6 +19,14 @@ export type NextOptions = {
   timeout?: number
 }
 
+export type EmitterOnOptions = {
+  signal?: AbortSignal
+}
+
+export type EmitterOnceOptions = {
+  signal?: AbortSignal
+}
+
 export class EmitterTimeoutError extends Error {
   constructor(event: string, timeout: number) {
     super(`Timeout waiting for ${event} event after ${timeout}ms`)
@@ -66,46 +74,43 @@ export function createEmitter<T extends Events>(options?: EmitterOptions) {
     onEvent(event, payload)
   }
 
-  function on(globalEventHandler: GlobalEventHandler<T>): () => void
-  function on<E extends Event>(event: E, handler: Handler<EventPayload<E>>): () => void
-  function on<E extends Event>(globalHandlerOrEvent: E | GlobalEventHandler<T>, handler?: Handler<EventPayload<E>>): () => void {
+  function on(globalEventHandler: GlobalEventHandler<T>, options?: EmitterOnOptions): () => void
+  function on<E extends Event>(event: E, handler: Handler<EventPayload<E>>, options?: EmitterOnOptions): () => void
+  function on<E extends Event>(globalHandlerOrEvent: E | GlobalEventHandler<T>, handlerOrOptions?: Handler<EventPayload<E>> | EmitterOnOptions, options?: EmitterOnOptions): () => void {
     if (isGlobalEventHandler(globalHandlerOrEvent)) {
-      globalHandlers.add(globalHandlerOrEvent)
-      
-      return () => off(globalHandlerOrEvent)
+      const globalHandlerOptions = typeof handlerOrOptions === 'object' ? handlerOrOptions : {}
+
+      return addGlobalHandler(globalHandlerOrEvent, globalHandlerOptions)
     }
 
+    const handler = typeof handlerOrOptions === 'function' ? handlerOrOptions : undefined
     const event = globalHandlerOrEvent
+    const handlerOptions = options ?? {}
 
     if (!handler) {
       throw new Error(`Handler must be given for ${String(event)} event`)
     }
 
-    const existing = handlers.get(event)
-
-    if (existing) {
-      existing.add(handler)
-    } else {
-      handlers.set(event, new Set([handler]))
-    }
-
-    return () => off(event, handler)
+    return addEventHandler(event, handler, handlerOptions)
   }
 
-  function once(globalEventHandler: GlobalEventHandler<T>): void
-  function once<E extends Event>(event: E, handler: Handler<EventPayload<E>>): void
-  function once<E extends Event>(globalHandlerOrEvent: E | GlobalEventHandler<T>, handler?: Handler<EventPayload<E>>): void {
+  function once(globalEventHandler: GlobalEventHandler<T>, options?: EmitterOnceOptions): () => void
+  function once<E extends Event>(event: E, handler: Handler<EventPayload<E>>, options?: EmitterOnceOptions): () => void
+  function once<E extends Event>(globalHandlerOrEvent: E | GlobalEventHandler<T>, handlerOrOptions?: Handler<EventPayload<E>> | EmitterOnceOptions, options?: EmitterOnceOptions): () => void {
     if (isGlobalEventHandler(globalHandlerOrEvent)) {
+      const globalHandlerOptions = typeof handlerOrOptions === 'object' ? handlerOrOptions : {}
+
       const callback: GlobalEventHandler<T> = (args) => {
         off(callback)
         globalHandlerOrEvent(args)
       }
 
-      on(callback)
-      return
+      return addGlobalHandler(callback, globalHandlerOptions)
     }
-    
+
     const event = globalHandlerOrEvent
+    const handler = typeof handlerOrOptions === 'function' ? handlerOrOptions : undefined
+    const handlerOptions = options ?? {}
 
     if (!handler) {
       throw new Error(`Handler must be given for ${String(event)} event`)
@@ -116,7 +121,7 @@ export function createEmitter<T extends Events>(options?: EmitterOptions) {
       handler(args)
     }
 
-    on(event, callback)
+    return addEventHandler(event, callback, handlerOptions)
   }
 
   function next(options?: NextOptions): Promise<GlobalEvent<T>>
@@ -194,6 +199,42 @@ export function createEmitter<T extends Events>(options?: EmitterOptions) {
     }))
   }
 
+  function addGlobalHandler(globalEventHandler: GlobalEventHandler<T>, options?: EmitterOnOptions): () => void {
+    const { signal } = options ?? {}
+
+    if(signal?.aborted) {
+      return () => {}
+    }
+
+    globalHandlers.add(globalEventHandler)
+
+    const offHandler = () => off(globalEventHandler)
+
+    signal?.addEventListener('abort', offHandler)
+
+    return offHandler
+  }
+
+  function addEventHandler<E extends Event>(event: E, handler: Handler<EventPayload<E>>, options?: EmitterOnOptions): () => void {
+    const { signal } = options ?? {}
+
+    if(signal?.aborted) {
+      return () => {}
+    }
+
+    if(!handlers.has(event)){
+      handlers.set(event, new Set())
+    }
+
+    handlers.get(event)?.add(handler)
+
+    const offHandler = () => off(event, handler)
+
+    signal?.addEventListener('abort', offHandler)
+
+    return offHandler
+  }
+
   return {
     on,
     off,