From fb72b77ac19491a9f06994014eed0c9145661a70 Mon Sep 17 00:00:00 2001
From: Adria Navarro <adria@budibase.com>
Date: Thu, 30 Nov 2023 18:34:56 +0100
Subject: [PATCH] Use jest.useFakeTimers

---
 .../backend-core/src/redis/redlockImpl.ts     | 24 ++++------------
 .../src/redis/tests/redlockImpl.spec.ts       | 28 +++++++++++--------
 packages/types/src/sdk/locks.ts               |  1 +
 3 files changed, 24 insertions(+), 29 deletions(-)

diff --git a/packages/backend-core/src/redis/redlockImpl.ts b/packages/backend-core/src/redis/redlockImpl.ts
index 36504ea916..a65dcd0c3f 100644
--- a/packages/backend-core/src/redis/redlockImpl.ts
+++ b/packages/backend-core/src/redis/redlockImpl.ts
@@ -101,20 +101,20 @@ function getLockName(opts: LockOptions) {
   return name
 }
 
+export const autoExtendPollingMs = Duration.fromSeconds(10).toMs()
+
 export async function doWithLock<T>(
   opts: LockOptions,
   task: () => Promise<T>
 ): Promise<RedlockExecution<T>> {
   const redlock = await getClient(opts.type, opts.customOptions)
   let lock: Redlock.Lock | undefined
-  let timeout
+  let timeout: NodeJS.Timeout | undefined
   try {
     const name = getLockName(opts)
 
     const ttl =
-      opts.type === LockType.AUTO_EXTEND
-        ? Duration.fromSeconds(10).toMs()
-        : opts.ttl
+      opts.type === LockType.AUTO_EXTEND ? autoExtendPollingMs : opts.ttl
 
     // create the lock
     lock = await redlock.lock(name, ttl)
@@ -123,21 +123,9 @@ export async function doWithLock<T>(
       // We keep extending the lock while the task is running
       const extendInIntervals = (): void => {
         timeout = setTimeout(async () => {
-          let isExpired = false
-          try {
-            lock = await lock!.extend(ttl)
-          } catch (err: any) {
-            isExpired = err.message.includes("Cannot extend lock on resource")
-            if (isExpired) {
-              console.error("The lock expired", { name })
-            } else {
-              throw err
-            }
-          }
+          lock = await lock!.extend(ttl, () => opts.onExtend && opts.onExtend())
 
-          if (!isExpired) {
-            extendInIntervals()
-          }
+          extendInIntervals()
         }, ttl / 2)
       }
 
diff --git a/packages/backend-core/src/redis/tests/redlockImpl.spec.ts b/packages/backend-core/src/redis/tests/redlockImpl.spec.ts
index 9d312721b1..6f894c2951 100644
--- a/packages/backend-core/src/redis/tests/redlockImpl.spec.ts
+++ b/packages/backend-core/src/redis/tests/redlockImpl.spec.ts
@@ -1,14 +1,15 @@
 import { LockName, LockType, LockOptions } from "@budibase/types"
-import tk from "timekeeper"
-import { doWithLock } from "../redlockImpl"
+import { autoExtendPollingMs, doWithLock } from "../redlockImpl"
 import { DBTestConfiguration, generator } from "../../../tests"
 
-tk.reset()
-
 describe("redlockImpl", () => {
+  beforeEach(() => {
+    jest.useFakeTimers()
+  })
+
   describe("doWithLock", () => {
     const config = new DBTestConfiguration()
-    const lockTtl = 50
+    const lockTtl = autoExtendPollingMs
 
     function runLockWithExecutionTime({
       opts,
@@ -21,7 +22,10 @@ describe("redlockImpl", () => {
     }) {
       return config.doInTenant(() =>
         doWithLock(opts, async () => {
-          await new Promise<void>(r => setTimeout(() => r(), executionTimeMs))
+          const interval = lockTtl / 10
+          for (let i = executionTimeMs; i > 0; i -= interval) {
+            await jest.advanceTimersByTimeAsync(interval)
+          }
           return task()
         })
       )
@@ -33,7 +37,7 @@ describe("redlockImpl", () => {
         const expectedResult = generator.guid()
         const mockTask = jest.fn().mockResolvedValue(expectedResult)
 
-        const opts = {
+        const opts: LockOptions = {
           name: LockName.PERSIST_WRITETHROUGH,
           type: lockType,
           ttl: lockTtl,
@@ -54,22 +58,24 @@ describe("redlockImpl", () => {
     it("should extend when type is autoextend", async () => {
       const expectedResult = generator.guid()
       const mockTask = jest.fn().mockResolvedValue(expectedResult)
+      const mockOnExtend = jest.fn()
 
-      const opts = {
+      const opts: LockOptions = {
         name: LockName.PERSIST_WRITETHROUGH,
         type: LockType.AUTO_EXTEND,
-        ttl: lockTtl,
+        onExtend: mockOnExtend,
       }
 
       const result = await runLockWithExecutionTime({
         opts,
         task: mockTask,
-        executionTimeMs: lockTtl * 3,
+        executionTimeMs: lockTtl * 2.5,
       })
 
       expect(result.executed).toBe(true)
       expect(result.executed && result.result).toBe(expectedResult)
       expect(mockTask).toHaveBeenCalledTimes(1)
+      expect(mockOnExtend).toHaveBeenCalledTimes(5)
     })
 
     it.each(Object.values(LockType).filter(t => t !== LockType.AUTO_EXTEND))(
@@ -77,7 +83,7 @@ describe("redlockImpl", () => {
       async (lockType: LockType) => {
         const mockTask = jest.fn().mockResolvedValue("mockResult")
 
-        const opts = {
+        const opts: LockOptions = {
           name: LockName.PERSIST_WRITETHROUGH,
           type: lockType,
           ttl: lockTtl,
diff --git a/packages/types/src/sdk/locks.ts b/packages/types/src/sdk/locks.ts
index b9470402c1..6ff91d4315 100644
--- a/packages/types/src/sdk/locks.ts
+++ b/packages/types/src/sdk/locks.ts
@@ -54,5 +54,6 @@ export type LockOptions = {
     }
   | {
       type: LockType.AUTO_EXTEND
+      onExtend?: () => void
     }
 )