feat: wasm worker manager (#966)

<!--
Pull requests are squashed and merged using:
- their title as the commit message
- their description as the commit body

Having a good title and description is important for the users to get
readable changelog.
-->

<!-- 1. Explain WHAT the change is about -->

- Closes
[MET-805](https://linear.app/metatypedev/issue/MET-805/worker-manager-implementation-for-wasm-runtime-and-python).

<!-- 3. Explain HOW users should update their code -->

#### Migration notes

---

- [x] The change comes with new or modified tests
- [x] Hard-to-understand functions have explanatory comments
- [x] End-user documentation is updated to reflect the change


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

- **New Features**
  - Added WebAssembly (WASM) worker management capabilities.
  - Introduced new type definitions for WASM runtime interactions.
  - Enhanced worker pool and worker manager functionality.
- Introduced a new `WasmWorker` class for managing WebAssembly workers.

- **Configuration Changes**
- Updated global configuration to include WASM worker-related settings.
- Replaced `substantial_worker_wait_timeout_ms` with more granular
worker configuration options.
- Added new environment variables for WASM worker management:
`MIN_WASM_WORKERS` and `MAX_WASM_WORKERS`.

- **Runtime Improvements**
  - Refactored runtime initialization processes.
  - Improved worker lifecycle management.
- Enhanced error handling and logging for various runtime environments.
  - Removed explicit `destroy` operations in some runtime components.

- **Dependency Updates**
  - Updated Metatype version to `0.5.1-rc.0`.

- **Performance Optimizations**
  - Added pre-warming steps in performance tests.
  - Streamlined worker initialization and management.

- **Code Quality**
  - Improved code formatting and readability.
  - Simplified runtime and worker management interfaces.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Natoandro <[email protected]>

Author: luckasRanarison

deno.lock

@@ -87,4 +87,6 @@ The following environment variables can be used to configure the typegate. `SYNC
 | DENO_WORKER_WAIT_TIMEOUT_MS        | Timeout for waiting for a free deno worker                                                               | 5000                     | 2000                                                                                     |
 | MIN_SUBSTANTIAL_WORKERS            | Minimal number of available substantial workers                                                          | 2                        | 4                                                                                        |
 | MAX_SUBSTANTIAL_WORKERS            | Maximal number of available substantial workers                                                          | 8                        | 16                                                                                       |
+| MIN_WASM_WORKERS                   | Minimal number of available wasm workers                                                                 | 2                        | 4                                                                                        |
+| MAX_WASM_WORKERS                   | Maximal number of available wasm workers                                                                 | 8                        | 16                                                                                       |
 | SUBSTANTIAL_WORKER_WAIT_TIMEOUT_MS | Timeout for waiting for a free substantial worker                                                        | 15000                    | 2000                                                                                     |

docs/metatype.dev/docs/reference/typegate/index.mdx

@@ -55,7 +55,6 @@ const Meta = {
   wasmtimeWit: getOp("op_wasmtime_wit"),
   wit_wire: {
     init: getOp("op_wit_wire_init"),
-    destroy: getOp("op_wit_wire_destroy"),
     handle: getOp("op_wit_wire_handle"),
   },
   substantial: {

src/typegate/engine/00_runtime.js

@@ -45,7 +45,6 @@ export type MetaNS = {
       args: WitWireInitArgs,
       cb: (op_name: string, json: string) => Promise<string>,
     ) => Promise<WitWireInitResponse>;
-    destroy: (instanceId: string) => Promise<void>;
     handle: (
       instanceId: string,
       args: WitWireReq,

src/typegate/engine/runtime.d.ts

@@ -40,7 +40,6 @@ deno_core::extension!(
         prisma::op_archive,
         wit_wire::op_wit_wire_init,
         wit_wire::op_wit_wire_handle,
-        wit_wire::op_wit_wire_destroy,
         grpc::op_grpc_register,
         grpc::op_grpc_unregister,
         grpc::op_call_grpc_method,

src/typegate/engine/src/ext.rs

@@ -96,11 +96,40 @@ impl Ctx {
 }
 
 struct Instance {
+    id: String,
     bindings: wit::WitWire,
     store: wasmtime::Store<InstanceState>,
     preopen_dir: PathBuf,
 }
 
+// An Instance's lifetime is tied to a thread (worker)
+// The Instance is dropped when a worker is terminated
+impl Drop for Instance {
+    fn drop(&mut self) {
+        let id = self.id.clone();
+        let preopen_dir = self.preopen_dir.clone();
+
+        debug!("destroying wit_wire instance {id}");
+
+        // FIXME: Does it really have to be async?
+        std::mem::drop(tokio::task::spawn(async move {
+            debug!("Removing preopened dir: {preopen_dir:?}");
+
+            match tokio::fs::remove_dir_all(&preopen_dir).await {
+                Err(err) if err.kind() != std::io::ErrorKind::NotFound => {
+                    error!(
+                        "error removing preopend dir for instance {id} at {:?}: {err}",
+                        preopen_dir
+                    )
+                }
+                _ => {
+                    debug!("preopened dir removed for instance {id}")
+                }
+            }
+        }));
+    }
+}
+
 struct InstanceState {
     table: wasmtime_wasi::ResourceTable,
     ctx: wasmtime_wasi::WasiCtx,
@@ -157,14 +186,6 @@ struct TypegateHost {
 struct SendPtr<T>(NonNull<T>);
 unsafe impl<T> Send for SendPtr<T> {}
 
-impl TypegateHost {
-    fn drop(self, scope: &mut v8::HandleScope) {
-        unsafe {
-            _ = v8::Global::from_raw(scope, self.js_fn.0);
-        }
-    }
-}
-
 #[wasmtime_wasi::async_trait]
 impl Host for TypegateHost {
     async fn hostcall(
@@ -389,8 +410,9 @@ pub async fn op_wit_wire_init<'scope>(
     })??; // <- note second try for the wit err. we have an into impl above
     assert!(res.ok);
     ctx.instances.insert(
-        instance_id,
+        instance_id.clone(),
         Instance {
+            id: instance_id,
             bindings,
             store,
             preopen_dir: work_dir,
@@ -399,34 +421,6 @@ pub async fn op_wit_wire_init<'scope>(
     Ok(WitWireInitResponse {})
 }
 
-#[deno_core::op2(fast)]
-pub fn op_wit_wire_destroy(
-    state: Rc<RefCell<OpState>>,
-    scope: &mut v8::HandleScope<'_>,
-    #[string] instance_id: String,
-) {
-    debug!("destroying wit_wire instance {instance_id}");
-    let ctx = {
-        let state = state.borrow();
-        let ctx = state.borrow::<Ctx>();
-        ctx.clone()
-    };
-
-    let Some((_id, instance)) = ctx.instances.remove(&instance_id) else {
-        return;
-    };
-    let tg_host = instance.store.into_data().tg_host;
-    tg_host.drop(scope);
-    std::mem::drop(tokio::task::spawn(async move {
-        if let Err(err) = tokio::fs::remove_dir_all(&instance.preopen_dir).await {
-            error!(
-                "error removing preopend dir for instance {_id} at {:?}: {err}",
-                instance.preopen_dir
-            )
-        }
-    }));
-}
-
 #[derive(Deserialize)]
 #[serde(crate = "serde")]
 pub struct WitWireReq {

src/typegate/engine/src/runtimes/wit_wire.rs

@@ -66,9 +66,13 @@ export const globalConfigSchema = z.object({
   // deno_idle_worker_timeout_ms: z.coerce.number().positive().optional(), // auto remove idle workers
   min_substantial_workers: z.coerce.number().positive().default(2),
   max_substantial_workers: z.coerce.number().positive().default(8),
-  substantial_worker_wait_timeout_ms: z.coerce.number().positive().default(
-    15000,
-  ),
+  min_wasm_workers: z.coerce.number().positive().default(2),
+  max_wasm_workers: z.coerce.number().positive().default(8),
+  wasm_worker_wait_timeout_ms: z.coerce.number().positive().default(5000),
+  substantial_worker_wait_timeout_ms: z.coerce
+    .number()
+    .positive()
+    .default(15000),
   // substantial_idle_worker_timeout_ms: z.coerce.number().positive().optional(), // auto remove idle workers
 });
 export type GlobalConfig = z.infer<typeof globalConfigSchema>;
@@ -88,8 +92,7 @@ export const typegateConfigBaseSchema = z.object({
       if (bytes.length != 64) {
         ctx.addIssue({
           code: z.ZodIssueCode.custom,
-          message:
-            `Base64 contains ${bytes.length} instead of 64 bytes (use openssl rand -base64 64 | tr -d '\n')`,
+          message: `Base64 contains ${bytes.length} instead of 64 bytes (use openssl rand -base64 64 | tr -d '\n')`,
         });
       }
       return bytes;

src/typegate/src/config/types.ts

@@ -34,14 +34,21 @@ const predefinedFuncs: Record<
   string,
   (param: any) => Resolver<Record<string, unknown>>
 > = {
-  identity: () => ({ _, ...args }) => args,
+  identity:
+    () =>
+    ({ _, ...args }) =>
+      args,
   true: () => () => true,
   false: () => () => false,
   allow: () => () => "ALLOW" as PolicyResolverOutput,
   deny: () => () => "DENY" as PolicyResolverOutput,
   pass: () => () => "PASS" as PolicyResolverOutput,
-  internal_policy: () => ({ _: { context } }) =>
-    context.provider === "internal" ? "ALLOW" : "PASS" as PolicyResolverOutput,
+  internal_policy:
+    () =>
+    ({ _: { context } }) =>
+      context.provider === "internal"
+        ? "ALLOW"
+        : ("PASS" as PolicyResolverOutput),
   context_check: ({ key, value }) => {
     let check: (value: any) => boolean;
     switch (value.type) {
@@ -63,7 +70,7 @@ const predefinedFuncs: Record<
       for (const segment of path) {
         value = value?.[segment];
       }
-      return check(value) ? "PASS" : "DENY" as PolicyResolverOutput;
+      return check(value) ? "PASS" : ("DENY" as PolicyResolverOutput);
     };
   },
 };
@@ -102,9 +109,8 @@ export class DenoRuntime extends Runtime {
     for (const m of materializers) {
       let matSecrets = (m.data.secrets as string[]) ?? [];
       if (m.name === "outjection") {
-        matSecrets = m.data.source === "secret"
-          ? [...getInjectionValues(m.data)]
-          : [];
+        matSecrets =
+          m.data.source === "secret" ? [...getInjectionValues(m.data)] : [];
       }
       for (const secretName of (m.data.secrets as []) ?? []) {
         secrets[secretName] = secretManager.secretOrFail(secretName);
@@ -137,7 +143,7 @@ export class DenoRuntime extends Runtime {
         const matData = mat.data;
         const entryPoint = artifacts[matData.entryPoint as string];
         const depMetas = (matData.deps as string[]).map((dep) =>
-          createArtifactMeta(typegraphName, artifacts[dep])
+          createArtifactMeta(typegraphName, artifacts[dep]),
         );
         const moduleMeta = createArtifactMeta(typegraphName, entryPoint);
 
@@ -273,7 +279,7 @@ export class DenoRuntime extends Runtime {
       const entryPoint =
         this.tg.meta.artifacts[modMat.data.entryPoint as string];
       const depMetas = (modMat.data.deps as string[]).map((dep) =>
-        createArtifactMeta(this.typegraphName, this.tg.meta.artifacts[dep])
+        createArtifactMeta(this.typegraphName, this.tg.meta.artifacts[dep]),
       );
       const moduleMeta = createArtifactMeta(this.typegraphName, entryPoint);
 

src/typegate/src/runtimes/deno/deno.ts

@@ -12,8 +12,10 @@ export interface DenoWorkerError extends BaseMessage {
 
 export type BaseDenoWorkerMessage = BaseMessage | DenoWorkerError;
 
-export class DenoWorker<M extends BaseMessage, E extends BaseDenoWorkerMessage>
-  extends BaseWorker<M, E> {
+export class DenoWorker<
+  M extends BaseMessage,
+  E extends BaseDenoWorkerMessage,
+> extends BaseWorker<M, E> {
   #worker: Worker;
   #workerId: string;
   constructor(workerId: string, workerPath: string) {
@@ -44,12 +46,10 @@ export class DenoWorker<M extends BaseMessage, E extends BaseDenoWorkerMessage>
     };
 
     this.#worker.onerror = async (event) => {
-      await handlerFn(
-        {
-          type: "WORKER_ERROR",
-          event,
-        } as E,
-      );
+      await handlerFn({
+        type: "WORKER_ERROR",
+        event,
+      } as E);
     };
   }
 

src/typegate/src/runtimes/patterns/worker_manager/deno.ts

@@ -29,10 +29,13 @@ export class BaseWorkerManager<
   M extends BaseMessage,
   E extends BaseMessage,
 > {
-  #activeTasks: Map<TaskId, {
-    worker: BaseWorker<M, E>;
-    taskSpec: T;
-  }> = new Map();
+  #activeTasks: Map<
+    TaskId,
+    {
+      worker: BaseWorker<M, E>;
+      taskSpec: T;
+    }
+  > = new Map();
   #tasksByName: Map<string, Set<TaskId>> = new Map();
   #startedAt: Map<TaskId, Date> = new Map();
   #pool: WorkerPool<T, M, E>;
@@ -94,11 +97,12 @@ export class BaseWorkerManager<
     if (activeTaskNames.length > 0) {
       if (options.destroy) {
         logger.warn(
-          `destroying workers for tasks ${
-            activeTaskNames.map((w) => `"${w}"`).join(", ")
-          }`,
+          `destroying workers for tasks ${activeTaskNames
+            .map((w) => `"${w}"`)
+            .join(", ")}`,
         );
       }
+
       for (const name of activeTaskNames) {
         this.deallocateWorkersByName(name, options);
       }
@@ -111,8 +115,8 @@ export class BaseWorkerManager<
   ) {
     const taskIds = this.#tasksByName.get(name);
     if (taskIds) {
-      for (const taskId of taskIds) {
-        this.deallocateWorker(name, taskId, options);
+      for (const id of taskIds) {
+        this.deallocateWorker(name, id, options);
       }
       return true;
     }
@@ -163,7 +167,6 @@ export class BaseWorkerManager<
   deinit() {
     this.deallocateAllWorkers({ destroy: true, shutdown: true });
     this.#pool.clear(); // this will be called more than once, but that is ok
-    return Promise.resolve();
   }
 }
 

src/typegate/src/runtimes/patterns/worker_manager/mod.ts

@@ -21,6 +21,7 @@ export type Consumer<T> = (x: T) => void;
 export interface WaitQueue<W> {
   push(consumer: Consumer<W>, onCancel: () => void): void;
   shift(produce: () => W): boolean;
+  clear(): void;
 }
 
 export function createSimpleWaitQueue<W>(): WaitQueue<W> {
@@ -37,6 +38,7 @@ export function createSimpleWaitQueue<W>(): WaitQueue<W> {
       }
       return false;
     },
+    clear() {},
   };
 }
 
@@ -76,15 +78,21 @@ export class WaitQueueWithTimeout<W> implements WaitQueue<W> {
     return false;
   }
 
+  clear() {
+    if (this.#timerId != null) {
+      clearTimeout(this.#timerId);
+    }
+  }
+
   #timeoutHandler() {
     this.#cancelNextEntry();
     this.#updateTimer();
   }
 
   #updateTimer() {
     if (this.#queue.length > 0) {
-      const timeoutMs = this.#queue[0].addedAt + this.#waitTimeoutMs -
-        Date.now();
+      const timeoutMs =
+        this.#queue[0].addedAt + this.#waitTimeoutMs - Date.now();
       if (timeoutMs <= 0) {
         this.#cancelNextEntry();
         this.#updateTimer();
@@ -97,13 +105,11 @@ export class WaitQueueWithTimeout<W> implements WaitQueue<W> {
   }
 
   #cancelNextEntry() {
-    this.#queue.shift()!.cancellationHandler();
+    this.#queue.shift()?.cancellationHandler();
   }
 
   [Symbol.dispose]() {
-    if (this.#timerId != null) {
-      clearTimeout(this.#timerId);
-    }
+    this.clear();
   }
 }
 
@@ -214,15 +220,12 @@ export class WorkerPool<
 
   clear() {
     logger.warn(
-      `destroying idle workers: ${
-        this.#idleWorkers
-          .map((w) => `"${w.id}"`)
-          .join(", ")
-      }`,
+      `destroying idle workers: ${this.#idleWorkers
+        .map((w) => `"${w.id}"`)
+        .join(", ")}`,
     );
-    for (const worker of this.#idleWorkers) {
-      worker.destroy();
-    }
+    this.#idleWorkers.forEach((worker) => worker.destroy());
     this.#idleWorkers = [];
+    this.#waitQueue.clear();
   }
 }