chadtmiller.com/pds.js

Switch repository

Loading…

// The blob edge cache in the Worker entry: public getBlob GETs are served
// from caches.default on repeat, a cold ranged request passes through while
// a background fetch stores the whole blob, and nothing but complete 200s
// ever enters the cache.
import { afterEach, describe, expect, it, vi } from 'vitest';
import worker, { blobCacheable } from '../src/index.js';

const BLOB_URL =
  'https://pds.example.com/xrpc/com.atproto.sync.getBlob?did=did%3Aplc%3Ame&cid=bafkreiblob';

describe('blobCacheable', () => {
  it('accepts only getBlob GETs', () => {
    expect(blobCacheable(new Request(BLOB_URL))).toBe(true);
    expect(blobCacheable(new Request(BLOB_URL, { method: 'POST' }))).toBe(
      false,
    );
    expect(
      blobCacheable(
        new Request('https://pds.example.com/xrpc/com.atproto.sync.getRepo'),
      ),
    ).toBe(false);
  });
});

describe('worker entry blob cache flow', () => {
  /** @type {Map<string, Response>} */
  const store = new Map();
  const fakeCaches = {
    default: {
      /** @param {Request} request */
      match: async (request) => store.get(request.url)?.clone(),
      /** @param {Request} request @param {Response} response */
      put: async (request, response) => {
        store.set(request.url, response);
      },
    },
  };

  /**
   * A DO stub that answers getBlob like the real handler: 206 for a ranged
   * request, 200 with the whole body otherwise, or a fixed error.
   * @param {{error?: Response}} [options]
   */
  function makeEnv({ error } = {}) {
    let doCalls = 0;
    const env = /** @type {any} */ ({
      HOSTNAME: 'pds.example.com',
      PDS: {
        idFromName: () => 'default',
        get: () => ({
          /** @param {Request} request */
          fetch: async (request) => {
            doCalls += 1;
            if (error) return error.clone();
            if (request.headers.get('Range')) {
              return new Response('par', {
                status: 206,
                headers: {
                  'Content-Range': 'bytes 0-2/9',
                  'Cache-Control': 'public, max-age=31536000, immutable',
                },
              });
            }
            return new Response('wholeblob', {
              status: 200,
              headers: {
                'Content-Length': '9',
                'Cache-Control': 'public, max-age=31536000, immutable',
              },
            });
          },
        }),
      },
    });
    return { env, doCalls: () => doCalls };
  }

  /** @type {Promise<unknown>[]} */
  const pending = [];
  const ctx = /** @type {any} */ ({
    /** @param {Promise<unknown>} p */
    waitUntil: (p) => {
      pending.push(p);
    },
  });

  afterEach(() => {
    store.clear();
    pending.length = 0;
    vi.unstubAllGlobals();
  });

  it('serves the second request from the cache', async () => {
    vi.stubGlobal('caches', fakeCaches);
    const { env, doCalls } = makeEnv();

    const first = await worker.fetch(new Request(BLOB_URL), env, ctx);
    expect(first.status).toBe(200);
    expect(await first.text()).toBe('wholeblob');
    await Promise.all(pending);
    expect(doCalls()).toBe(1);

    const second = await worker.fetch(new Request(BLOB_URL), env, ctx);
    expect(second.status).toBe(200);
    expect(second.headers.get('X-Blob-Cache')).toBe('hit');
    expect(await second.text()).toBe('wholeblob');
    expect(doCalls()).toBe(1);
  });

  it('passes a cold ranged request through and warms the whole blob', async () => {
    vi.stubGlobal('caches', fakeCaches);
    const { env, doCalls } = makeEnv();

    const ranged = await worker.fetch(
      new Request(BLOB_URL, { headers: { Range: 'bytes=0-2' } }),
      env,
      ctx,
    );
    expect(ranged.status).toBe(206);
    expect(await ranged.text()).toBe('par');
    await Promise.all(pending);

    // The 206 itself is not stored; the background fetch stored the 200.
    expect(doCalls()).toBe(2);
    expect(store.get(BLOB_URL)?.status).toBe(200);

    const later = await worker.fetch(new Request(BLOB_URL), env, ctx);
    expect(later.headers.get('X-Blob-Cache')).toBe('hit');
    expect(doCalls()).toBe(2);
  });

  it('never caches an error and never warms off one', async () => {
    vi.stubGlobal('caches', fakeCaches);
    const { env, doCalls } = makeEnv({
      error: Response.json({ error: 'BlobNotFound' }, { status: 404 }),
    });

    await worker.fetch(new Request(BLOB_URL), env, ctx);
    await worker.fetch(new Request(BLOB_URL), env, ctx);
    await Promise.all(pending);

    expect(store.size).toBe(0);
    expect(doCalls()).toBe(2);
  });

  it('sends non-blob PDS traffic straight to the Durable Object', async () => {
    vi.stubGlobal('caches', fakeCaches);
    const { env, doCalls } = makeEnv();

    const response = await worker.fetch(
      new Request('https://pds.example.com/xrpc/com.atproto.sync.getBlob', {
        method: 'POST',
      }),
      env,
      ctx,
    );
    expect(response.status).toBe(200);
    expect(doCalls()).toBe(1);
    expect(store.size).toBe(0);
  });
});