Lazy plugs (#596)

* Manifest caching and lazy loading of plug workers
* Fixes #546 Plug unloading after time out

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,
   );
 

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,
   ) {

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;
+  }
+}

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();
   }
 }

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,
   );
 

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;
+    }
   }
 }

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;