Add 15s timeout for document saving in collaboration service

Co-authored-by: tommoor <380914+tommoor@users.noreply.github.com>

server/collaboration/PersistenceExtension.test.ts

@@ -0,0 +1,56 @@
+import { promiseTimeout } from "@shared/utils/timers";
+
+// Mock the dependencies
+jest.mock("@server/logging/Logger");
+jest.mock("@server/logging/tracing", () => ({
+  trace: () => () => {},
+}));
+jest.mock("@server/models/Document");
+jest.mock("@server/storage/database");
+jest.mock("@server/storage/redis");
+jest.mock("../commands/documentCollaborativeUpdater");
+
+describe("PersistenceExtension", () => {
+  describe("timeout functionality", () => {
+    it("should timeout after 15 seconds", async () => {
+      // Create a promise that takes longer than 15 seconds
+      const slowPromise = new Promise((resolve) => {
+        setTimeout(resolve, 20000);
+      });
+
+      // Wrap it with a 15 second timeout
+      const timedPromise = promiseTimeout(slowPromise, 15000);
+
+      // Expect it to throw a timeout error
+      await expect(timedPromise).rejects.toThrow(
+        "Operation timed out after 15000ms"
+      );
+    }, 20000);
+
+    it("should complete successfully if operation finishes in time", async () => {
+      // Create a promise that completes quickly
+      const fastPromise = new Promise((resolve) => {
+        setTimeout(() => resolve("success"), 100);
+      });
+
+      // Wrap it with a 15 second timeout
+      const timedPromise = promiseTimeout(fastPromise, 15000);
+
+      // Expect it to resolve successfully
+      await expect(timedPromise).resolves.toBe("success");
+    });
+
+    it("should use custom error message when provided", async () => {
+      // Create a promise that takes longer than timeout
+      const slowPromise = new Promise((resolve) => {
+        setTimeout(resolve, 2000);
+      });
+
+      const customMessage = "Custom timeout message";
+      const timedPromise = promiseTimeout(slowPromise, 1000, customMessage);
+
+      // Expect it to throw with custom message
+      await expect(timedPromise).rejects.toThrow(customMessage);
+    }, 5000);
+  });
+});

server/collaboration/PersistenceExtension.ts

@@ -5,6 +5,7 @@ import {
   Extension,
 } from "@hocuspocus/server";
 import * as Y from "yjs";
+import { promiseTimeout } from "@shared/utils/timers";
 import Logger from "@server/logging/Logger";
 import { trace } from "@server/logging/tracing";
 import Document from "@server/models/Document";
@@ -127,13 +128,17 @@ export default class PersistenceExtension implements Extension {
     }
 
     try {
-      await documentCollaborativeUpdater({
-        documentId,
-        ydoc: document,
-        sessionCollaboratorIds,
-        isLastConnection: clientsCount === 0,
-        clientVersion,
-      });
+      await promiseTimeout(
+        documentCollaborativeUpdater({
+          documentId,
+          ydoc: document,
+          sessionCollaboratorIds,
+          isLastConnection: clientsCount === 0,
+          clientVersion,
+        }),
+        15000,
+        `Document save operation timed out after 15s for document ${documentId}`
+      );
     } catch (err) {
       Logger.error("Unable to persist document", err, {
         documentId,

shared/utils/timers.test.ts

@@ -0,0 +1,70 @@
+import { promiseTimeout, sleep } from "./timers";
+
+describe("timers", () => {
+  describe("promiseTimeout", () => {
+    it("should timeout after specified milliseconds", async () => {
+      // Create a promise that takes longer than timeout
+      const slowPromise = sleep(2000);
+
+      // Wrap it with a 100ms timeout
+      const timedPromise = promiseTimeout(slowPromise, 100);
+
+      // Expect it to throw a timeout error
+      await expect(timedPromise).rejects.toThrow(
+        "Operation timed out after 100ms"
+      );
+    });
+
+    it("should complete successfully if operation finishes in time", async () => {
+      // Create a promise that completes quickly
+      const fastPromise = Promise.resolve("success");
+
+      // Wrap it with a 1 second timeout
+      const timedPromise = promiseTimeout(fastPromise, 1000);
+
+      // Expect it to resolve successfully
+      await expect(timedPromise).resolves.toBe("success");
+    });
+
+    it("should use custom error message when provided", async () => {
+      // Create a promise that takes longer than timeout
+      const slowPromise = sleep(2000);
+
+      const customMessage = "Custom timeout message";
+      const timedPromise = promiseTimeout(slowPromise, 100, customMessage);
+
+      // Expect it to throw with custom message
+      await expect(timedPromise).rejects.toThrow(customMessage);
+    });
+
+    it("should preserve the resolved value", async () => {
+      const result = { data: "test", value: 42 };
+      const promise = Promise.resolve(result);
+
+      const timedPromise = promiseTimeout(promise, 1000);
+
+      await expect(timedPromise).resolves.toEqual(result);
+    });
+
+    it("should preserve rejection from original promise", async () => {
+      const error = new Error("Original error");
+      const promise = Promise.reject(error);
+
+      const timedPromise = promiseTimeout(promise, 1000);
+
+      await expect(timedPromise).rejects.toThrow("Original error");
+    });
+  });
+
+  describe("sleep", () => {
+    it("should resolve after specified time", async () => {
+      const start = Date.now();
+      await sleep(100);
+      const elapsed = Date.now() - start;
+
+      // Allow some margin for timing
+      expect(elapsed).toBeGreaterThanOrEqual(90);
+      expect(elapsed).toBeLessThan(200);
+    });
+  });
+});

shared/utils/timers.ts

@@ -6,3 +6,25 @@
 export function sleep(ms = 1) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
+
+/**
+ * Wraps a promise with a timeout. If the promise doesn't resolve within the
+ * specified time, it will be rejected with a timeout error.
+ *
+ * @param promise The promise to wrap with a timeout.
+ * @param timeoutMs The timeout duration in milliseconds.
+ * @param errorMessage Optional custom error message for the timeout.
+ * @returns A promise that resolves with the original promise's value or rejects on timeout.
+ */
+export function promiseTimeout<T>(
+  promise: Promise<T>,
+  timeoutMs: number,
+  errorMessage = `Operation timed out after ${timeoutMs}ms`
+): Promise<T> {
+  return Promise.race([
+    promise,
+    new Promise<T>((_, reject) =>
+      setTimeout(() => reject(new Error(errorMessage)), timeoutMs)
+    ),
+  ]);
+}