如何使用 Firebase Firestore 獲取悲觀鎖？

Question

我有一些 firebase function 執行以下操作

1. 接受一些記錄的部分有效載荷作為輸入
2. 從第三方讀取該記錄
3. 使用部分有效載荷記錄並將其存儲在第三方的補丁

第三方有自己的數據庫，不支持任何類型的部分更新/鎖定。 用戶可能會同時調用此端點，因此我需要實施悲觀鎖/FIFO 隊列以確保 function 永遠不會針對該特定記錄同時運行。 否則，一個用戶的更新可能會破壞另一個用戶。 請注意，這不是性能問題，因為並發調用不太可能超過 10 個。

以下是我想象的如何通過添加步驟 0 和 4 以高級術語實現此鎖定：

0：獲取具有特定 ID 的鎖，如果該鎖已被獲取，則繼續輪詢直到它空閑或鎖過期

...

4：現在記錄補丁已經完成釋放鎖

我如何使用 Firebase Firstore 實現這樣的鎖定機制？

Answer 1

我找到了一個方法，你可以使用 firebase 事務來檢查一行是否存在，如果不存在則可以獲取鎖：

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