chadtmiller.com/pds.js

Switch repository

Loading…

// The R2 streamed writer against a fake bucket: a small blob lands with one
// put, a large one stages through a multipart upload at a temporary key and
// commits by streaming to the final key, since the R2 binding has no
// server-side copy. The fake refuses a write to a key that another write
// holds, as R2 does with error 10058.
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { createBlobs } from '../src/index.js';

const DID = 'did:plc:r2writer';
const CID = 'bafkreir2demo';
const PART = 8 * 1024 * 1024;

/**
 * @param {ReadableStream} stream
 * @returns {Promise<Uint8Array>}
 */
async function drain(stream) {
  const reader = stream.getReader();
  /** @type {Uint8Array[]} */
  const chunks = [];
  let size = 0;
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    size += value.length;
  }
  const out = new Uint8Array(size);
  let at = 0;
  for (const chunk of chunks) {
    out.set(chunk, at);
    at += chunk.length;
  }
  return out;
}

function makeBucket() {
  /** @type {Map<string, {data: Uint8Array, customMetadata?: Record<string, string>}>} */
  const store = new Map();
  /** @type {{key: string, partSizes: number[], completed: boolean, aborted: boolean}[]} */
  const uploads = [];
  /** @type {Set<string>} */
  const writing = new Set();
  /** @type {string[]} */
  const writes = [];

  /**
   * @param {string} key
   * @param {string} op
   * @param {() => Promise<void>} apply
   */
  async function write(key, op, apply) {
    if (writing.has(key)) {
      throw new Error(
        `${op}: Reduce your concurrent request rate for the same object. (10058)`,
      );
    }
    writing.add(key);
    try {
      await new Promise((resolve) => setTimeout(resolve, 5));
      await apply();
      writes.push(`${op} ${key}`);
    } finally {
      writing.delete(key);
    }
  }

  return {
    store,
    uploads,
    writes,

    /**
     * @param {string} key
     * @param {Uint8Array|ReadableStream} value
     * @param {{customMetadata?: Record<string, string>}} [opts]
     */
    async put(key, value, opts) {
      await write(key, 'put', async () => {
        const data =
          value instanceof ReadableStream ? await drain(value) : value;
        store.set(key, { data, customMetadata: opts?.customMetadata });
      });
    },

    /** @param {string} key */
    async head(key) {
      const held = store.get(key);
      return held ? { size: held.data.length } : null;
    },

    /**
     * @param {string} key
     * @param {{range?: {offset: number, length: number}}} [options]
     */
    async get(key, options) {
      const held = store.get(key);
      if (!held) return null;
      const data = options?.range
        ? held.data.subarray(
            options.range.offset,
            options.range.offset + options.range.length,
          )
        : held.data;
      return {
        size: held.data.length,
        customMetadata: held.customMetadata,
        body: new Response(/** @type {any} */ (data)).body,
        arrayBuffer: async () => data.buffer,
      };
    },

    /** @param {string} key */
    async delete(key) {
      await write(key, 'delete', async () => {
        store.delete(key);
      });
    },

    /** @param {string} key */
    async createMultipartUpload(key) {
      /** @type {{n: number, data: Uint8Array}[]} */
      const parts = [];
      const upload = {
        key,
        partSizes: /** @type {number[]} */ ([]),
        completed: false,
        aborted: false,
      };
      uploads.push(upload);
      return {
        /**
         * @param {number} n
         * @param {Uint8Array} data
         */
        async uploadPart(n, data) {
          parts.push({ n, data: data.slice() });
          upload.partSizes.push(data.length);
          return { partNumber: n, etag: `etag-${n}` };
        },
        /** @param {unknown[]} list */
        async complete(list) {
          upload.completed = true;
          const size = parts.reduce((sum, p) => sum + p.data.length, 0);
          const data = new Uint8Array(size);
          let at = 0;
          for (const part of parts.sort((a, b) => a.n - b.n)) {
            data.set(part.data, at);
            at += part.data.length;
          }
          store.set(key, { data });
          return { size: list.length };
        },
        async abort() {
          upload.aborted = true;
        },
      };
    },
  };
}

/**
 * @param {number} size
 * @returns {Uint8Array}
 */
function bytes(size) {
  const out = new Uint8Array(size);
  for (let i = 0; i < size; i += 1) out[i] = (i * 17 + 3) % 256;
  return out;
}

beforeEach(() => {
  // The commit path streams through FixedLengthStream, a workerd global. An
  // identity transform stands in; the length claim goes unchecked here.
  vi.stubGlobal(
    'FixedLengthStream',
    class extends TransformStream {
      /** @param {number} _length */
      constructor(_length) {
        super();
      }
    },
  );
});

afterEach(() => {
  vi.unstubAllGlobals();
});

describe('R2 createWriter', () => {
  it('lands a small blob with one put and no multipart upload', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));
    const data = bytes(1000);

    const writer = await /** @type {any} */ (blobs).createWriter(DID);
    await writer.write(data.subarray(0, 400));
    await writer.write(data.subarray(400));
    await writer.commit(CID, 'video/mp4');

    expect(bucket.uploads).toEqual([]);
    expect([...bucket.store.keys()]).toEqual([`${DID}/${CID}`]);
    const stored = await blobs.get(DID, CID);
    expect(stored?.data).toEqual(data);
    expect(stored?.mimeType).toBe('video/mp4');
  });

  it('stages a large blob in uniform parts and streams it to the cid', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));
    const data = bytes(2 * PART + 1024);

    const writer = await /** @type {any} */ (blobs).createWriter(DID);
    for (let at = 0; at < data.length; at += 1024 * 1024) {
      await writer.write(data.subarray(at, at + 1024 * 1024));
    }
    await writer.commit(CID, 'video/mp4');

    expect(bucket.uploads.length).toBe(1);
    expect(bucket.uploads[0].partSizes).toEqual([PART, PART, 1024]);
    expect(bucket.uploads[0].completed).toBe(true);
    expect([...bucket.store.keys()]).toEqual([`${DID}/${CID}`]);
    const stored = await blobs.get(DID, CID);
    // Buffer.equals, since toEqual walks a big array element by element
    expect(Buffer.from(stored?.data ?? []).equals(Buffer.from(data))).toBe(
      true,
    );
    expect(stored?.mimeType).toBe('video/mp4');
  });

  it('commits a blob that is an exact multiple of the part size', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));
    const data = bytes(PART);

    const writer = await /** @type {any} */ (blobs).createWriter(DID);
    await writer.write(data);
    await writer.commit(CID, 'application/octet-stream');

    expect(bucket.uploads[0].partSizes).toEqual([PART]);
    const stored = await blobs.get(DID, CID);
    expect(Buffer.from(stored?.data ?? []).equals(Buffer.from(data))).toBe(
      true,
    );
  });

  it('abort ends the multipart upload and leaves nothing stored', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));

    const writer = await /** @type {any} */ (blobs).createWriter(DID);
    await writer.write(bytes(PART + 5));
    await writer.abort();

    expect(bucket.uploads.length).toBe(1);
    expect(bucket.uploads[0].aborted).toBe(true);
    expect(bucket.store.size).toBe(0);
  });

  it('abort before any part is a no-op', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));

    const writer = await /** @type {any} */ (blobs).createWriter(DID);
    await writer.write(bytes(10));
    await writer.abort();

    expect(bucket.uploads).toEqual([]);
    expect(bucket.store.size).toBe(0);
  });
});

describe('R2 getStream', () => {
  it('streams the whole object with its type', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));
    const data = bytes(100);
    await blobs.put(DID, CID, data, 'video/mp4');

    const found = await /** @type {any} */ (blobs).getStream(DID, CID);

    expect(found.mimeType).toBe('video/mp4');
    expect(await drain(found.body)).toEqual(data);
  });

  it('asks the bucket for the translated range', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));
    const data = bytes(100);
    await blobs.put(DID, CID, data, 'video/mp4');

    const found = await /** @type {any} */ (blobs).getStream(DID, CID, {
      start: 10,
      end: 19,
    });

    expect(await drain(found.body)).toEqual(data.subarray(10, 20));
  });

  it('answers null for an object the bucket lacks', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));

    expect(await /** @type {any} */ (blobs).getStream(DID, CID)).toBe(null);
  });
});

describe('R2 writes to one key', () => {
  it('stores the same blob uploaded five times at once with one put', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));
    const data = bytes(64);

    await Promise.all(
      Array.from({ length: 5 }, () => blobs.put(DID, CID, data, 'text/plain')),
    );

    expect(bucket.writes).toEqual([`put ${DID}/${CID}`]);
    expect((await blobs.get(DID, CID))?.data).toEqual(data);
  });

  it('streams the same blob from five writers at once with one put', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));
    const data = bytes(64);

    await Promise.all(
      Array.from({ length: 5 }, async () => {
        const writer = await /** @type {any} */ (blobs).createWriter(DID);
        await writer.write(data);
        await writer.commit(CID, 'text/plain');
      }),
    );

    expect(bucket.writes).toEqual([`put ${DID}/${CID}`]);
  });

  it('orders a delete and a put of one key, the put last', async () => {
    const bucket = makeBucket();
    const blobs = createBlobs(/** @type {any} */ (bucket));
    const data = bytes(64);
    await blobs.put(DID, CID, data, 'text/plain');

    await Promise.all([
      blobs.delete(DID, CID),
      blobs.put(DID, CID, data, 'text/plain'),
    ]);

    expect(bucket.writes).toEqual([
      `put ${DID}/${CID}`,
      `delete ${DID}/${CID}`,
      `put ${DID}/${CID}`,
    ]);
    expect((await blobs.get(DID, CID))?.data).toEqual(data);
  });

  it('writes again after R2 refuses a put with 10058', async () => {
    const bucket = makeBucket();
    const put = bucket.put;
    let refusals = 1;
    bucket.put = async (key, value, opts) => {
      if (refusals-- > 0) {
        throw new Error(
          'put: Reduce your concurrent request rate for the same object. (10058)',
        );
      }
      return put(key, value, opts);
    };
    const blobs = createBlobs(/** @type {any} */ (bucket));

    await blobs.put(DID, CID, bytes(64), 'text/plain');

    expect(bucket.writes).toEqual([`put ${DID}/${CID}`]);
  });

  it('passes on an error other than 10058 without a retry', async () => {
    const bucket = makeBucket();
    let calls = 0;
    bucket.put = async () => {
      calls += 1;
      throw new Error('put: Internal error (10001)');
    };
    const blobs = createBlobs(/** @type {any} */ (bucket));

    await expect(blobs.put(DID, CID, bytes(64), 'text/plain')).rejects.toThrow(
      '(10001)',
    );
    expect(calls).toBe(1);
  });
});