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

Question

I have some firebase function which does the following
Accepts a partial payload of some record as input
Reads that record from a third party
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