TryLock with Postgres

locks.ts

import { hash } from 'bun';
import { sql } from './db';
import type { Nullable } from './types';

export const LOCKS_TABLE = 'locks';
const DEFAULT_LOCK_TIMEOUT = 60 * 1000; // 1 minute
let tableSetupDone = false;

type LockOptions = {
  timeout?: number;
};

/**
 * Tries to acquire lock and returns immediately.
 *
 * This is opposed to blocking lock where the function await until it can acquire the lock.
 */
export async function tryLock(
  lockName: string,
  { timeout = DEFAULT_LOCK_TIMEOUT }: LockOptions = {},
): Promise<[boolean, () => Promise<void>]> {
  const now = Date.now();

  const hasLock = await sql.begin(async (sql) => {
    if (!tableSetupDone) {
      // Setup table
      await sql`CREATE TABLE IF NOT EXISTS ${sql(LOCKS_TABLE)} (key TEXT UNIQUE PRIMARY KEY, created_at NUMERIC NOT NULL, expires_at NUMERIC NOT NULL)`;
      await sql`CREATE INDEX IF NOT EXISTS locks_key ON locks (key)`;
      tableSetupDone = true;
    }

    // Lazy delete expired locks
    await sql`DELETE FROM ${sql(LOCKS_TABLE)} WHERE expires_at <= ${now}`;

    // Check if lock exists
    const rows = await sql`SELECT * FROM ${sql(LOCKS_TABLE)} WHERE key = ${lockName}`;
    if (rows.length > 0) return false;

    // Acquire lock
    const expiresAt = now + timeout;
    await sql`INSERT INTO ${sql(LOCKS_TABLE)} ${sql({ key: lockName, created_at: now, expires_at: expiresAt })}`;
    return true;
  });

  async function release() {
    // Using the created_at timestamp when deleting insures that lock owner cannot release the lock after it has been expired (see related test).
    await sql`DELETE FROM locks WHERE key = ${lockName} AND created_at = ${now}`;
  }

  return [hasLock, release];
}

/**
 * Executes the given given function conditionally if it can acquire the lock.
 * Lock name is derived from given function.
 * This basically protects the given function from being executed multiple times at the same time, like a distributed throttling.
 */
export async function withTryLock<T = unknown>(
  fn: () => Promise<T>,
  { lockName, ...lockOptions }: { lockName?: string } & LockOptions = {},
): Promise<Nullable<T>> {
  const key = lockName ?? String(hash(fn.toString()));
  const [hasLock, release] = await tryLock(key, lockOptions);
  if (!hasLock) return undefined;

  const output = await fn();
  await release();
  return output;
}