Merge pull request #5 from samthor/cache-exp

mux

Author: samthor

src/cache.ts

@@ -68,3 +68,19 @@ export function once<T = void>(fn: () => T) {
     return result;
   };
 }
+
+/**
+ * Lazy {@link WeakMap} factory.
+ */
+export function lazyWeak<W extends Object, V>(fn: (w: W) => V): (w: W) => V {
+  const wm = new WeakMap<W, V>();
+
+  return (w: W): V => {
+    let value = wm.get(w);
+    if (value === undefined) {
+      value = fn(w);
+      wm.set(w, value);
+    }
+    return value;
+  };
+}

src/generator.ts

@@ -1,7 +1,7 @@
-import { promiseWithResolvers, unresolvedPromise } from './promise.js';
-import { AbortSignalArgs } from './types.js';
-import type { WorkQueue } from './queue.js';
-import { promiseForSignal, symbolAbortSignal } from './internal.js';
+import { promiseWithResolvers, unresolvedPromise } from './promise.ts';
+import type { AbortSignalArgs } from './types.ts';
+import type { WorkQueue } from './queue.ts';
+import { promiseForSignal, symbolAbortSignal } from './internal.ts';
 
 /**
  * Combines the N passed async generators into a single generator which yields items in order,
@@ -209,3 +209,23 @@ export class AsyncGeneratorCache<T, Y> {
     return this._knownValues;
   }
 }
+
+/**
+ * Helper which iterates over a {@link AsyncGenerator}.
+ *
+ * Allows relatively easy access to the return value.
+ */
+export async function forEachAsync<T, TReturn, TNext>(
+  async: AsyncGenerator<T, TReturn, TNext>,
+  cb: (x: T) => Promise<TNext> | TNext,
+): Promise<TReturn> {
+  let result: IteratorResult<T, TReturn>;
+  let next: [TNext] | [] = [];
+
+  while (!(result = await async.next(...next)).done) {
+    const internalResult = await cb(result.value);
+    next = [internalResult];
+  }
+
+  return result.value;
+}

src/index.ts

@@ -15,7 +15,7 @@ export * from './listener.ts';
 export * from './maps.js';
 export * from './matcher.js';
 export * from './memoize.js';
-export * from './multiplex.ts';
+export * from './mux.ts';
 export * from './notify.ts';
 export * from './object-utils.js';
 export * from './primitives.js';

src/internal.ts

@@ -1,4 +1,4 @@
-import { unresolvedPromise } from './promise.js';
+import { unresolvedPromise } from './promise.ts';
 
 /**
  * Returns a {@link Promise} that rejects with {@link symbolAbortSignal} when aborted.

src/limit.ts

@@ -0,0 +1,69 @@
+/**
+ * Configuration for {@link buildLimiter}.
+ */
+export type LimitConfig = {
+  /**
+   * Maximum and initial total of tokens in this limiter.
+   *
+   * This must be `>=1` to be valid.
+   */
+  b: number;
+
+  /**
+   * Rate of renewal of tokens, per second.
+   *
+   * This must be `>0` to be valid.
+   */
+  r: number;
+};
+
+/**
+ * Fetches a {@link Promise} which resolves when a token is available.
+ *
+ * Throws if the {@link AbortSignal} is or becomes aborted before a token is available.
+ */
+export type Limiter = (signal: AbortSignal) => Promise<void>;
+
+/**
+ * Builds a rate-limiter which can be called to consume a token.
+ *
+ * Throws if the config is invalid (cowardly refuses to make an invalid limiter).
+ * Without a passed config, uses a default of 100 tokens, renews at 10/sec.
+ */
+export function buildLimiter(c?: LimitConfig): Limiter {
+  const maxTokens = (c?.b ?? 100.0) * 1.0;
+  const rateOfIncrease = (c?.r ?? 10) * 1.0;
+
+  if (maxTokens <= 0 || rateOfIncrease <= 0) {
+    throw new Error(`invalid LimitConfig, no requests allowed`);
+  }
+
+  let last = performance.now();
+  let tokens = maxTokens;
+
+  return async (signal: AbortSignal): Promise<void> => {
+    for (;;) {
+      signal.throwIfAborted();
+
+      const now = performance.now();
+      const delta = (now - last) / 1000.0;
+      const increase = delta * rateOfIncrease;
+      tokens = Math.min(maxTokens, tokens + increase);
+      last = now;
+      const secondsToWait = (1.0 - tokens) / rateOfIncrease;
+
+      if (secondsToWait <= 0.0) {
+        tokens -= 1.0;
+        return; // success
+      }
+
+      await new Promise<void>((r) => {
+        const timeout = setTimeout(r, secondsToWait * 1000.0);
+        signal.addEventListener('abort', () => {
+          clearTimeout(timeout);
+          r();
+        });
+      });
+    }
+  };
+}

src/listener.ts

@@ -44,7 +44,7 @@ export type NamedListeners<T extends Record<string, any>> = {
 export type SoloListener<V> = {
   addListener(listener: (data: V) => void, signal: AbortSignal): void;
   dispatch(data: V): boolean;
-  any(handler: (signal: AbortSignal) => void, signal: AbortSignal);
+  any(handler: (signal: AbortSignal) => void, signal?: AbortSignal);
   hasAny(): boolean;
 };
 

src/mux.ts

@@ -0,0 +1,195 @@
+import { soloListener } from './listener.ts';
+import { objectNotifyAll, objectWait } from './notify.ts';
+import { buildLinkQueue } from './queue.ts';
+
+export type CallToken = Object;
+
+export type MuxTask<Out> = { token: CallToken } & ({ signal: AbortSignal } | { data: Out });
+
+type InternalMuxTask<In, Out> =
+  | { token: CallToken } & (
+      | { signal: AbortSignal; respond: (x: SimpleResponse<In>) => void }
+      | { data: Out }
+    );
+
+export interface MuxSession<In, Out> {
+  signal: AbortSignal;
+
+  /**
+   * Are there any active calls as part of this {@link MuxSession}?
+   * The runner should probably not return if this is the case.
+   */
+  hasActive(): boolean;
+
+  /**
+   * Synchronously return the next task.
+   * This may return `undefined` even if the session is still active, check {@link hasActive}.
+   */
+  nextTask(): MuxTask<Out> | undefined;
+
+  /**
+   * Wait for another task to become ready.
+   */
+  waitForTask(): Promise<true>;
+
+  /**
+   * Process a message received for this {@link CallToken}.
+   */
+  handle(token: CallToken, data: In): void;
+
+  /**
+   * Process a remote shutdown of this {@link CallToken}.
+   */
+  stop(token: CallToken, error?: Error): void;
+}
+
+export type MuxFn<In, Out> = (session: MuxSession<In, Out>) => Promise<any>;
+
+type SimpleResponse<X> = { data?: X; error?: Error };
+
+class MuxSessionImpl<In, Out> implements MuxSession<In, Out> {
+  readonly controller = new AbortController();
+  public readonly signal = this.controller.signal;
+
+  readonly active = new Map<CallToken, (x: SimpleResponse<In>) => void>();
+
+  constructor(private readonly tasks: InternalMuxTask<In, Out>[]) {}
+
+  hasActive(): boolean {
+    return this.active.size !== 0;
+  }
+
+  nextTask(): MuxTask<Out> | undefined {
+    const t = this.tasks.shift();
+    if (!t) {
+      return;
+    }
+
+    if ('respond' in t) {
+      if (t.signal.aborted) {
+        return this.nextTask();
+      }
+      t.signal.addEventListener('abort', () => this.active.delete(t.token));
+      this.active.set(t.token, t.respond);
+      return { token: t.token, signal: t.signal };
+    } else if (!this.active.has(t.token)) {
+      // silently drop messages for already-dead calls
+      return this.nextTask();
+    }
+
+    return t;
+  }
+
+  async waitForTask(): Promise<true> {
+    while (!this.tasks.length) {
+      await objectWait(this.tasks);
+    }
+    return true;
+  }
+
+  handle(token: CallToken, data: In): void {
+    const fn = this.active.get(token);
+    fn?.({ data });
+  }
+
+  stop(token: CallToken, error?: Error): void {
+    const fn = this.active.get(token);
+    fn?.({ error });
+    this.active.delete(token);
+  }
+}
+
+export interface Call<In, Out> {
+  /**
+   * Listen to data from the other side.
+   */
+  readonly gen: AsyncGenerator<In, Error | void, void>;
+
+  /**
+   * Send the packet.
+   * This may queue the packet rather than send immediately.
+   */
+  send(data: Out): void;
+}
+
+export type MuxCall<In, Out> = {
+  call(signal: AbortSignal): Call<In, Out>;
+  addListener(cb: (error: Error) => void, signal: AbortSignal): void;
+};
+
+/**
+ * Builds a {@link MuxCall}.
+ */
+export function buildMux<In, Out>(runner: MuxFn<In, Out>): MuxCall<In, Out> {
+  let activeRunner: Promise<any> | undefined;
+  const tasks: InternalMuxTask<In, Out>[] = [];
+
+  const listener = soloListener<Error>();
+
+  const maybeStartRunner = () => {
+    if (!tasks.length || activeRunner !== undefined) {
+      return;
+    }
+
+    const session = new MuxSessionImpl(tasks);
+    const response: { error?: Error } = {};
+
+    activeRunner = runner(session)
+      .catch((error) => {
+        response.error = error;
+        listener.dispatch(error);
+      })
+      .finally(() => {
+        activeRunner = undefined;
+        maybeStartRunner();
+
+        // kill all still-active because the runner shut down
+        session.active.forEach((cb) => cb(response));
+      });
+  };
+
+  const call = (signal: AbortSignal): Call<In, Out> => {
+    if (signal.aborted) {
+      const gen = (async function* () {
+        try {
+          signal.throwIfAborted();
+        } catch (e: any) {
+          return e instanceof Error ? e : new Error(e);
+        }
+      })();
+      return { send() {}, gen };
+    }
+
+    const queue = buildLinkQueue<SimpleResponse<In>>();
+    const listener = queue.join(signal);
+
+    const gen = (async function* () {
+      for (;;) {
+        const m = await listener.next();
+        if (m?.data !== undefined) {
+          yield m?.data as In;
+        }
+        if (m?.error || m?.data === undefined) {
+          return m?.error as Error | void;
+        }
+      }
+    })();
+
+    // something object-like
+    const token: CallToken = new Object();
+    tasks.push({ token, signal, respond: queue.push.bind(queue) });
+    objectNotifyAll(tasks);
+
+    const send = (out: Out) => {
+      if (!signal.aborted) {
+        tasks.push({ token, data: out });
+      }
+    };
+
+    maybeStartRunner();
+
+    return { gen, send };
+  };
+
+  return { call, addListener: listener.addListener };
+}

src/queue.ts

@@ -1,5 +1,5 @@
-import { promiseWithResolvers } from './promise.js';
-import { promiseVoidForSignal } from './signal.js';
+import { promiseWithResolvers } from './promise.ts';
+import { promiseVoidForSignal } from './signal.ts';
 
 export class WorkQueue<T> {
   private pending: T[] = [];