fix: cache data stores to fix potential race condition

Multiple concurrent requests for the same user pool would create multiple data stores backed by the same file, and on write they would stomp on each other.

integration-tests/aws-sdk/setup.ts

@@ -13,6 +13,7 @@ import {
   TriggersService,
 } from "../../src/services";
 import { CognitoServiceFactoryImpl } from "../../src/services/cognitoService";
+import { NoOpCache } from "../../src/services/dataStore/cache";
 import { StormDBDataStoreFactory } from "../../src/services/dataStore/stormDb";
 import { otp } from "../../src/services/otp";
 import { UserPoolServiceFactoryImpl } from "../../src/services/userPoolService";
@@ -38,7 +39,10 @@ export const withCognitoSdk =
       dataDirectory = await mkdtemp("/tmp/cognito-local:");
       const ctx = { logger };
 
-      const dataStoreFactory = new StormDBDataStoreFactory(dataDirectory);
+      const dataStoreFactory = new StormDBDataStoreFactory(
+        dataDirectory,
+        new NoOpCache()
+      );
       const cognitoServiceFactory = new CognitoServiceFactoryImpl(
         dataDirectory,
         clock,
@@ -66,7 +70,10 @@ export const withCognitoSdk =
         port: 0,
       });
 
-      const address = httpServer.address()!;
+      const address = httpServer.address();
+      if (!address) {
+        throw new Error("HttpServer has no address");
+      }
       const url =
         typeof address === "string"
           ? address

integration-tests/cognitoService.test.ts

@@ -7,6 +7,7 @@ import {
 } from "../src/services/cognitoService";
 import fs from "fs";
 import { promisify } from "util";
+import { NoOpCache } from "../src/services/dataStore/cache";
 import { StormDBDataStoreFactory } from "../src/services/dataStore/stormDb";
 import { UserPoolServiceFactoryImpl } from "../src/services/userPoolService";
 
@@ -21,7 +22,10 @@ describe("Cognito Service", () => {
     dataDirectory = await mkdtemp("/tmp/cognito-local:");
 
     const clock = new DateClock();
-    const dataStoreFactory = new StormDBDataStoreFactory(dataDirectory);
+    const dataStoreFactory = new StormDBDataStoreFactory(
+      dataDirectory,
+      new NoOpCache()
+    );
 
     factory = new CognitoServiceFactoryImpl(
       dataDirectory,

integration-tests/dataStore.test.ts

@@ -1,7 +1,8 @@
-import StormDB from "stormdb";
-import { TestContext } from "../src/__tests__/testContext";
 import fs from "fs";
+import StormDB from "stormdb";
 import { promisify } from "util";
+import { TestContext } from "../src/__tests__/testContext";
+import { InMemoryCache, NoOpCache } from "../src/services/dataStore/cache";
 import { DataStoreFactory } from "../src/services/dataStore/factory";
 import { StormDBDataStoreFactory } from "../src/services/dataStore/stormDb";
 
@@ -15,7 +16,7 @@ describe("Data Store", () => {
 
   beforeEach(async () => {
     path = await mkdtemp("/tmp/cognito-local:");
-    factory = new StormDBDataStoreFactory(path);
+    factory = new StormDBDataStoreFactory(path, new NoOpCache());
   });
 
   afterEach(() =>
@@ -303,5 +304,29 @@ describe("Data Store", () => {
 
       expect(result).toEqual(date);
     });
+
+    it("supports caching data stores across creates for the same id", async () => {
+      const factory = new StormDBDataStoreFactory(path, new InMemoryCache());
+
+      const dataStore1 = await factory.create(TestContext, "example", {});
+      const dataStore2 = await factory.create(TestContext, "example", {});
+      const dataStore3 = await factory.create(TestContext, "example2", {});
+
+      await dataStore1.set(TestContext, "One", 1);
+      await dataStore2.set(TestContext, "Two", 2);
+      await dataStore3.set(TestContext, "Three", 3);
+
+      expect(await dataStore1.getRoot(TestContext)).toEqual({
+        One: 1,
+        Two: 2,
+      });
+      expect(await dataStore2.getRoot(TestContext)).toEqual({
+        One: 1,
+        Two: 2,
+      });
+      expect(await dataStore3.getRoot(TestContext)).toEqual({
+        Three: 3,
+      });
+    });
   });
 });

integration-tests/userPoolService.test.ts

@@ -1,14 +1,9 @@
 import fs from "fs";
 import { promisify } from "util";
 import { TestContext } from "../src/__tests__/testContext";
-import {
-  CognitoService,
-  CognitoServiceImpl,
-  DateClock,
-  UserPoolService,
-  UserPoolServiceImpl,
-} from "../src/services";
+import { CognitoService, DateClock, UserPoolService } from "../src/services";
 import { CognitoServiceFactoryImpl } from "../src/services/cognitoService";
+import { NoOpCache } from "../src/services/dataStore/cache";
 import { StormDBDataStoreFactory } from "../src/services/dataStore/stormDb";
 import { UserPoolServiceFactoryImpl } from "../src/services/userPoolService";
 
@@ -25,7 +20,10 @@ describe("User Pool Service", () => {
   beforeEach(async () => {
     dataDirectory = await mkdtemp("/tmp/cognito-local:");
     const clock = new DateClock();
-    const dataStoreFactory = new StormDBDataStoreFactory(dataDirectory);
+    const dataStoreFactory = new StormDBDataStoreFactory(
+      dataDirectory,
+      new NoOpCache()
+    );
 
     cognitoClient = await new CognitoServiceFactoryImpl(
       dataDirectory,

src/server/defaults.ts

@@ -7,6 +7,7 @@ import {
   TriggersService,
 } from "../services";
 import { CognitoServiceFactoryImpl } from "../services/cognitoService";
+import { InMemoryCache } from "../services/dataStore/cache";
 import { StormDBDataStoreFactory } from "../services/dataStore/stormDb";
 import { ConsoleMessageSender } from "../services/messageDelivery/consoleMessageSender";
 import { MessageDeliveryService } from "../services/messageDelivery/messageDelivery";
@@ -28,14 +29,17 @@ export const createDefaultServer = async (
   const config = await loadConfig(
     ctx,
     // the config gets a separate factory because it's stored in a different directory
-    new StormDBDataStoreFactory(configDirectory)
+    new StormDBDataStoreFactory(configDirectory, new InMemoryCache())
   );
 
   logger.debug({ config }, "Loaded config");
 
   const clock = new DateClock();
 
-  const dataStoreFactory = new StormDBDataStoreFactory(dataDirectory);
+  const dataStoreFactory = new StormDBDataStoreFactory(
+    dataDirectory,
+    new InMemoryCache()
+  );
 
   const cognitoServiceFactory = new CognitoServiceFactoryImpl(
     dataDirectory,

src/services/dataStore/cache.ts

@@ -0,0 +1,27 @@
+import { DataStore } from "./dataStore";
+
+export type DataStoreCache = {
+  get(key: string): DataStore | null;
+  set(key: string, value: DataStore): void;
+};
+
+export class InMemoryCache implements DataStoreCache {
+  private readonly cache: Record<string, DataStore> = {};
+
+  get(key: string): DataStore | null {
+    return this.cache[key];
+  }
+
+  set(key: string, value: DataStore): void {
+    this.cache[key] = value;
+  }
+}
+
+export class NoOpCache implements DataStoreCache {
+  get(): DataStore | null {
+    return null;
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-empty-function
+  set(): void {}
+}

src/services/dataStore/stormDb.ts

@@ -2,6 +2,7 @@ import fs from "fs";
 import StormDB from "stormdb";
 import { promisify } from "util";
 import { Context } from "../context";
+import { DataStoreCache } from "./cache";
 import { DataStore } from "./dataStore";
 import { DataStoreFactory } from "./factory";
 
@@ -98,9 +99,11 @@ const createStormDBInstance = (directory: string, id: string): StormDB => {
 
 export class StormDBDataStoreFactory implements DataStoreFactory {
   private readonly directory: string;
+  private readonly cache: DataStoreCache;
 
-  public constructor(directory: string) {
+  public constructor(directory: string, dataStoreCache: DataStoreCache) {
     this.directory = directory;
+    this.cache = dataStoreCache;
   }
 
   public async create(
@@ -111,6 +114,12 @@ export class StormDBDataStoreFactory implements DataStoreFactory {
     ctx.logger.debug({ id }, "createDataStore");
     await mkdir(this.directory, { recursive: true });
 
+    const cachedDb = this.cache.get(id);
+    if (cachedDb) {
+      ctx.logger.debug({ id }, "Using cached data store");
+      return cachedDb;
+    }
+
     ctx.logger.debug({ id }, "Creating new data store");
     const db = createStormDBInstance(this.directory, id);
 
@@ -119,6 +128,10 @@ export class StormDBDataStoreFactory implements DataStoreFactory {
     ctx.logger.debug({ store: db.value() }, "DataStore.save");
     await db.save();
 
-    return new StormDBDataStore(db);
+    const dataStore = new StormDBDataStore(db);
+
+    this.cache.set(id, dataStore);
+
+    return dataStore;
   }
 }