How can I use Firebase Firestore to acquire a pessimistic lock?

Question

I have some firebase function which does the following

1. Accepts a partial payload of some record as input
2. Reads that record from a third party
3. Patches that record with the partial payload and stores it on the third party

The third party has their own database and doesn't support any kind of partial updates / locking. Users are likely to be calling this endpoint concurrently so I need to implement a pessimistic lock / FIFO queue to ensure the function never runs concurrently for that specific record. Otherwise, one user's update could clobber another. Note, this is not a performance concern because it's very unlikely to have more than 10 concurrent calls.

Here is how I imagine I could implement this locking in high level terms by adding step 0 and 4:

0: Acquire a lock with a specific ID, if that lock is already acquired keep polling until it's free or the lock expires

...

4: Release the lock now that the record patch has completed

How could I implement such a locking mechanism with Firebase Firstore?

Answer 1

I found a way, you can use firebase transactions to check if a row exists or not, if it does not exists then you can acquire the lock:

import { getFirestore } from 'firebase-admin/firestore'

export type Lock = {
  lockId: string
  release: () => Promise<void>
}

export const Locks = {
  /** Use firestore to acquire a lock, by writing the expiration time to the `locks/${lockId}` document.
   * If the document already exists and the expiration time is in the future, then exit without acquiring the lock.
   * If the document does not exist or the expiration time is in the past, then the acquire the lock and return a Lock object. */
  async acquire(lockId: string, { lockDurationMs = 1000 * 60 } = {}): Promise<Lock | undefined> {
    const db = getFirestore()
    const docRef = db.collection('locks').doc(lockId)
    const maybeLock = await db.runTransaction(async transaction => {
      const doc = await transaction.get(docRef)
      const expiration = doc.data()?.expiration || 0
      if (expiration > Date.now()) {
        // Someone else has the lock
        return
      }
      const newExpiration = Date.now() + lockDurationMs
      await transaction.set(docRef, { expiration: newExpiration }, { merge: true })
      const lock: Lock = {
        lockId,
        release: async () => {
          await docRef.set({ expiration: null }, { merge: true })
        },
      }
      return lock
    })
    return maybeLock
  },
  /** Poll for a lock to be acquired, and return it when it is.
   * Will return undefined if the lock is not acquired within the timeout. */
  async pollUntilAcquired(
    lockId: string,
    { lockDurationMs = 1000 * 60, pollIntervalMs = 1000, timeoutMs = 1000 * 60 * 2 } = {}
  ) {
    const lock = await pollFor(() => this.acquire(lockId, { lockDurationMs }), {
      interval: pollIntervalMs,
      // If we don't get the lock within 2 minutes, then something is wrong
      timeout: timeoutMs,
    })
    return lock
  },
}

// #region Tests
if (process.env.TESTING) {
  async function testLock() {
    const { initFirebaseAdmin } = require('./initFirebaseAdmin')
    initFirebaseAdmin('staging')

    const db = getFirestore()
    const locks = db.collection('locks')

    async function getTestLock() {
      const doc = await locks.doc('testLock').get()
      const data = doc.data()
      return data?.expiration || 0
    }

    {
      // TEST 1: Acquire a lock, check it inserts the row, then release it and check it deletes the row
      const lock = await Locks.acquire('testLock', { lockDurationMs: 1000 })
      const expiration = await getTestLock()
      assert(expiration > Date.now(), 'Lock was not acquired')
      await lock?.release()
      const expiration2 = await getTestLock()
      assert(expiration2 === 0, 'Lock was not released')
    }

    {
      // TEST 2: Acquire a lock, then try to acquire it again and check it fails
      const lock = await Locks.acquire('testLock', { lockDurationMs: 1000 })
      const lock2 = await Locks.acquire('testLock', { lockDurationMs: 1000 })
      assert(!lock2, 'Lock was acquired twice')
      await lock?.release()
    }

    {
      // TEST 3: Acquire a lock, then try to acquire it again with polling and check it succeeds
      const lock = await Locks.acquire('testLock', { lockDurationMs: 1000 })
      const lock2 = await Locks.pollUntilAcquired('testLock', { lockDurationMs: 1000 })
      assert(lock2, 'Lock was not acquired after polling, but it should have succeeded')
      await lock?.release()
      await lock2?.release()
    }
  }
  void testLock().catch(err => {
    console.error(err)
    process.exit(1)
  })
}

// #endregion

// #region Utils: you can skip this region of code, it's only included so that the snippet is self-contained

function sleep(ms: number) {
  return new Promise(resolve => setTimeout(resolve, ms))
}

function isUndefined(x: unknown): x is undefined {
  return typeof x === 'undefined'
}

/**
 * If neither timeout nor attempts is provided, defaults to 30
 * attempts.
 * If only timeout is provided, attempts will be infinite.
 * If only attempts is provided, timeout will be infinite.
 * If both are provided, both will be used to limit the poll.
 */
const pollFor = async <T>(
  fn: () => Promise<T | undefined> | T | undefined,
  opts?: {
    /** Defaults to 0 - in milliseconds */
    interval?: number
    /** In milliseconds  */
    timeout?: number
    attempts?: number
  }
) => {
  let { interval = 0, timeout, attempts } = opts || {} // eslint-disable-line
  if (!isUndefined(timeout) && isUndefined(attempts)) attempts = Infinity
  attempts = isUndefined(attempts) ? 30 : attempts
  timeout = isUndefined(timeout) ? Infinity : timeout

  const start = Date.now()
  for (let i = 1; i < attempts + 1; i++) {
    const result = await fn()
    if (result !== undefined) return result
    if (i > attempts) return
    if (Date.now() - start > timeout) return
    await sleep(interval)
  }

  return
}

function assert(
  predicate: any,
  message: string,
  extra?: Record<string, unknown>
): asserts predicate {
  if (!predicate) {
    throw new AssertionError(message, extra)
  }
}

class AssertionError extends Error {
  extra: Record<string, unknown> | undefined
  constructor(message: string, extra: Record<string, unknown> | undefined = undefined) {
    super(message)
    this.name = 'AssertionError'
    this.extra = extra
  }
  toJSON() {
    return {
      message: this.message,
      name: this.name,
      extra: this.extra,
    }
  }
}

// #endregion