@ricsam/isolate-daemon 0.1.5 → 0.1.6

package/dist/mjs/connection.mjs

@@ -7,7 +7,6 @@ import {
   buildFrame,
   MessageType,
   ErrorCode,
-  STREAM_THRESHOLD,
   STREAM_CHUNK_SIZE,
   STREAM_DEFAULT_CREDIT,
   marshalValue,
@@ -56,18 +55,26 @@ function handleConnection(socket, state) {
     for (const isolateId of connection.isolates) {
       const instance = state.isolates.get(isolateId);
       if (instance) {
-        try {
-          if (instance.playwrightHandle) {
-            instance.playwrightHandle.dispose();
-          }
-          instance.runtime.dispose();
-        } catch {}
-        state.isolates.delete(isolateId);
+        if (instance.namespaceId != null && !instance.isDisposed) {
+          softDeleteRuntime(instance, state);
+        } else if (!instance.isDisposed) {
+          try {
+            if (instance.playwrightHandle) {
+              instance.playwrightHandle.dispose();
+            }
+            instance.runtime.dispose();
+          } catch {}
+          state.isolates.delete(isolateId);
+        }
       }
     }
     for (const [, pending] of connection.pendingCallbacks) {
+      if (pending.timeoutId) {
+        clearTimeout(pending.timeoutId);
+      }
       pending.reject(new Error("Connection closed"));
     }
+    connection.pendingCallbacks.clear();
     state.connections.delete(socket);
   });
   socket.on("error", (err) => {
@@ -193,10 +200,147 @@ async function handleMessage(message, connection, state) {
       sendError(connection.socket, message.requestId ?? 0, ErrorCode.UNKNOWN_MESSAGE_TYPE, `Unknown message type: ${message.type}`);
   }
 }
+function softDeleteRuntime(instance, state) {
+  instance.isDisposed = true;
+  instance.disposedAt = Date.now();
+  instance.ownerConnection = null;
+  instance.callbacks.clear();
+  instance.runtime.timers.clearAll();
+  instance.runtime.console.reset();
+  instance.pendingCallbacks.length = 0;
+  instance.returnedCallbacks?.clear();
+  instance.returnedPromises?.clear();
+  instance.returnedIterators?.clear();
+}
+function reuseNamespacedRuntime(instance, connection, message, state) {
+  instance.ownerConnection = connection.socket;
+  instance.isDisposed = false;
+  instance.disposedAt = undefined;
+  instance.lastActivity = Date.now();
+  connection.isolates.add(instance.isolateId);
+  const callbacks = message.options.callbacks;
+  if (instance.callbackContext) {
+    instance.callbackContext.connection = connection;
+    instance.callbackContext.consoleOnEntry = callbacks?.console?.onEntry?.callbackId;
+    instance.callbackContext.fetch = callbacks?.fetch?.callbackId;
+    instance.callbackContext.moduleLoader = callbacks?.moduleLoader?.callbackId;
+    instance.callbackContext.fs = {
+      readFile: callbacks?.fs?.readFile?.callbackId,
+      writeFile: callbacks?.fs?.writeFile?.callbackId,
+      stat: callbacks?.fs?.stat?.callbackId,
+      readdir: callbacks?.fs?.readdir?.callbackId,
+      unlink: callbacks?.fs?.unlink?.callbackId,
+      mkdir: callbacks?.fs?.mkdir?.callbackId,
+      rmdir: callbacks?.fs?.rmdir?.callbackId
+    };
+    instance.callbackContext.custom.clear();
+    if (callbacks?.custom) {
+      for (const [name, reg] of Object.entries(callbacks.custom)) {
+        if (reg) {
+          instance.callbackContext.custom.set(name, reg.callbackId);
+        }
+      }
+    }
+  }
+  instance.callbacks.clear();
+  if (callbacks?.console?.onEntry) {
+    instance.callbacks.set(callbacks.console.onEntry.callbackId, {
+      ...callbacks.console.onEntry,
+      name: "onEntry"
+    });
+  }
+  if (callbacks?.fetch) {
+    instance.callbacks.set(callbacks.fetch.callbackId, callbacks.fetch);
+  }
+  if (callbacks?.fs) {
+    for (const [name, reg] of Object.entries(callbacks.fs)) {
+      if (reg) {
+        instance.callbacks.set(reg.callbackId, { ...reg, name });
+      }
+    }
+  }
+  if (callbacks?.moduleLoader) {
+    instance.moduleLoaderCallbackId = callbacks.moduleLoader.callbackId;
+    instance.callbacks.set(callbacks.moduleLoader.callbackId, callbacks.moduleLoader);
+  }
+  if (callbacks?.custom) {
+    for (const [name, reg] of Object.entries(callbacks.custom)) {
+      if (reg) {
+        instance.callbacks.set(reg.callbackId, { ...reg, name });
+      }
+    }
+  }
+  instance.returnedCallbacks = new Map;
+  instance.returnedPromises = new Map;
+  instance.returnedIterators = new Map;
+  instance.nextLocalCallbackId = 1e6;
+  if (callbacks?.custom) {
+    const newCallbackIdMap = {};
+    for (const [name, reg] of Object.entries(callbacks.custom)) {
+      if (reg) {
+        newCallbackIdMap[name] = reg.callbackId;
+      }
+    }
+    try {
+      instance.runtime.context.global.setSync("__customFnCallbackIds", new ivm.ExternalCopy(newCallbackIdMap).copyInto());
+    } catch {}
+  }
+}
+function evictOldestDisposedRuntime(state) {
+  let oldest = null;
+  let oldestTime = Infinity;
+  for (const [, instance] of state.isolates) {
+    if (instance.isDisposed && instance.disposedAt !== undefined) {
+      if (instance.disposedAt < oldestTime) {
+        oldestTime = instance.disposedAt;
+        oldest = instance;
+      }
+    }
+  }
+  if (oldest) {
+    try {
+      if (oldest.playwrightHandle) {
+        oldest.playwrightHandle.dispose();
+      }
+      oldest.runtime.dispose();
+    } catch {}
+    state.isolates.delete(oldest.isolateId);
+    if (oldest.namespaceId != null) {
+      state.namespacedRuntimes.delete(oldest.namespaceId);
+    }
+    return true;
+  }
+  return false;
+}
 async function handleCreateRuntime(message, connection, state) {
+  const namespaceId = message.options.namespaceId;
+  if (namespaceId != null) {
+    const existing = state.namespacedRuntimes.get(namespaceId);
+    if (existing) {
+      if (!existing.isDisposed) {
+        if (existing.ownerConnection === connection.socket) {
+          sendOk(connection.socket, message.requestId, {
+            isolateId: existing.isolateId,
+            reused: true
+          });
+          return;
+        }
+        sendError(connection.socket, message.requestId, ErrorCode.SCRIPT_ERROR, `Namespace "${namespaceId}" already has an active runtime`);
+        return;
+      }
+      reuseNamespacedRuntime(existing, connection, message, state);
+      sendOk(connection.socket, message.requestId, {
+        isolateId: existing.isolateId,
+        reused: true
+      });
+      return;
+    }
+  }
   if (state.isolates.size >= state.options.maxIsolates) {
-    sendError(connection.socket, message.requestId, ErrorCode.ISOLATE_MEMORY_LIMIT, `Maximum isolates (${state.options.maxIsolates}) reached`);
-    return;
+    if (!evictOldestDisposedRuntime(state)) {
+      sendError(connection.socket, message.requestId, ErrorCode.ISOLATE_MEMORY_LIMIT, `Maximum isolates (${state.options.maxIsolates}) reached`);
+      return;
+    }
   }
   try {
     const isolateId = randomUUID();
@@ -206,32 +350,61 @@ async function handleCreateRuntime(message, connection, state) {
     const moduleLoaderCallback = message.options.callbacks?.moduleLoader;
     const customCallbacks = message.options.callbacks?.custom;
     const pendingCallbacks = [];
+    const callbackContext = {
+      connection,
+      consoleOnEntry: consoleCallbacks?.onEntry?.callbackId,
+      fetch: fetchCallback?.callbackId,
+      moduleLoader: moduleLoaderCallback?.callbackId,
+      fs: {
+        readFile: fsCallbacks?.readFile?.callbackId,
+        writeFile: fsCallbacks?.writeFile?.callbackId,
+        stat: fsCallbacks?.stat?.callbackId,
+        readdir: fsCallbacks?.readdir?.callbackId,
+        unlink: fsCallbacks?.unlink?.callbackId,
+        mkdir: fsCallbacks?.mkdir?.callbackId,
+        rmdir: fsCallbacks?.rmdir?.callbackId
+      },
+      custom: new Map(customCallbacks ? Object.entries(customCallbacks).map(([name, reg]) => [name, reg.callbackId]) : [])
+    };
     const runtime = await createInternalRuntime({
       memoryLimitMB: message.options.memoryLimitMB ?? state.options.defaultMemoryLimitMB,
       cwd: message.options.cwd,
-      console: consoleCallbacks?.onEntry ? {
+      console: {
         onEntry: (entry) => {
-          const promise = invokeClientCallback(connection, consoleCallbacks.onEntry.callbackId, [entry]).catch(() => {});
+          const conn = callbackContext.connection;
+          const callbackId = callbackContext.consoleOnEntry;
+          if (!conn || callbackId === undefined)
+            return;
+          const promise = invokeClientCallback(conn, callbackId, [entry]).catch(() => {});
           pendingCallbacks.push(promise);
         }
-      } : undefined,
-      fetch: fetchCallback ? {
+      },
+      fetch: {
         onFetch: async (request) => {
+          const conn = callbackContext.connection;
+          const callbackId = callbackContext.fetch;
+          if (!conn || callbackId === undefined) {
+            throw new Error("Fetch callback not available");
+          }
           const serialized = await serializeRequest(request);
-          const result = await invokeClientCallback(connection, fetchCallback.callbackId, [serialized]);
+          const result = await invokeClientCallback(conn, callbackId, [serialized]);
           return deserializeResponse(result);
         }
-      } : undefined,
-      fs: fsCallbacks ? {
+      },
+      fs: {
         getDirectory: async (path) => {
+          const conn = callbackContext.connection;
+          if (!conn) {
+            throw new Error("FS callbacks not available");
+          }
           return createCallbackFileSystemHandler({
-            connection,
-            callbacks: fsCallbacks,
+            connection: conn,
+            callbackContext,
             invokeClientCallback,
             basePath: path
           });
         }
-      } : undefined
+      }
     });
     const instance = {
       isolateId,
@@ -244,7 +417,10 @@ async function handleCreateRuntime(message, connection, state) {
       returnedCallbacks: new Map,
       returnedPromises: new Map,
       returnedIterators: new Map,
-      nextLocalCallbackId: 1e6
+      nextLocalCallbackId: 1e6,
+      namespaceId,
+      isDisposed: false,
+      callbackContext
     };
     if (moduleLoaderCallback) {
       instance.moduleLoaderCallbackId = moduleLoaderCallback.callbackId;
@@ -329,6 +505,9 @@ async function handleCreateRuntime(message, connection, state) {
     state.isolates.set(isolateId, instance);
     connection.isolates.add(isolateId);
     state.stats.totalIsolatesCreated++;
+    if (namespaceId != null) {
+      state.namespacedRuntimes.set(namespaceId, instance);
+    }
     instance.runtime.fetch.onWebSocketCommand((cmd) => {
       let data;
       if (cmd.data instanceof ArrayBuffer) {
@@ -349,7 +528,7 @@ async function handleCreateRuntime(message, connection, state) {
       };
       sendMessage(connection.socket, wsCommandMsg);
     });
-    sendOk(connection.socket, message.requestId, { isolateId });
+    sendOk(connection.socket, message.requestId, { isolateId, reused: false });
   } catch (err) {
     const error = err;
     sendError(connection.socket, message.requestId, ErrorCode.SCRIPT_ERROR, error.message, { name: error.name, stack: error.stack });
@@ -366,12 +545,16 @@ async function handleDisposeRuntime(message, connection, state) {
     return;
   }
   try {
-    if (instance.playwrightHandle) {
-      instance.playwrightHandle.dispose();
-    }
-    instance.runtime.dispose();
-    state.isolates.delete(message.isolateId);
     connection.isolates.delete(message.isolateId);
+    if (instance.namespaceId != null) {
+      softDeleteRuntime(instance, state);
+    } else {
+      if (instance.playwrightHandle) {
+        instance.playwrightHandle.dispose();
+      }
+      instance.runtime.dispose();
+      state.isolates.delete(message.isolateId);
+    }
     sendOk(connection.socket, message.requestId);
   } catch (err) {
     const error = err;
@@ -428,22 +611,21 @@ async function handleDispatchRequest(message, connection, state) {
       body: requestBody
     });
     const response = await instance.runtime.fetch.dispatchRequest(request);
-    const contentLength = response.headers.get("content-length");
-    const knownSize = contentLength ? parseInt(contentLength, 10) : null;
-    if (knownSize !== null && knownSize > STREAM_THRESHOLD) {
+    if (response.body) {
       await sendStreamedResponse(connection, message.requestId, response);
     } else {
-      const clonedResponse = response.clone();
-      try {
-        const serialized = await serializeResponse(response);
-        if (serialized.body && serialized.body.length > STREAM_THRESHOLD) {
-          await sendStreamedResponse(connection, message.requestId, clonedResponse);
-        } else {
-          sendOk(connection.socket, message.requestId, { response: serialized });
+      const headers = [];
+      response.headers.forEach((value, key) => {
+        headers.push([key, value]);
+      });
+      sendOk(connection.socket, message.requestId, {
+        response: {
+          status: response.status,
+          statusText: response.statusText,
+          headers,
+          body: null
         }
-      } catch {
-        await sendStreamedResponse(connection, message.requestId, clonedResponse);
-      }
+      });
     }
   } catch (err) {
     const error = err;
@@ -1018,7 +1200,8 @@ async function setupCustomFunctions(context, customCallbacks, connection, instan
         }
         result = await callback(...args);
       } else {
-        result = await invokeClientCallback(connection, callbackId, args);
+        const conn = instance.callbackContext?.connection || connection;
+        result = await invokeClientCallback(conn, callbackId, args);
       }
       const ctx = createMarshalContext();
       const marshalledResult = await marshalValue({ ok: true, value: result }, ctx);
@@ -1034,6 +1217,11 @@ async function setupCustomFunctions(context, customCallbacks, connection, instan
   });
   global.setSync("__customFn_invoke", invokeCallbackRef);
   context.evalSync(ISOLATE_MARSHAL_CODE);
+  const callbackIdMap = {};
+  for (const [name, registration] of Object.entries(customCallbacks)) {
+    callbackIdMap[name] = registration.callbackId;
+  }
+  global.setSync("__customFnCallbackIds", new ivm.ExternalCopy(callbackIdMap).copyInto());
   for (const [name, registration] of Object.entries(customCallbacks)) {
     if (name.includes(":")) {
       continue;
@@ -1041,10 +1229,11 @@ async function setupCustomFunctions(context, customCallbacks, connection, instan
     if (registration.type === "sync") {
       context.evalSync(`
         globalThis.${name} = function(...args) {
+          const callbackId = globalThis.__customFnCallbackIds["${name}"];
           const argsJson = JSON.stringify(__marshalForHost(args));
           const resultJson = __customFn_invoke.applySyncPromise(
             undefined,
-            [${registration.callbackId}, argsJson]
+            [callbackId, argsJson]
           );
           const result = JSON.parse(resultJson);
           if (result.ok) {
@@ -1067,10 +1256,11 @@ async function setupCustomFunctions(context, customCallbacks, connection, instan
       context.evalSync(`
         globalThis.${name} = function(...args) {
           // Start the iterator and get the iteratorId
+          const startCallbackId = globalThis.__customFnCallbackIds["${name}:start"];
           const argsJson = JSON.stringify(__marshalForHost(args));
           const startResultJson = __customFn_invoke.applySyncPromise(
             undefined,
-            [${startReg.callbackId}, argsJson]
+            [startCallbackId, argsJson]
           );
           const startResult = JSON.parse(startResultJson);
           if (!startResult.ok) {
@@ -1083,10 +1273,11 @@ async function setupCustomFunctions(context, customCallbacks, connection, instan
           return {
             [Symbol.asyncIterator]() { return this; },
             async next() {
+              const nextCallbackId = globalThis.__customFnCallbackIds["${name}:next"];
               const argsJson = JSON.stringify(__marshalForHost([iteratorId]));
               const resultJson = __customFn_invoke.applySyncPromise(
                 undefined,
-                [${nextReg.callbackId}, argsJson]
+                [nextCallbackId, argsJson]
               );
               const result = JSON.parse(resultJson);
               if (!result.ok) {
@@ -1098,19 +1289,21 @@ async function setupCustomFunctions(context, customCallbacks, connection, instan
               return { done: val.done, value: val.value };
             },
             async return(v) {
+              const returnCallbackId = globalThis.__customFnCallbackIds["${name}:return"];
               const argsJson = JSON.stringify(__marshalForHost([iteratorId, v]));
               const resultJson = __customFn_invoke.applySyncPromise(
                 undefined,
-                [${returnReg.callbackId}, argsJson]
+                [returnCallbackId, argsJson]
               );
               const result = JSON.parse(resultJson);
               return { done: true, value: result.ok ? __unmarshalFromHost(result.value) : undefined };
             },
             async throw(e) {
+              const throwCallbackId = globalThis.__customFnCallbackIds["${name}:throw"];
               const argsJson = JSON.stringify(__marshalForHost([iteratorId, { message: e?.message, name: e?.name }]));
               const resultJson = __customFn_invoke.applySyncPromise(
                 undefined,
-                [${throwReg.callbackId}, argsJson]
+                [throwCallbackId, argsJson]
               );
               const result = JSON.parse(resultJson);
               if (!result.ok) {
@@ -1127,10 +1320,11 @@ async function setupCustomFunctions(context, customCallbacks, connection, instan
     } else if (registration.type === "async") {
       context.evalSync(`
         globalThis.${name} = async function(...args) {
+          const callbackId = globalThis.__customFnCallbackIds["${name}"];
           const argsJson = JSON.stringify(__marshalForHost(args));
           const resultJson = __customFn_invoke.applySyncPromise(
             undefined,
-            [${registration.callbackId}, argsJson]
+            [callbackId, argsJson]
           );
           const result = JSON.parse(resultJson);
           if (result.ok) {
@@ -1179,22 +1373,6 @@ async function serializeRequest(request) {
     body
   };
 }
-async function serializeResponse(response) {
-  const headers = [];
-  response.headers.forEach((value, key) => {
-    headers.push([key, value]);
-  });
-  let body = null;
-  if (response.body) {
-    body = new Uint8Array(await response.arrayBuffer());
-  }
-  return {
-    status: response.status,
-    statusText: response.statusText,
-    headers,
-    body
-  };
-}
 function deserializeResponse(data) {
   return new Response(data.body, {
     status: data.status,
@@ -1355,14 +1533,21 @@ async function handleRunTests(message, connection, state) {
   instance.lastActivity = Date.now();
   try {
     const timeout = message.timeout ?? 30000;
+    let timeoutId;
     const timeoutPromise = new Promise((_, reject) => {
-      setTimeout(() => reject(new Error("Test timeout")), timeout);
+      timeoutId = setTimeout(() => reject(new Error("Test timeout")), timeout);
     });
-    const results = await Promise.race([
-      runTestsInContext(instance.runtime.context),
-      timeoutPromise
-    ]);
-    sendOk(connection.socket, message.requestId, results);
+    try {
+      const results = await Promise.race([
+        runTestsInContext(instance.runtime.context),
+        timeoutPromise
+      ]);
+      sendOk(connection.socket, message.requestId, results);
+    } finally {
+      if (timeoutId) {
+        clearTimeout(timeoutId);
+      }
+    }
   } catch (err) {
     const error = err;
     sendError(connection.socket, message.requestId, ErrorCode.SCRIPT_ERROR, error.message, { name: error.name, stack: error.stack });
@@ -1477,4 +1662,4 @@ export {
   handleConnection
 };
 
-//# debugId=8FEF67C0571D870E64756E2164756E21
+//# debugId=EC4C1FF56C0445D064756E2164756E21