chadtmiller.com/pds.js

Switch repository

Loading…

/**
 * com.atproto.sync.subscribeRepos consumer, narrowed to a set of collections.
 *
 * A commit event carries its own blocks as a CAR, so a record write arrives
 * whole. Nothing has to be read back from the PDS to learn what was written.
 */

import { parseCarFile } from '@pdsjs/core/car';
import { cborDecodeFirst, cidToString } from '@pdsjs/core/repo';

/** Backoff bounds for reconnecting after the socket closes. */
const RECONNECT_MIN_MS = 1_000;
const RECONNECT_MAX_MS = 30_000;

/**
 * @typedef {Object} FirehoseFrame
 * @property {{op: number, t?: string}} header
 * @property {Record<string, *>} body
 */

/**
 * Split one websocket message into its header and body items.
 * @param {Uint8Array} bytes
 * @returns {FirehoseFrame}
 */
export function decodeFrame(bytes) {
  const [header, consumed] = cborDecodeFirst(bytes);
  const [body] = cborDecodeFirst(bytes.subarray(consumed));
  return { header, body };
}

/**
 * @typedef {Object} RecordWrite
 * @property {string} did - repository the record belongs to
 * @property {string} collection
 * @property {string} rkey
 * @property {string} cid
 * @property {string} uri
 * @property {number} seq
 * @property {Record<string, *>} value
 */

/**
 * Record writes in one commit body that land in one of the given collections.
 * Deletes carry no cid and are left out; a repository record only disappears
 * when the repository is deleted, which starts no run.
 * @param {Record<string, *>} body - decoded #commit body
 * @param {string[]} collections
 * @returns {RecordWrite[]}
 */
export function commitWrites(body, collections) {
  const ops = Array.isArray(body.ops) ? body.ops : [];
  /** Path prefix to the collection it names. */
  const prefixes = collections.map(
    (/** @type {string} */ collection) =>
      /** @type {[string, string]} */ ([`${collection}/`, collection]),
  );
  /** @type {Array<{op: *, collection: string, rkey: string}>} */
  const wanted = [];
  for (const op of ops) {
    if (!op.cid || op.action === 'delete') continue;
    const match = prefixes.find(([prefix]) => op.path.startsWith(prefix));
    if (!match) continue;
    wanted.push({
      op,
      collection: match[1],
      rkey: op.path.slice(match[0].length),
    });
  }
  if (wanted.length === 0) return [];

  const blocks =
    body.blocks instanceof Uint8Array && body.blocks.length > 0
      ? parseCarFile(body.blocks).blocks
      : new Map();

  /** @type {RecordWrite[]} */
  const writes = [];
  for (const { op, collection, rkey } of wanted) {
    const cid = cidToString(op.cid);
    const block = blocks.get(cid);
    // tooBig commits omit the blocks; the record is then unreadable here.
    if (!block) continue;
    const [value] = cborDecodeFirst(block);
    writes.push({
      did: body.repo,
      collection,
      rkey,
      cid,
      uri: `at://${body.repo}/${op.path}`,
      seq: body.seq,
      value,
    });
  }
  return writes;
}

/**
 * @param {unknown} data - websocket message payload
 * @returns {Promise<Uint8Array>}
 */
async function toBytes(data) {
  if (data instanceof ArrayBuffer) return new Uint8Array(data);
  if (data instanceof Uint8Array) return data;
  if (typeof Blob !== 'undefined' && data instanceof Blob) {
    return new Uint8Array(await data.arrayBuffer());
  }
  throw new Error('firehose sent a non-binary frame');
}

/**
 * @typedef {Object} SubscribeOptions
 * @property {string} service - PDS base URL
 * @property {string[]} collections
 * @property {number} [cursor] - sequence to resume after
 * @property {AbortSignal} [signal]
 * @property {(write: RecordWrite) => Promise<void>|void} onWrite
 * @property {(seq: number) => Promise<void>|void} [onSeq] - every event, so a
 *   restart resumes past the ones that started no work
 * @property {(message: string) => void} [onNotice]
 */

/**
 * Consume one PDS's firehose until the signal aborts, reconnecting from the
 * last sequence seen. Resolves only when aborted.
 * @param {SubscribeOptions} options
 * @returns {Promise<void>}
 */
export async function subscribeRepos(options) {
  const {
    service,
    collections,
    signal,
    onWrite,
    onSeq = () => {},
    onNotice = () => {},
  } = options;
  const socketUrl = new URL(
    '/xrpc/com.atproto.sync.subscribeRepos',
    service.replace(/^http/, 'ws'),
  );
  let cursor = options.cursor;
  let backoff = RECONNECT_MIN_MS;

  while (!signal?.aborted) {
    if (cursor !== undefined)
      socketUrl.searchParams.set('cursor', String(cursor));
    const closed = await new Promise((resolve) => {
      const socket = new WebSocket(socketUrl);
      socket.binaryType = 'arraybuffer';
      /** @type {Promise<void>} */
      let queue = Promise.resolve();

      const abort = () => socket.close();
      signal?.addEventListener('abort', abort, { once: true });
      const finish = (/** @type {string|null} */ reason) => {
        signal?.removeEventListener('abort', abort);
        resolve(reason);
      };

      socket.addEventListener('open', () => {
        backoff = RECONNECT_MIN_MS;
        onNotice(`watching ${service} from sequence ${cursor ?? 'now'}`);
      });
      // Each message is handled to completion before the next one starts, so
      // a run cannot be published out of order with the events that caused it.
      socket.addEventListener('message', (event) => {
        queue = queue.then(async () => {
          const frame = decodeFrame(await toBytes(event.data));
          if (frame.header.op !== 1) {
            onNotice(`firehose error: ${frame.body.error ?? 'unknown'}`);
            return;
          }
          if (frame.header.t === '#commit') {
            for (const write of commitWrites(frame.body, collections)) {
              await onWrite(write);
            }
          }
          if (typeof frame.body.seq === 'number') {
            cursor = frame.body.seq;
            await onSeq(cursor);
          }
        });
        queue = queue.catch((err) => {
          onNotice(
            `event dropped: ${err instanceof Error ? err.message : err}`,
          );
        });
      });
      socket.addEventListener('error', () => finish('connection failed'));
      socket.addEventListener('close', () => {
        queue.then(() => finish(null));
      });
    });

    if (signal?.aborted) return;
    if (closed) onNotice(`${service}: ${closed}`);
    await new Promise((resolve) => setTimeout(resolve, backoff));
    backoff = Math.min(backoff * 2, RECONNECT_MAX_MS);
  }
}