perf: Add cache for document structure (#9196)

* Normalize Collection.findByPk

* Add caching of documentStructure

* fix: Do not set cache before transaction is flushed

* Mock Redis

server/models/Collection.ts

@@ -14,7 +14,8 @@ import {
   EmptyResultError,
   type CreateOptions,
   type UpdateOptions,
-  ScopeOptions,
+  type ScopeOptions,
+  type SaveOptions,
 } from "sequelize";
 import {
   Sequelize,
@@ -39,6 +40,7 @@ import {
   BeforeCreate,
   BeforeUpdate,
   DefaultScope,
+  AfterSave,
 } from "sequelize-typescript";
 import isUUID from "validator/lib/isUUID";
 import type { CollectionSort, ProsemirrorData } from "@shared/types";
@@ -48,6 +50,7 @@ import { sortNavigationNodes } from "@shared/utils/collections";
 import slugify from "@shared/utils/slugify";
 import { CollectionValidation } from "@shared/validations";
 import { ValidationError } from "@server/errors";
+import { CacheHelper } from "@server/utils/CacheHelper";
 import removeIndexCollision from "@server/utils/removeIndexCollision";
 import { generateUrlId } from "@server/utils/url";
 import { ValidateIndex } from "@server/validation";
@@ -334,6 +337,34 @@ class Collection extends ParanoidModel<
     if (!model.content) {
       model.content = await DocumentHelper.toJSON(model);
     }
+    if (model.changed("documentStructure")) {
+      await CacheHelper.clearData(
+        CacheHelper.getCollectionDocumentsKey(model.id)
+      );
+    }
+  }
+
+  @AfterSave
+  static async cacheDocumentStructure(
+    model: Collection,
+    options: SaveOptions<Collection>
+  ) {
+    if (model.changed("documentStructure")) {
+      const setData = () =>
+        CacheHelper.setData(
+          CacheHelper.getCollectionDocumentsKey(model.id),
+          model.documentStructure,
+          60
+        );
+
+      if (options.transaction) {
+        return (options.transaction.parent || options.transaction).afterCommit(
+          setData
+        );
+      }
+
+      await setData();
+    }
   }
 
   @BeforeDestroy

server/routes/api/collections/collections.ts

@@ -38,6 +38,7 @@ import {
   presentFileOperation,
 } from "@server/presenters";
 import { APIContext } from "@server/types";
+import { CacheHelper } from "@server/utils/CacheHelper";
 import { RateLimiterStrategy } from "@server/utils/RateLimiter";
 import { collectionIndexing } from "@server/utils/indexing";
 import pagination from "../middlewares/pagination";
@@ -143,13 +144,25 @@ router.post(
     const { user } = ctx.state.auth;
     const collection = await Collection.findByPk(id, {
       userId: user.id,
-      includeDocumentStructure: true,
     });
 
     authorize(user, "readDocument", collection);
 
+    const documentStructure = await CacheHelper.getDataOrSet(
+      CacheHelper.getCollectionDocumentsKey(collection.id),
+      async () =>
+        (
+          await Collection.findByPk(collection.id, {
+            attributes: ["documentStructure"],
+            includeDocumentStructure: true,
+            rejectOnEmpty: true,
+          })
+        ).documentStructure,
+      60
+    );
+
     ctx.body = {
-      data: collection.documentStructure || [],
+      data: documentStructure || [],
     };
   }
 );

server/storage/__mocks__/redis.ts

@@ -0,0 +1,47 @@
+import { EventEmitter } from "events";
+
+// Create a mock Redis client with all needed methods mocked
+class RedisMock extends EventEmitter {
+  constructor() {
+    super();
+  }
+
+  get = jest.fn().mockResolvedValue(null);
+  set = jest.fn().mockResolvedValue("OK");
+  del = jest.fn().mockResolvedValue(1);
+  keys = jest.fn().mockResolvedValue([]);
+  ping = jest.fn().mockResolvedValue("PONG");
+  disconnect = jest.fn();
+  setMaxListeners = jest.fn();
+}
+
+// Mock the RedisAdapter class
+class RedisAdapter extends RedisMock {
+  constructor(_url: string | undefined, _options = {}) {
+    super();
+  }
+
+  private static client: RedisAdapter;
+  private static subscriber: RedisAdapter;
+
+  public static get defaultClient(): RedisAdapter {
+    return (
+      this.client ||
+      (this.client = new this(undefined, {
+        connectionNameSuffix: "client",
+      }))
+    );
+  }
+
+  public static get defaultSubscriber(): RedisAdapter {
+    return (
+      this.subscriber ||
+      (this.subscriber = new this(undefined, {
+        maxRetriesPerRequest: null,
+        connectionNameSuffix: "subscriber",
+      }))
+    );
+  }
+}
+
+export default RedisAdapter;

server/test/setup.ts

@@ -7,6 +7,11 @@ require("@server/storage/database");
 
 jest.mock("bull");
 
+// Enable mocks for Redis-related modules
+jest.mock("@server/storage/redis");
+jest.mock("@server/utils/MutexLock");
+jest.mock("@server/utils/CacheHelper");
+
 // This is needed for the relative manual mock to be picked up
 jest.mock("../queues");
 
@@ -34,7 +39,9 @@ jest.mock("@aws-sdk/s3-request-presigner", () => ({
   getSignedUrl: jest.fn(),
 }));
 
-afterAll(() => Redis.defaultClient.disconnect());
+afterAll(() => {
+  Redis.defaultClient.disconnect();
+});
 
 beforeEach(() => {
   env.URL = sharedEnv.URL = "https://app.outline.dev";

server/utils/CacheHelper.ts

@@ -125,4 +125,8 @@ export class CacheHelper {
   public static getUnfurlKey(teamId: string, url = "") {
     return `unfurl:${teamId}:${url}`;
   }
+
+  public static getCollectionDocumentsKey(collectionId: string) {
+    return `cd:${collectionId}`;
+  }
 }

server/utils/__mocks__/CacheHelper.ts

@@ -0,0 +1,54 @@
+import { Day } from "@shared/utils/time";
+
+/**
+ * A Mock Helper class for server-side cache management
+ */
+export class CacheHelper {
+  // Default expiry time for cache data in seconds
+  private static defaultDataExpiry = Day.seconds;
+
+  /**
+   * Mocked method that resolves with the callback result
+   */
+  public static async getDataOrSet<T>(
+    key: string,
+    callback: () => Promise<T | undefined>,
+    _expiry: number,
+    _lockTimeout: number
+  ): Promise<T | undefined> {
+    return await callback();
+  }
+
+  /**
+   * Mocked method that resolves with undefined
+   */
+  public static async getData<T>(_key: string): Promise<T | undefined> {
+    return undefined;
+  }
+
+  /**
+   * Mocked method that resolves with void
+   */
+  public static async setData<T>(_key: string, _data: T, _expiry?: number) {
+    return;
+  }
+
+  /**
+   * Mocked method that resolves with void
+   */
+  public static async clearData(_prefix: string) {
+    return;
+  }
+
+  /**
+   * These are real methods that don't require mocking as they don't
+   * interact with Redis directly
+   */
+  public static getUnfurlKey(teamId: string, url = "") {
+    return `unfurl:${teamId}:${url}`;
+  }
+
+  public static getCollectionDocumentsKey(collectionId: string) {
+    return `cd:${collectionId}`;
+  }
+}

server/utils/__mocks__/MutexLock.ts

@@ -0,0 +1,18 @@
+export class MutexLock {
+  // Default expiry time for acquiring lock in milliseconds
+  public static defaultLockTimeout = 4000;
+
+  /**
+   * Returns the mock redlock instance
+   */
+  public static get lock() {
+    return {
+      acquire: jest.fn().mockResolvedValue({
+        release: jest.fn().mockResolvedValue(true),
+        expiration: Date.now() + 10000,
+      }),
+    };
+  }
+
+  private static redlock: any;
+}

yarn.lock

@@ -4136,18 +4136,7 @@
     "@smithy/util-utf8" "^4.0.0"
     tslib "^2.6.2"
 
-"@smithy/credential-provider-imds@^4.0.2":
-  version "4.0.2"
-  resolved "https://registry.yarnpkg.com/@smithy/credential-provider-imds/-/credential-provider-imds-4.0.2.tgz#1ec34a04842fa69996b151a695b027f0486c69a8"
-  integrity sha512-32lVig6jCaWBHnY+OEQ6e6Vnt5vDHaLiydGrwYMW9tPqO688hPGTYRamYJ1EptxEC2rAwJrHWmPoKRBl4iTa8w==
-  dependencies:
-    "@smithy/node-config-provider" "^4.0.2"
-    "@smithy/property-provider" "^4.0.2"
-    "@smithy/types" "^4.2.0"
-    "@smithy/url-parser" "^4.0.2"
-    tslib "^2.6.2"
-
-"@smithy/credential-provider-imds@^4.0.4":
+"@smithy/credential-provider-imds@^4.0.2", "@smithy/credential-provider-imds@^4.0.4":
   version "4.0.4"
   resolved "https://registry.yarnpkg.com/@smithy/credential-provider-imds/-/credential-provider-imds-4.0.4.tgz#01315ab90c4cb3e017c1ee2c6e5f958aeaa7cf78"
   integrity sha512-jN6M6zaGVyB8FmNGG+xOPQB4N89M1x97MMdMnm1ESjljLS3Qju/IegQizKujaNcy2vXAvrz0en8bobe6E55FEA==