Durable writes using promises (with retries and reverts!)

atomic-write.test.ts

import { Write, atomicWrite } from '../../util/atomic-write';

describe('write', () => {
  const createWrite = (retries: number, shouldSucceed: boolean): Write => {
    let attempts = 0;

    return {
      id: `${shouldSucceed ? 'success' : 'failure'}-${retries}-retries-${Math.floor(Math.random() * 100)}`,
      perform: jest.fn(() => (shouldSucceed && (attempts += 1) > retries ? Promise.resolve() : Promise.reject())),
      revert: jest.fn(),
      retries,
    };
  };

  type TestCase = {
    name: string;
    writes: Write[];
  };

  it.each<TestCase>([
    { name: '1/1 successful, 0 retries', writes: [createWrite(0, true)] },
    { name: '1/1 successful, 3 retries', writes: [createWrite(3, true)] },
    { name: '10/10 successful, 0 retries', writes: new Array(10).fill(0).map(() => createWrite(0, true)) },
    { name: '0/1 successful, 0 retries', writes: [createWrite(0, false)] },
    { name: '0/10 successful, 3 retries', writes: new Array(10).fill(0).map(() => createWrite(3, false)) },
    {
      name: '1/10 successful, 0 retries',
      writes: new Array(9)
        .fill(0)
        .map(() => createWrite(0, false))
        .concat(createWrite(0, true)),
    },
    {
      name: '9/10 successful, 0 retries',
      writes: new Array(9)
        .fill(0)
        .map(() => createWrite(0, true))
        .concat(createWrite(0, false)),
    },
    {
      name: '5/10 successful, 3 retries',
      writes: new Array(10).fill(0).map((_, index) => createWrite(3, index % 2 === 0)),
    },
  ])('performs an atomic write ($name)', async ({ writes }) => {
    const results = await atomicWrite(...writes);

    const expectedSuccesses = writes.filter(({ id }) => id.includes('success'));
    const successes = results.filter(({ success }) => success);
    const expectedFailures = writes.filter(({ id }) => !id.includes('success'));
    const failures = results.filter(({ success }) => !success);

    expect(results.map(({ id }) => id)).toEqual(writes.map(({ id }) => id));
    expect(successes.map(({ id }) => id)).toEqual(expectedSuccesses.map(({ id }) => id));
    expect(failures.map(({ id }) => id)).toEqual(expectedFailures.map(({ id }) => id));

    writes.forEach(({ perform, retries = 0 }) => {
      expect(perform).toHaveBeenCalledTimes(retries + 1);
    });

    if (!writes.every(({ id }) => id.includes('success'))) {
      writes.forEach(({ id, revert }) => {
        if (id.includes('success')) {
          expect(revert).toHaveBeenCalledTimes(1);
        } else {
          expect(revert).not.toHaveBeenCalled();
        }
      });
    }
  });
});

atomic-write.ts

export type Perform = () => Promise<any>;
export type Revert = () => Promise<void>;

/**
 * Represents an atomic write.
 */
export type Write = {
  id: string;
  retries?: number;
  perform: Perform;
  revert: Revert;
};

type WriteResultSuccess = {
  id: string;
  success: true;
  value: any;
};

type WriteResultFailure = {
  id: string;
  success: false;
};

export type WriteResult = WriteResultSuccess | WriteResultFailure;

/**
 * Executes all given `writes` in parallel. If any write fails, reverts the ones that succeeded.
 * Returns a list containing the result of each write attempt.
 */
export async function atomicWrite(...writes: Write[]): Promise<WriteResult[]> {
  const writesWithRetries = writes.map(({ perform, retries = 0 }) => withRetries(perform, retries));
  const results = await Promise.allSettled(writesWithRetries.map((write) => write()));

  if (!results.every(({ status }) => status === 'fulfilled')) {
    const reverts = writes.filter((_, index) => results[index].status === 'fulfilled');
    await Promise.allSettled(reverts.map(({ revert }) => revert()));
  }

  return writes.map(({ id }, index) => ({
    id,
    success: results[index].status !== 'rejected',
    value: (results[index] as any).value,
  }));
}

/**
 * Wraps the given `action` function and calls it `retries` times until it succeeds.
 */
function withRetries(action: Perform, retries: number): () => Promise<any> {
  return async () => {
    let attemptsRemaining = retries + 1;
    while (attemptsRemaining > 0) {
      try {
        return await action();
      } catch (err) {
        attemptsRemaining -= 1;
        if (attemptsRemaining <= 0) {
          throw err;
        }
      }
    }
  };
}