feat: add withDatastoreLock() helper (#5057)

Makes it easier to acquire and release a datastore lock.

---------

Co-authored-by: Jeff Ching <chingor@google.com>

packages/datastore-lock/src/datastore-lock.ts

@@ -207,3 +207,60 @@ export class DatastoreLock {
     else return false;
   }
 }
+
+/**
+ * Parameters for creating a DatastoreLock.
+ */
+export interface DatastoreLockDetails {
+  /** A unique identifier for the lock.
+   * This identifier is used to identify the lock in the datastore. */
+  lockId: string;
+  /** The name of the resource that the lock is protecting.
+   * Usually a url of a github resource. */
+  target: string;
+  /** If the current process crashes or otherwise fails to release the lock,
+   * it will be automatically released in lockExpiry milliseconds. */
+  lockExpiry?: number;
+  /** Milliseconds to wait while trying to acquire the lock. */
+  lockAcquireTimeout?: number;
+}
+
+/**
+ * Construct a new DatastoreLock.
+ */
+export function datastoreLockFromDetails(
+  details: DatastoreLockDetails
+): DatastoreLock {
+  return new DatastoreLock(
+    details.lockId,
+    details.target,
+    details.lockExpiry,
+    details.lockAcquireTimeout
+  );
+}
+
+/**
+ * Execute a function while holding a lock.
+ *
+ * @param details The Datastore lock's details.
+ * @param f The function to execute while holding the lock.
+ * @returns the value returned by f().
+ */
+export async function withDatastoreLock<R>(
+  details: DatastoreLockDetails,
+  f: () => Promise<R>
+): Promise<R> {
+  const lock = datastoreLockFromDetails(details);
+  const acquired = await lock.acquire();
+  if (!acquired) {
+    // throw an error and expect gcf-utils infrastructure to retry
+    throw new Error(
+      `Failed to acquire lock in ${details.lockAcquireTimeout}ms for ${details.target}.`
+    );
+  }
+  try {
+    return await f();
+  } finally {
+    lock.release();
+  }
+}

packages/datastore-lock/test/datastore-lock.ts

@@ -17,7 +17,12 @@ import {describe, it, before, after} from 'mocha';
 import nock from 'nock';
 
 import DataStoreEmulator from 'google-datastore-emulator';
-import {DatastoreLock} from '../src/datastore-lock';
+import {
+  DatastoreLock,
+  DatastoreLockDetails,
+  datastoreLockFromDetails,
+  withDatastoreLock,
+} from '../src/datastore-lock';
 
 describe('datastore-lock', () => {
   let emulator: DataStoreEmulator;
@@ -90,4 +95,37 @@ describe('datastore-lock', () => {
       assert.strictEqual(await l2.peek(), false);
     });
   });
+
+  describe('withDatastoreLock', () => {
+    const details: DatastoreLockDetails = {
+      lockId: 'withDatastoreLock',
+      target: 'test',
+    };
+
+    it('works', async () => {
+      await withDatastoreLock(details, async () => {
+        const l = datastoreLockFromDetails(details);
+        assert(await l.peek());
+      });
+    });
+
+    it('times out', async () => {
+      const details: DatastoreLockDetails = {
+        lockId: 'withDatastoreLock',
+        target: 'test',
+      };
+      try {
+        await withDatastoreLock(details, async () => {
+          await withDatastoreLock(
+            {...details, lockAcquireTimeout: 100},
+            async () => {
+              throw new Error('This code should never execute.');
+            }
+          );
+        });
+      } catch (e) {
+        assert(String(e).includes('100ms'));
+      }
+    });
+  });
 });