we have a socket

deno.json

···
   },
 
   "imports": {
+    "@repo/": "./src/",
     "@deno/vite-plugin": "npm:@deno/vite-plugin@^1.0.4",
     "@oak/oak": "jsr:@oak/oak@^17.1.4",
     "@std/assert": "jsr:@std/assert@^1.0.12",
···
     "lib": [
       "dom",
       "dom.iterable",
-      "deno.ns"
+      "deno.ns",
+      "esnext"
     ],
     "jsx": "react-jsx",
     "jsxImportSource": "react"

src/common/aborts.ts

···
+export type CancellableAbortController = AbortController & { cancel: () => void }
+
+export function makeAbort(): CancellableAbortController {
+  let cancelled    = false
+  const controller = new AbortController();
+
+  return {
+    cancel: () => { cancelled = true },
+    abort:  (r) => { !cancelled && controller.abort(r) },
+    signal: controller.signal
+  };
+}
+
+export function makeTimeoutAbort(ms: number): CancellableAbortController {
+  const controller = new AbortController();
+  const { promise, triggered, cancel } = makeTimeoutPromise(ms);
+
+  controller.signal.addEventListener("abort", () => {
+    if (!triggered())
+      cancel();
+  });
+
+  promise.then(() => {
+    if (!controller.signal.aborted)
+      controller.abort();
+  });
+
+  return {
+    cancel,
+    abort: controller.abort.bind(controller),
+    signal: controller.signal,
+  };
+}
+
+export function makeTimeoutPromise(ms: number) {
+  const { promise, resolve } = Promise.withResolvers<void>();
+
+  let triggered = false
+  const timeout = setTimeout(() => {
+    triggered = true;
+    resolve();
+  }, ms);
+
+  return {
+    promise,
+    cancel: () => clearTimeout(timeout),
+    triggered: () => triggered,
+  };
+}

src/common/blocking-atom.ts

···
+// simple blocking atom, for waiting for a value
+// cribbed mostly from https://github.com/ComFreek/async-playground
+
+import { Semaphore } from "./semaphore.ts";
+
+export class BlockingAtom<T> {
+  #sema = new Semaphore();
+  #item: T | undefined;
+
+  set(item: T) {
+    this.#item = item;
+    this.#sema.free();
+  }
+
+  async get(signal?: AbortSignal) {
+    if (await this.#sema.take(signal))
+      return this.#item;
+
+    return undefined;
+  }
+}

src/common/blocking-queue.ts

···
+// simple blocking queue, for turning streams into async pulls
+// cribbed mostly from https://github.com/ComFreek/async-playground
+
+import { Semaphore } from "./semaphore.ts";
+
+export class BlockingQueue<T> {
+  #sema = new Semaphore();
+  #items: T[] = [];
+  #maxsize: number | undefined;
+
+  constructor(maxsize: number | undefined = 1000) {
+    this.#maxsize = maxsize ? maxsize : undefined;
+  }
+
+  get size() {
+    return this.#items.length;
+  }
+
+  enqueue(...elements: T[]) {
+    for (const el of elements) {
+      if (this.#maxsize && this.size >= this.#maxsize)
+        throw Error("out of room");
+
+      this.#items.push(el);
+      this.#sema.free();
+    }
+  }
+
+  async dequeue(signal?: AbortSignal) {
+    if (await this.#sema.take(signal)) return this.poll();
+
+    throw Error("canceled dequeue");
+  }
+
+  poll() {
+    if (this.size > 0) {
+      return this.#items.shift() as T;
+    } else {
+      throw Error("no elements");
+    }
+  }
+}

src/common/once.ts

···
+export function onceMaker<F extends (...args: any[]) => void>(outercb?: () => void) {
+  let done = false;
+
+  return {
+    fired: () => done,
+
+    once: (innercb: F) => (...args: Parameters<F>): void => {
+      if (!done) {
+        done = true;
+
+        outercb?.();
+        innercb(...args);
+      }
+    },
+
+    many: (innercb: F) => (...args: Parameters<F>): void => {
+      if (done) return;
+      return innercb(...args);
+    }
+  };
+}

src/common/semaphore.ts

···
+// simple counting semaphore, for blocking async ops
+// cribbed mostly from https://github.com/ComFreek/async-playground
+
+export class Semaphore {
+  #counter = 0;
+  #resolvers: ((taken: boolean) => void)[] = [];
+
+  constructor(count = 0) {
+    this.#counter = count;
+  }
+
+  take(signal?: AbortSignal) {
+    return new Promise<boolean>((resolve) => {
+      if (signal?.aborted) return resolve(false);
+
+      // if there's resources available, use them
+
+      this.#counter--;
+      if (this.#counter >= 0) return resolve(true);
+
+      // otherwise add to pending
+      // and explicitly remove the resolver from the list on abort
+
+      this.#resolvers.push(resolve);
+      signal?.addEventListener("abort", () => {
+        const index = this.#resolvers.indexOf(resolve);
+        if (index >= 0) {
+          this.#resolvers.splice(index, 1);
+          this.#counter++;
+        }
+
+        resolve(false);
+      });
+    });
+  }
+
+  poll() {
+    if (this.#counter <= 0) return false;
+
+    this.#counter--;
+    return true;
+  }
+
+  free() {
+    this.#counter++;
+
+    if (this.#resolvers.length > 0) {
+      const resolver = this.#resolvers.shift();
+      resolver && queueMicrotask(() => resolver(true));
+    }
+  }
+}

src/common/socket.ts

···
+import { makeAbort, makeTimeoutAbort } from "./aborts.ts";
+import { BlockingAtom } from "./blocking-atom.ts";
+import { BlockingQueue } from "./blocking-queue.ts";
+import { onceMaker } from "./once.ts";
+
+export type StreamConfig = typeof STREAM_CONFIG_DEFAULT;
+const STREAM_CONFIG_DEFAULT = {
+  signal: undefined as AbortSignal | undefined,
+  queueSize: 1000,
+  threshPause: 800,
+  threshResume: 200,
+  onPause: () => {},
+  onResume: () => {},
+};
+
+type StreamEvent =
+  | ['yield', MessageEvent['data']]
+  | ['end']
+  | ['error', Error]
+  ;
+
+export async function* streamSocket(ws: WebSocket, config_: Partial<StreamConfig> = {}): AsyncGenerator<unknown> {
+  const config = { ...STREAM_CONFIG_DEFAULT, ...config_ }
+  const queue = new BlockingQueue<StreamEvent>(config.queueSize);
+  const gate  = onceMaker();
+
+  const onMessage = gate.many((m: MessageEvent) => {
+    queue.enqueue(['yield', m.data])
+
+    if (queue.size > config.threshPause)
+      config.onPause?.();
+  });
+
+  const onClose = gate.once(() => queue.enqueue(['end']))
+
+  const onError   = gate.once((e: Error) => {
+    // todo: why are we getting this on client shutdown?
+    if (e.message === "Unexpected EOF")
+      queue.enqueue(['end'])
+    else
+      queue.enqueue(['error', e]);
+  });
+
+  try {
+    ws.addEventListener("message", onMessage)
+    ws.addEventListener("error", onError)
+    ws.addEventListener("close", onClose)
+
+    while (true) {
+      const [event, value] = await queue.dequeue(config.signal)
+      if (queue.size < config.threshResume)
+        config.onResume?.();
+
+      switch (event) {
+        case "yield":
+          yield value;
+          continue;
+
+        case "end":
+          return;
+
+        case "error":
+          throw value;
+      }
+    }
+  } catch (e) {
+    if (!config.signal?.aborted) throw e;
+  } finally {
+    ws.removeEventListener("message", onMessage)
+    ws.removeEventListener("error", onError)
+    ws.removeEventListener("close", onClose)
+  }
+}
+
+export async function takeSocket(ws: WebSocket, ms?: number): Promise<MessageEvent['data']> {
+  const atom  = new BlockingAtom<MessageEvent['data']>();
+  const abort = ms ? makeTimeoutAbort(ms) : makeAbort();
+  const gate  = onceMaker();
+
+  // callback functions - consts so we can `off` them on clean up
+
+  const onMessage = gate.once((m: MessageEvent) => atom.set(m.data));
+  const onClose   = gate.once(() => abort.abort('closed'));
+  const onError   = gate.once((e: Error) => {
+    if (e.message == "Unexpected EOF")
+      // todo: why are we getting this on client shutdown?
+      abort.abort('closed')
+    else
+      abort.abort(e)
+  });
+
+  try {
+    ws.addEventListener("message", onMessage)
+    ws.addEventListener("close", onClose)
+    ws.addEventListener("error", onError)
+
+    return await atom.get(abort.signal);
+  } finally {
+    if (!abort.signal.aborted)
+      abort.cancel()
+
+    ws.removeEventListener("message", onMessage)
+    ws.removeEventListener("close", onClose)
+    ws.removeEventListener("error", onError)
+  }
+}

src/common/strict-map.ts

···
+export class StrictMap<K, V> extends Map<K, V> {
+  require(key: K): V {
+    if (!this.has(key)) throw Error(`key is required but not in the map`);
+
+    return this.get(key) as V;
+  }
+
+  ensure(key: K, maker: () => V): V {
+    if (!this.has(key)) {
+      this.set(key, maker());
+    }
+
+    return this.get(key) as V;
+  }
+
+  update(key: K, update: (v?: V) => V) {
+    const current = this.get(key);
+    this.set(key, update(current));
+  }
+}

src/server/main.ts

···
 import { Application } from "@oak/oak/application";
 import { parseArgs } from "jsr:@std/cli/parse-args";
 
-import { apiMiddleware } from "./routes.api.ts";
-import { socketMiddleware } from "./routes.socket.ts";
+import { apiMiddleware } from "./api/middleware.ts";
+import { socketMiddleware } from "./socket/middleware.ts";
 import { makeSpaRoute, makeStaticRoute } from "./routes.static.ts";
 import { notFound } from "./routes.error.ts";
 

src/server/routes.socket.ts

···
-import { Middleware } from "@oak/oak/middleware";
-
-export const socketMiddleware: (path: string) => Middleware =
-  (pathname) => async (ctx, next) => {
-    // must be an upgradable websocket connection to api root
-    if (ctx.request.url.pathname !== pathname || !ctx.isUpgradable) {
-      return await next();
-    }
-
-    // the we'll upgrade and connect a handler
-    const ws = ctx.upgrade();
-    ws.onopen = () => {
-      console.log("Connected to client");
-      ws.send("Hello from server!");
-    };
-
-    ws.onmessage = (m) => {
-      console.log("Got message from client: ", m.data);
-      ws.send(m.data as string);
-      ws.close();
-    };
-
-    ws.onclose = () => console.log("Disconncted from client");
-  };

src/server/socket/handler.ts

···
+import { streamSocket, takeSocket } from "@repo/common/socket.ts";
+
+export async function socketHandler(this: WebSocket) {
+  try {
+    this.send('your name?')
+
+    const name = await takeSocket(this, 5000)
+    if (!name) return
+
+    this.send(`welcome, ${name}`)
+    for await (const message of streamSocket(this)) {
+      this.send(`${name} said: ${message}`)
+    }
+  } catch (e) {
+    console.error(e)
+  } finally {
+    if (this.readyState !== this.CLOSED)
+      this.close()
+
+    console.log("kthxbye");
+  }
+}

src/server/socket/middleware.ts

···
+import { Middleware } from "@oak/oak/middleware";
+import { socketHandler } from "./handler.ts";
+
+export const socketMiddleware: (path: string) => Middleware = (pathname) => async (ctx, next) => {
+  // must be an upgradable websocket connection to api root
+  if (ctx.request.url.pathname !== pathname || !ctx.isUpgradable)
+    return await next()
+
+  // then we'll upgrade and connect a handler
+  const ws  = ctx.upgrade();
+  ws.onopen = socketHandler;
+};