chadtmiller.com/pds.js

Switch repository

Loading…

import { zstdDecompressSync } from 'node:zlib';
import { cborEncodeDagCbor } from '@pdsjs/core/repo';
import Database from 'better-sqlite3';
import { afterEach, describe, expect, it, vi } from 'vitest';
import { WebSocketServer } from 'ws';
import { createAppViewStore } from '../src/store.js';
import { watch } from '../src/watch.js';
import { blockFrame } from './helpers/archive-block.js';

const COLLECTION = 'dev.pdsjs.git.pull';
const DID = 'did:plc:stranger';
const RECORD = {
  subject: { uri: 'at://did:plc:owner/dev.pdsjs.git.repo/proj' },
};

/**
 * @param {Uint8Array} bytes
 * @returns {Promise<Uint8Array>}
 */
const decompress = async (bytes) =>
  new Uint8Array(zstdDecompressSync(Buffer.from(bytes)));

/**
 * One live commit frame, in the envelope Jetstream v2 sends.
 * @param {number} seq
 * @param {string} rkey
 */
const liveFrame = (seq, rkey) =>
  JSON.stringify({
    $type: 'message',
    payload: {
      $type: 'network.bsky.jetstream.subscribeEvents#commit',
      did: DID,
      seq,
      operation: 'create',
      collection: COLLECTION,
      rkey,
      rev: '3krev',
      record: RECORD,
      cid: 'bafylive',
    },
  });

/**
 * The archive answering one plan page and one block. Every plan request
 * records its parsed body, so a test can assert the sequence asked for.
 * @param {number} seq - the sequence of the single archived create
 * @param {string} rkey
 * @param {number} sealedTipSeq
 */
function archive(seq, rkey, sealedTipSeq) {
  const frame = blockFrame([
    {
      seq,
      kind: 1,
      did: DID,
      collection: COLLECTION,
      rkey,
      rev: '3krev',
      payload: cborEncodeDagCbor(RECORD),
    },
  ]);
  /** @type {Array<{url: string, body: Record<string, *>, authorization: string|null}>} */
  const plans = [];
  /** @type {string[]} */
  const blocks = [];
  /** @type {typeof globalThis.fetch} */
  const fetch = async (input, init) => {
    const url = String(input);
    if (url.includes('planSnapshot')) {
      plans.push({
        url,
        body: JSON.parse(String(init?.body)),
        authorization: new Headers(init?.headers).get('authorization'),
      });
      return new Response(
        JSON.stringify({
          sealedTipSeq,
          plannedThroughSeq: sealedTipSeq,
          segments: [
            {
              name: 'seg_0000000000.jss',
              index: 0,
              mode: 'blocks',
              blocks: [{ first: 0, last: 0 }],
            },
          ],
        }),
        { status: 200 },
      );
    }
    blocks.push(url);
    return new Response(/** @type {BodyInit} */ (frame), { status: 200 });
  };
  return { fetch, plans, blocks };
}

describe('watch with an archive replay', () => {
  /** @type {WebSocketServer} */
  let server;

  afterEach(() => {
    server?.close();
    vi.unstubAllGlobals();
  });

  /**
   * @param {(socket: import('ws').WebSocket, connections: number) => void} onConnection
   * @returns {string}
   */
  const serve = (onConnection) => {
    server = new WebSocketServer({ port: 0 });
    let connections = 0;
    server.on('connection', (socket) => {
      connections += 1;
      onConnection(socket, connections);
    });
    const { port } = /** @type {import('node:net').AddressInfo} */ (
      server.address()
    );
    return `ws://127.0.0.1:${port}`;
  };

  it('replays the archive from the floor, then tails from the sealed tip', async () => {
    const store = createAppViewStore(new Database(':memory:'));
    const { fetch, plans, blocks } = archive(1000, '3karchived', 2000);
    vi.stubGlobal('fetch', fetch);
    const service = serve((socket) => socket.send(liveFrame(3000, '3klive')));

    const controller = new AbortController();
    /** @type {string[]} */
    const notices = [];
    await watch({
      store,
      service,
      collections: [COLLECTION],
      signal: controller.signal,
      onNotice: (message) => notices.push(message),
      replay: { afterSeq: 500, decompress },
      onSeq: (seq) => {
        if (seq >= 3000) controller.abort();
      },
    });

    expect(
      store.getRecord(`at://${DID}/${COLLECTION}/3karchived`),
    ).toMatchObject({ value: RECORD });
    expect(store.getRecord(`at://${DID}/${COLLECTION}/3klive`)).toMatchObject({
      value: RECORD,
    });
    expect(store.cursor('jetstream')).toBe(3000);
    expect(plans).toHaveLength(1);
    expect(plans[0].body.afterSeq).toBe(500);
    expect(plans[0].authorization).toBe(null);
    expect(blocks).toHaveLength(1);
    // Without a replay service of its own the archive is the live host.
    expect(plans[0].url).toBe(
      `${service}/xrpc/network.bsky.jetstream.planSnapshot`,
    );
    expect(notices).toContain('replaying the archive from 500');
    expect(notices.some((message) => message.includes('from 2000'))).toBe(true);
  });

  it('replays the gap when the server refuses the stored cursor', async () => {
    const store = createAppViewStore(new Database(':memory:'));
    store.setCursor('jetstream', 42);
    const { fetch, plans } = archive(50, '3kgap', 100);
    vi.stubGlobal('fetch', fetch);
    // Three handshakes that carry no event are how a refused cursor reads.
    const service = serve((socket, connections) => {
      if (connections <= 3) socket.close();
      else socket.send(liveFrame(200, '3klive'));
    });

    const controller = new AbortController();
    /** @type {string[]} */
    const notices = [];
    await watch({
      store,
      service,
      collections: [COLLECTION],
      signal: controller.signal,
      onNotice: (message) => notices.push(message),
      replay: {
        service: 'https://archive.example',
        apiKey: 'secret',
        decompress,
      },
      onSeq: (seq) => {
        if (seq >= 200) controller.abort();
      },
    });

    expect(store.getRecord(`at://${DID}/${COLLECTION}/3kgap`)).toMatchObject({
      value: RECORD,
    });
    expect(store.cursor('jetstream')).toBe(200);
    expect(plans).toHaveLength(1);
    expect(plans[0].url).toBe(
      'https://archive.example/xrpc/network.bsky.jetstream.planSnapshot',
    );
    expect(plans[0].body.afterSeq).toBe(42);
    expect(plans[0].authorization).toBe('Bearer secret');
    expect(notices).toContain('replaying the archive from 42');
    expect(notices.some((message) => message.includes('from 100'))).toBe(true);
  });
});