Effection WebSocket Resource

use-websocket.ts

import { type Stream, type Operation, action, each, once, resource, spawn, createChannel, createSignal } from "effection";

export interface WebSocketHandle extends Stream<MessageEvent, CloseEvent> {
  send(value: string): Operation<void>;
  close(code?: number, reason?: string): Operation<void>;
}

export function useWebSocket(url: string | URL, protocols?: string | string[]) {
  return resource<WebSocketHandle>(function*(provide) {
    let input = createChannel<string, {code?: number; reason?: string}>();
    let output = createSignal<MessageEvent, CloseEvent>();
    let socket = new WebSocket(url, protocols);

    yield* spawn(function*() {
      let cause = yield* once(socket, 'error');
      throw new Error("WebSocket error", { cause });
    });

    yield* spawn(function*() {
      let inputs = yield* input.output;
      let next = yield* inputs.next();

      while (!next.done) {
        socket.send(next.value);
        next = yield* inputs.next();
      }
      let { code, reason } = next.value;
      socket.close(code, reason);
    })

    socket.onmessage = output.send;
    socket.onclose = output.close;

    yield* once(socket, 'open');

    let handle: WebSocketHandle = {
      send: input.input.send,
      close: (code, reason) => input.input.close({ code, reason }),
      [Symbol.iterator]: output.stream[Symbol.iterator],
    };

    try {
      yield* action(function*(resolve) {
        // remote close
        yield* spawn(function*() {
          for (let _ of yield* each(output.stream)) { yield* each.next; }
          resolve();
        });

        // local close
        yield* spawn(function*() {
          for (let _ of yield* each(input.output)) { yield* each.next; }
          resolve();
        });

        yield* provide(handle);
      });
    } finally {
      socket.close(1001);
      if (socket.readyState !== socket.CLOSED) {
        yield* once(socket, 'close');
      }
    }
  });
}