Throttle persisting to disk

package.json

@@ -32,6 +32,7 @@
     "dedent": "^1.5.3",
     "dotenv": "^16.4.7",
     "execa": "^9.5.2",
+    "exit-hook": "^4.0.0",
     "hono": "^4.5.4",
     "langsmith": "^0.2.15",
     "open": "^10.1.0",

pnpm-lock.yaml

@@ -47,6 +47,9 @@ importers:
       execa:
         specifier: ^9.5.2
         version: 9.5.2
+      exit-hook:
+        specifier: ^4.0.0
+        version: 4.0.0
       hono:
         specifier: ^4.5.4
         version: 4.6.16
@@ -730,6 +733,10 @@ packages:
     resolution: {integrity: sha512-EHlpxMCpHWSAh1dgS6bVeoLAXGnJNdR93aabr4QCGbzOM73o5XmRfM/e5FUqsw3aagP8S8XEWUWFAxnRBnAF0Q==}
     engines: {node: ^18.19.0 || >=20.5.0}
 
+  exit-hook@4.0.0:
+    resolution: {integrity: sha512-Fqs7ChZm72y40wKjOFXBKg7nJZvQJmewP5/7LtePDdnah/+FH9Hp5sgMujSCMPXlxOAW2//1jrW9pnsY7o20vQ==}
+    engines: {node: '>=18'}
+
   expect-type@1.1.0:
     resolution: {integrity: sha512-bFi65yM+xZgk+u/KRIpekdSYkTB5W1pEf0Lt8Q8Msh7b+eQ7LXVtIB1Bkm4fvclDEL1b2CZkMhv2mOeF8tMdkA==}
     engines: {node: '>=12.0.0'}
@@ -1634,6 +1641,8 @@ snapshots:
       strip-final-newline: 4.0.0
       yoctocolors: 2.1.1
 
+  exit-hook@4.0.0: {}
+
   expect-type@1.1.0: {}
 
   fecha@4.2.3: {}

src/cli/dev.entrypoint.mts

@@ -1,11 +1,14 @@
 import "../preload.mjs";
 
+import { asyncExitHook } from "exit-hook";
 import * as process from "node:process";
 import { startServer, StartServerSchema } from "../server.mjs";
 import { connectToServer } from "./utils/ipc/client.mjs";
 import { Client as LangSmithClient } from "langsmith";
 import { logger } from "../logging.mjs";
 
+logger.info(`Starting server...`);
+
 const [ppid, payload] = process.argv.slice(-2);
 const sendToParent = await connectToServer(+ppid);
 
@@ -19,7 +22,7 @@ const isTracingEnabled = () => {
   return value === "true";
 };
 
-const [host, organizationId] = await Promise.all([
+const [{ host, cleanup }, organizationId] = await Promise.all([
   startServer(StartServerSchema.parse(JSON.parse(payload))),
   (async () => {
     if (isTracingEnabled()) {
@@ -34,4 +37,5 @@ const [host, organizationId] = await Promise.all([
   })(),
 ]);
 
+asyncExitHook(cleanup, { wait: 1000 });
 sendToParent?.({ host, organizationId });

src/server.mts

@@ -57,26 +57,32 @@ export const StartServerSchema = z.object({
 });
 
 export async function startServer(options: z.infer<typeof StartServerSchema>) {
-  // TODO: Implement safe exit
-  logger.info(`Initializing storage`);
-  await Promise.all([
+  logger.info(`Initializing storage...`);
+  const callbacks = await Promise.all([
     opsConn.initialize(options.cwd),
     checkpointer.initialize(options.cwd),
     graphStore.initialize(options.cwd),
   ]);
 
+  const cleanup = async () => {
+    logger.info(`Flushing to persistent storage, exiting...`);
+    await Promise.all(callbacks.map((c) => c.flush()));
+  };
+
   logger.info(`Registering graphs from ${options.cwd}`);
   await registerFromEnv(options.graphs, { cwd: options.cwd });
 
   logger.info(`Starting ${options.nWorkers} workers`);
   for (let i = 0; i < options.nWorkers; i++) queue();
 
-  return new Promise<string>((resolve) => {
-    serve(
-      { fetch: app.fetch, port: options.port, hostname: options.host },
-      (c) => {
-        resolve(`${c.address}:${c.port}`);
-      }
-    );
-  });
+  return new Promise<{ host: string; cleanup: () => Promise<void> }>(
+    (resolve) => {
+      serve(
+        { fetch: app.fetch, port: options.port, hostname: options.host },
+        (c) => {
+          resolve({ host: `${c.address}:${c.port}`, cleanup });
+        }
+      );
+    }
+  );
 }

src/storage/persist.mts

@@ -19,7 +19,7 @@ export class FileSystemPersistence<Schema> {
   private defaultSchema: () => Schema;
   private name: string;
 
-  private flushChain: Promise<void> = Promise.resolve<void>(undefined);
+  private flushTimeout: NodeJS.Timeout | undefined = undefined;
 
   constructor(name: `.${string}.json`, defaultSchema: () => Schema) {
     this.name = name;
@@ -35,27 +35,26 @@ export class FileSystemPersistence<Schema> {
       this.data = this.defaultSchema();
     }
 
-    this.flushChain = this.flushChain.then(async () => {
-      if (this.data == null || this.filepath == null) return;
-      await fs.mkdir(path.dirname(this.filepath), { recursive: true });
-    });
+    await fs
+      .mkdir(path.dirname(this.filepath), { recursive: true })
+      .catch(() => void 0);
 
     return this;
   }
 
+  protected async persist() {
+    if (this.data == null || this.filepath == null) return;
+    clearTimeout(this.flushTimeout);
+    await fs.writeFile(this.filepath, superjson.stringify(this.data), "utf-8");
+  }
+
   protected schedulePersist() {
-    this.flushChain = this.flushChain.then(async () => {
-      if (this.data == null || this.filepath == null) return;
-      await fs.writeFile(
-        this.filepath,
-        superjson.stringify(this.data),
-        "utf-8"
-      );
-    });
+    clearTimeout(this.flushTimeout);
+    this.flushTimeout = setTimeout(() => this.persist(), 3000);
   }
 
   async flush() {
-    await this.flushChain;
+    await this.persist();
   }
 
   async with<T>(fn: (data: Schema) => T) {