From 8527528af400d28c5d04c439c557b8c27ae5b2b2 Mon Sep 17 00:00:00 2001
From: Zef Hemel <zef@zef.me>
Date: Wed, 6 Dec 2023 18:44:48 +0100
Subject: [PATCH] Lazy plugs (#596)

* Manifest caching and lazy loading of plug workers
* Fixes #546 Plug unloading after time out
---
 common/spaces/evented_space_primitives.ts |  8 ++-
 plug-api/silverbullet-syscall/space.ts    |  2 +-
 plugos/hooks/endpoint.test.ts             |  2 +
 plugos/lib/datastore.ts                   |  2 +-
 plugos/manifest_cache.ts                  | 49 +++++++++++++++++
 plugos/plug.ts                            | 54 ++++++++++++------
 plugos/runtime.test.ts                    |  2 +
 plugos/sandbox.ts                         | 43 +++++++++++----
 plugos/system.ts                          | 21 ++++++-
 plugs/markdown/markdown_render.test.ts    |  5 +-
 plugs/plug-manager/plugmanager.ts         |  2 +-
 server/server_system.ts                   | 27 +++++++--
 server/syscalls/space.ts                  |  2 +-
 web/client_system.ts                      | 67 ++++++++++++++++-------
 web/space.ts                              |  5 +-
 web/syscalls/space.ts                     |  2 +-
 16 files changed, 226 insertions(+), 67 deletions(-)
 create mode 100644 plugos/manifest_cache.ts

diff --git a/common/spaces/evented_space_primitives.ts b/common/spaces/evented_space_primitives.ts
index 216e53f..1917345 100644
--- a/common/spaces/evented_space_primitives.ts
+++ b/common/spaces/evented_space_primitives.ts
@@ -114,7 +114,13 @@ export class EventedSpacePrimitives implements SpacePrimitives {
         meta,
       );
       if (!selfUpdate) {
-        await this.dispatchEvent("file:changed", name, true);
+        await this.dispatchEvent(
+          "file:changed",
+          name,
+          true,
+          undefined,
+          newMeta.lastModified,
+        );
       }
       this.spaceSnapshot[name] = newMeta.lastModified;
 
diff --git a/plug-api/silverbullet-syscall/space.ts b/plug-api/silverbullet-syscall/space.ts
index 9a0675f..688886f 100644
--- a/plug-api/silverbullet-syscall/space.ts
+++ b/plug-api/silverbullet-syscall/space.ts
@@ -23,7 +23,7 @@ export function deletePage(name: string): Promise<void> {
   return syscall("space.deletePage", name);
 }
 
-export function listPlugs(): Promise<string[]> {
+export function listPlugs(): Promise<FileMeta[]> {
   return syscall("space.listPlugs");
 }
 
diff --git a/plugos/hooks/endpoint.test.ts b/plugos/hooks/endpoint.test.ts
index 1520e81..5c68bbc 100644
--- a/plugos/hooks/endpoint.test.ts
+++ b/plugos/hooks/endpoint.test.ts
@@ -18,6 +18,8 @@ Deno.test("Run a plugos endpoint server", async () => {
 
   await system.load(
     new URL(`file://${workerPath}`),
+    "test",
+    0,
     createSandbox,
   );
 
diff --git a/plugos/lib/datastore.ts b/plugos/lib/datastore.ts
index 17a57b9..8560193 100644
--- a/plugos/lib/datastore.ts
+++ b/plugos/lib/datastore.ts
@@ -9,7 +9,7 @@ import { KvPrimitives } from "./kv_primitives.ts";
  */
 export class DataStore {
   constructor(
-    private kv: KvPrimitives,
+    readonly kv: KvPrimitives,
     private prefix: KvKey = [],
     private functionMap: FunctionMap = builtinFunctions,
   ) {
diff --git a/plugos/manifest_cache.ts b/plugos/manifest_cache.ts
new file mode 100644
index 0000000..4099d47
--- /dev/null
+++ b/plugos/manifest_cache.ts
@@ -0,0 +1,49 @@
+import { KvPrimitives } from "./lib/kv_primitives.ts";
+import { Plug } from "./plug.ts";
+import { Manifest } from "./types.ts";
+
+export interface ManifestCache<T> {
+  getManifest(plug: Plug<T>, hash: number): Promise<Manifest<T>>;
+}
+
+export class KVPrimitivesManifestCache<T> implements ManifestCache<T> {
+  constructor(private kv: KvPrimitives, private manifestPrefix: string) {
+  }
+
+  async getManifest(plug: Plug<T>, hash: number): Promise<Manifest<T>> {
+    const [cached] = await this.kv.batchGet([[
+      this.manifestPrefix,
+      plug.name,
+    ]]);
+    if (cached && cached.hash === hash) {
+      // console.log("Using KV cached manifest for", plug.name);
+      return cached.manifest;
+    }
+    await plug.sandbox.init();
+    const manifest = plug.sandbox.manifest!;
+    await this.kv.batchSet([{
+      key: [this.manifestPrefix, plug.name],
+      value: { manifest, hash },
+    }]);
+    return manifest;
+  }
+}
+
+export class InMemoryManifestCache<T> implements ManifestCache<T> {
+  private cache = new Map<string, {
+    manifest: Manifest<T>;
+    hash: number;
+  }>();
+
+  async getManifest(plug: Plug<T>, hash: number): Promise<Manifest<T>> {
+    const cached = this.cache.get(plug.workerUrl.href);
+    if (cached && cached.hash === hash) {
+      // console.log("Using memory cached manifest for", plug.name);
+      return cached.manifest;
+    }
+    await plug.sandbox.init();
+    const manifest = plug.sandbox.manifest!;
+    this.cache.set(plug.name!, { manifest, hash });
+    return manifest;
+  }
+}
diff --git a/plugos/plug.ts b/plugos/plug.ts
index fbe8ff2..df2e124 100644
--- a/plugos/plug.ts
+++ b/plugos/plug.ts
@@ -9,34 +9,40 @@ export class Plug<HookT> {
   public grantedPermissions: string[] = [];
   public sandbox: Sandbox<HookT>;
 
-  // Resolves once the worker has been loaded
+  // Resolves once the plug's manifest is available
   ready: Promise<void>;
 
   // Only available after ready resolves
   public manifest?: Manifest<HookT>;
   public assets?: AssetBundle;
 
+  // Time of last function invocation
+  unloadTimeout?: number;
+
   constructor(
     private system: System<HookT>,
     public workerUrl: URL,
+    readonly name: string,
+    private hash: number,
     private sandboxFactory: (plug: Plug<HookT>) => Sandbox<HookT>,
   ) {
     this.runtimeEnv = system.env;
 
-    // Kick off worker
-    this.sandbox = this.sandboxFactory(this);
-    this.ready = this.sandbox.ready.then(() => {
-      this.manifest = this.sandbox.manifest!;
-      this.assets = new AssetBundle(
-        this.manifest.assets ? this.manifest.assets as AssetJson : {},
-      );
-      // TODO: These need to be explicitly granted, not just taken
-      this.grantedPermissions = this.manifest.requiredPermissions || [];
-    });
-  }
+    this.scheduleUnloadTimeout();
 
-  get name(): string | undefined {
-    return this.manifest?.name;
+    this.sandbox = this.sandboxFactory(this);
+    // Retrieve the manifest asynchonously, which may either come from a cache or be loaded from the worker
+    this.ready = system.options.manifestCache!.getManifest(this, this.hash)
+      .then(
+        (manifest) => {
+          this.manifest = manifest;
+          this.assets = new AssetBundle(
+            manifest.assets ? manifest.assets as AssetJson : {},
+          );
+          // TODO: These need to be explicitly granted, not just taken
+          this.grantedPermissions = manifest.requiredPermissions || [];
+        },
+      );
   }
 
   // Invoke a syscall
@@ -54,11 +60,26 @@ export class Plug<HookT> {
     return !funDef.env || !this.runtimeEnv || funDef.env === this.runtimeEnv;
   }
 
+  scheduleUnloadTimeout() {
+    if (!this.system.options.plugFlushTimeout) {
+      return;
+    }
+    // Reset the unload timeout, if set
+    if (this.unloadTimeout) {
+      clearTimeout(this.unloadTimeout);
+    }
+    this.unloadTimeout = setTimeout(() => {
+      this.stop();
+    }, this.system.options.plugFlushTimeout);
+  }
+
   // Invoke a function
   async invoke(name: string, args: any[]): Promise<any> {
     // Ensure the worker is fully up and running
     await this.ready;
 
+    this.scheduleUnloadTimeout();
+
     // Before we access the manifest
     const funDef = this.manifest!.functions[name];
     if (!funDef) {
@@ -90,8 +111,7 @@ export class Plug<HookT> {
   }
 
   stop() {
-    if (this.sandbox) {
-      this.sandbox.stop();
-    }
+    console.log("Stopping sandbox for", this.name);
+    this.sandbox.stop();
   }
 }
diff --git a/plugos/runtime.test.ts b/plugos/runtime.test.ts
index 78710d4..d1a2c7f 100644
--- a/plugos/runtime.test.ts
+++ b/plugos/runtime.test.ts
@@ -34,6 +34,8 @@ Deno.test("Run a deno sandbox", async () => {
 
   const plug = await system.load(
     new URL(`file://${workerPath}`),
+    "test",
+    0,
     createSandbox,
   );
 
diff --git a/plugos/sandbox.ts b/plugos/sandbox.ts
index a9da751..e5cf058 100644
--- a/plugos/sandbox.ts
+++ b/plugos/sandbox.ts
@@ -9,26 +9,38 @@ export type SandboxFactory<HookT> = (plug: Plug<HookT>) => Sandbox<HookT>;
  * Effectively this wraps a web worker, the reason to have this split from Plugs is to allow plugs to manage multiple sandboxes, e.g. for performance in the future
  */
 export class Sandbox<HookT> {
-  private worker: Worker;
+  private worker?: Worker;
   private reqId = 0;
   private outstandingInvocations = new Map<
     number,
     { resolve: (result: any) => void; reject: (e: any) => void }
   >();
 
-  public ready: Promise<void>;
+  // public ready: Promise<void>;
   public manifest?: Manifest<HookT>;
 
   constructor(
     readonly plug: Plug<HookT>,
-    workerOptions = {},
+    private workerOptions = {},
   ) {
-    this.worker = new Worker(plug.workerUrl, {
-      ...workerOptions,
+  }
+
+  /**
+   * Should only invoked lazily (either by invoke, or by a ManifestCache to load the manifest)
+   */
+  init(): Promise<void> {
+    console.log("Booting up worker for", this.plug.name);
+    if (this.worker) {
+      // Should not happen
+      console.warn("Double init of sandbox");
+    }
+    this.worker = new Worker(this.plug.workerUrl, {
+      ...this.workerOptions,
       type: "module",
     });
-    this.ready = new Promise((resolve) => {
-      this.worker.onmessage = (ev) => {
+
+    return new Promise((resolve) => {
+      this.worker!.onmessage = (ev) => {
         if (ev.data.type === "manifest") {
           this.manifest = ev.data.manifest;
           resolve();
@@ -46,14 +58,14 @@ export class Sandbox<HookT> {
         try {
           const result = await this.plug.syscall(data.name!, data.args!);
 
-          this.worker.postMessage({
+          this.worker!.postMessage({
             type: "sysr",
             id: data.id,
             result: result,
           } as WorkerMessage);
         } catch (e: any) {
           // console.error("Syscall fail", e);
-          this.worker.postMessage({
+          this.worker!.postMessage({
             type: "sysr",
             id: data.id,
             error: e.message,
@@ -76,9 +88,13 @@ export class Sandbox<HookT> {
     }
   }
 
-  invoke(name: string, args: any[]): Promise<any> {
+  async invoke(name: string, args: any[]): Promise<any> {
+    if (!this.worker) {
+      // Lazy initialization
+      await this.init();
+    }
     this.reqId++;
-    this.worker.postMessage({
+    this.worker!.postMessage({
       type: "inv",
       id: this.reqId,
       name,
@@ -90,6 +106,9 @@ export class Sandbox<HookT> {
   }
 
   stop() {
-    this.worker.terminate();
+    if (this.worker) {
+      this.worker.terminate();
+      this.worker = undefined;
+    }
   }
 }
diff --git a/plugos/system.ts b/plugos/system.ts
index 7296f38..7e61dc4 100644
--- a/plugos/system.ts
+++ b/plugos/system.ts
@@ -3,6 +3,7 @@ import { EventEmitter } from "./event.ts";
 import type { SandboxFactory } from "./sandbox.ts";
 import { Plug } from "./plug.ts";
 import { deepObjectMerge } from "$sb/lib/json.ts";
+import { InMemoryManifestCache, ManifestCache } from "./manifest_cache.ts";
 
 export interface SysCallMapping {
   [key: string]: (ctx: SyscallContext, ...args: any) => Promise<any> | any;
@@ -28,13 +29,27 @@ type Syscall = {
   callback: SyscallSignature;
 };
 
+export type SystemOptions = {
+  manifestCache?: ManifestCache<any>;
+  plugFlushTimeout?: number;
+};
+
 export class System<HookT> extends EventEmitter<SystemEvents<HookT>> {
   protected plugs = new Map<string, Plug<HookT>>();
   protected registeredSyscalls = new Map<string, Syscall>();
   protected enabledHooks = new Set<Hook<HookT>>();
 
-  constructor(readonly env?: string) {
+  /**
+   * @param env either an environment or undefined for hybrid mode
+   */
+  constructor(
+    readonly env: string | undefined,
+    readonly options: SystemOptions = {},
+  ) {
     super();
+    if (!options.manifestCache) {
+      options.manifestCache = new InMemoryManifestCache();
+    }
   }
 
   get loadedPlugs(): Map<string, Plug<HookT>> {
@@ -94,11 +109,13 @@ export class System<HookT> extends EventEmitter<SystemEvents<HookT>> {
 
   async load(
     workerUrl: URL,
+    name: string,
+    hash: number,
     sandboxFactory: SandboxFactory<HookT>,
     // Mapping plug name -> manifest overrides
     manifestOverrides?: Record<string, Partial<Manifest<HookT>>>,
   ): Promise<Plug<HookT>> {
-    const plug = new Plug(this, workerUrl, sandboxFactory);
+    const plug = new Plug(this, workerUrl, name, hash, sandboxFactory);
 
     // Wait for worker to boot, and pass back its manifest
     await plug.ready;
diff --git a/plugs/markdown/markdown_render.test.ts b/plugs/markdown/markdown_render.test.ts
index 78a475b..220ca71 100644
--- a/plugs/markdown/markdown_render.test.ts
+++ b/plugs/markdown/markdown_render.test.ts
@@ -5,16 +5,19 @@ import { System } from "../../plugos/system.ts";
 import { createSandbox } from "../../plugos/environments/deno_sandbox.ts";
 import { loadMarkdownExtensions } from "../../common/markdown_parser/markdown_ext.ts";
 import { renderMarkdownToHtml } from "./markdown_render.ts";
-import { assertEquals } from "../../test_deps.ts";
 
 Deno.test("Markdown render", async () => {
   const system = new System<any>("server");
   await system.load(
     new URL("../../dist_plug_bundle/_plug/editor.plug.js", import.meta.url),
+    "editor",
+    0,
     createSandbox,
   );
   await system.load(
     new URL("../../dist_plug_bundle/_plug/tasks.plug.js", import.meta.url),
+    "tasks",
+    0,
     createSandbox,
   );
   const lang = buildMarkdown(loadMarkdownExtensions(system));
diff --git a/plugs/plug-manager/plugmanager.ts b/plugs/plug-manager/plugmanager.ts
index 5b9d9e6..12bda32 100644
--- a/plugs/plug-manager/plugmanager.ts
+++ b/plugs/plug-manager/plugmanager.ts
@@ -68,7 +68,7 @@ export async function updatePlugsCommand() {
 
     const allPlugNames = [...builtinPlugNames, ...allCustomPlugNames];
     // And delete extra ones
-    for (const existingPlug of await space.listPlugs()) {
+    for (const { name: existingPlug } of await space.listPlugs()) {
       const plugName = existingPlug.substring(
         "_plug/".length,
         existingPlug.length - ".plug.js".length,
diff --git a/server/server_system.ts b/server/server_system.ts
index 8c6980a..d8a9a40 100644
--- a/server/server_system.ts
+++ b/server/server_system.ts
@@ -33,11 +33,14 @@ import { languageSyscalls } from "../common/syscalls/language.ts";
 import { handlebarsSyscalls } from "../common/syscalls/handlebars.ts";
 import { codeWidgetSyscalls } from "../web/syscalls/code_widget.ts";
 import { CodeWidgetHook } from "../web/hooks/code_widget.ts";
+import { KVPrimitivesManifestCache } from "../plugos/manifest_cache.ts";
 
 const fileListInterval = 30 * 1000; // 30s
 
+const plugNameExtractRegex = /\/(.+)\.plug\.js$/;
+
 export class ServerSystem {
-  system: System<SilverBulletHooks> = new System("server");
+  system!: System<SilverBulletHooks>;
   spacePrimitives!: SpacePrimitives;
   denoKv!: Deno.Kv;
   listInterval?: number;
@@ -52,6 +55,18 @@ export class ServerSystem {
 
   // Always needs to be invoked right after construction
   async init(awaitIndex = false) {
+    this.denoKv = await Deno.openKv(this.dbPath);
+    const kvPrimitives = new DenoKvPrimitives(this.denoKv);
+    this.ds = new DataStore(kvPrimitives);
+
+    this.system = new System(
+      "server",
+      {
+        manifestCache: new KVPrimitivesManifestCache(kvPrimitives, "manifest"),
+        plugFlushTimeout: 5 * 60 * 1000, // 5 minutes
+      },
+    );
+
     // Event hook
     const eventHook = new EventHook();
     this.system.addHook(eventHook);
@@ -60,9 +75,6 @@ export class ServerSystem {
     const cronHook = new CronHook(this.system);
     this.system.addHook(cronHook);
 
-    this.denoKv = await Deno.openKv(this.dbPath);
-    this.ds = new DataStore(new DenoKvPrimitives(this.denoKv));
-
     // Endpoint hook
     this.system.addHook(new EndpointHook(this.app, "/_/"));
 
@@ -179,10 +191,13 @@ export class ServerSystem {
   }
 
   async loadPlugFromSpace(path: string): Promise<Plug<SilverBulletHooks>> {
-    const plugJS = (await this.spacePrimitives.readFile(path)).data;
+    const { meta, data } = await this.spacePrimitives.readFile(path);
+    const plugName = path.match(plugNameExtractRegex)![1];
     return this.system.load(
       // Base64 encoding this to support `deno compile` mode
-      new URL(base64EncodedDataUrl("application/javascript", plugJS)),
+      new URL(base64EncodedDataUrl("application/javascript", data)),
+      plugName,
+      meta.lastModified,
       createSandbox,
     );
   }
diff --git a/server/syscalls/space.ts b/server/syscalls/space.ts
index 0ad7df6..17660f5 100644
--- a/server/syscalls/space.ts
+++ b/server/syscalls/space.ts
@@ -29,7 +29,7 @@ export function spaceSyscalls(space: Space): SysCallMapping {
     "space.deletePage": async (_ctx, name: string) => {
       await space.deletePage(name);
     },
-    "space.listPlugs": (): Promise<string[]> => {
+    "space.listPlugs": (): Promise<FileMeta[]> => {
       return space.listPlugs();
     },
     "space.listAttachments": async (): Promise<AttachmentMeta[]> => {
diff --git a/web/client_system.ts b/web/client_system.ts
index 4b1545b..3d43c0a 100644
--- a/web/client_system.ts
+++ b/web/client_system.ts
@@ -38,6 +38,12 @@ import { languageSyscalls } from "../common/syscalls/language.ts";
 import { handlebarsSyscalls } from "../common/syscalls/handlebars.ts";
 import { codeWidgetSyscalls } from "./syscalls/code_widget.ts";
 import { clientCodeWidgetSyscalls } from "./syscalls/client_code_widget.ts";
+import {
+  InMemoryManifestCache,
+  KVPrimitivesManifestCache,
+} from "../plugos/manifest_cache.ts";
+
+const plugNameExtractRegex = /\/(.+)\.plug\.js$/;
 
 export class ClientSystem {
   commandHook: CommandHook;
@@ -51,11 +57,18 @@ export class ClientSystem {
     private client: Client,
     private mq: MessageQueue,
     private ds: DataStore,
-    // private dbPrefix: string,
     private eventHook: EventHook,
   ) {
     // Only set environment to "client" when running in thin client mode, otherwise we run everything locally (hybrid)
-    this.system = new System(client.syncMode ? undefined : "client");
+    this.system = new System(
+      client.syncMode ? undefined : "client",
+      {
+        manifestCache: new KVPrimitivesManifestCache<SilverBulletHooks>(
+          ds.kv,
+          "manifest",
+        ),
+      },
+    );
 
     this.system.addHook(this.eventHook);
 
@@ -93,22 +106,28 @@ export class ClientSystem {
     this.slashCommandHook = new SlashCommandHook(this.client);
     this.system.addHook(this.slashCommandHook);
 
-    this.eventHook.addLocalListener("file:changed", async (path: string) => {
-      if (path.startsWith("_plug/") && path.endsWith(".plug.js")) {
-        console.log("Plug updated, reloading:", path);
-        this.system.unload(path);
-        const plug = await this.system.load(
-          new URL(`/${path}`, location.href),
-          createSandbox,
-          this.client.settings.plugOverrides,
-        );
-        if ((plug.manifest! as Manifest).syntax) {
-          // If there are syntax extensions, rebuild the markdown parser immediately
-          this.updateMarkdownParser();
+    this.eventHook.addLocalListener(
+      "file:changed",
+      async (path: string, _selfUpdate, _oldHash, newHash) => {
+        if (path.startsWith("_plug/") && path.endsWith(".plug.js")) {
+          const plugName = plugNameExtractRegex.exec(path)![1];
+          console.log("Plug updated, reloading", plugName, "from", path);
+          this.system.unload(path);
+          const plug = await this.system.load(
+            new URL(`/${path}`, location.href),
+            plugName,
+            newHash,
+            createSandbox,
+            this.client.settings.plugOverrides,
+          );
+          if ((plug.manifest! as Manifest).syntax) {
+            // If there are syntax extensions, rebuild the markdown parser immediately
+            this.updateMarkdownParser();
+          }
+          this.client.debouncedPlugsUpdatedEvent();
         }
-        this.client.debouncedPlugsUpdatedEvent();
-      }
-    });
+      },
+    );
 
     // Debugging
     // this.eventHook.addLocalListener("file:listed", (files) => {
@@ -177,15 +196,23 @@ export class ClientSystem {
     await space.updatePageList();
     await this.system.unloadAll();
     console.log("(Re)loading plugs");
-    await Promise.all((await space.listPlugs()).map(async (plugName) => {
+    await Promise.all((await space.listPlugs()).map(async (plugMeta) => {
       try {
+        const plugName = plugNameExtractRegex.exec(plugMeta.name)![1];
         await this.system.load(
-          new URL(plugName, location.origin),
+          new URL(plugMeta.name, location.origin),
+          plugName,
+          plugMeta.lastModified,
           createSandbox,
           this.client.settings.plugOverrides,
         );
       } catch (e: any) {
-        console.error("Could not load plug", plugName, "error:", e.message);
+        console.error(
+          "Could not load plug",
+          plugMeta.name,
+          "error:",
+          e.message,
+        );
       }
     }));
   }
diff --git a/web/space.ts b/web/space.ts
index cf709c6..a4ed117 100644
--- a/web/space.ts
+++ b/web/space.ts
@@ -103,14 +103,13 @@ export class Space {
     return this.cachedPageList;
   }
 
-  async listPlugs(): Promise<string[]> {
+  async listPlugs(): Promise<FileMeta[]> {
     const files = await this.spacePrimitives.fetchFileList();
     return files
       .filter((fileMeta) =>
         fileMeta.name.startsWith(plugPrefix) &&
         fileMeta.name.endsWith(".plug.js")
-      )
-      .map((fileMeta) => fileMeta.name);
+      );
   }
 
   async readPage(name: string): Promise<{ text: string; meta: PageMeta }> {
diff --git a/web/syscalls/space.ts b/web/syscalls/space.ts
index 85f9866..a0dc57d 100644
--- a/web/syscalls/space.ts
+++ b/web/syscalls/space.ts
@@ -33,7 +33,7 @@ export function spaceSyscalls(editor: Client): SysCallMapping {
       console.log("Deleting page");
       await editor.space.deletePage(name);
     },
-    "space.listPlugs": (): Promise<string[]> => {
+    "space.listPlugs": (): Promise<FileMeta[]> => {
       return editor.space.listPlugs();
     },
     "space.listAttachments": async (): Promise<AttachmentMeta[]> => {