feat(workerd): use new transport (#144)

hi-ogawa · Nov 7, 2024 · 633b74d · 633b74d
1 parent b5c4861
commit 633b74d
Show file tree

Hide file tree

Showing 12 changed files with 53 additions and 139 deletions.
diff --git a/examples/browser-cli/src/cli.ts b/examples/browser-cli/src/cli.ts
@@ -7,7 +7,8 @@ import {
   createServer,
   parseAstAsync,
 } from "vite";
-import type { FetchFunction, ModuleRunner } from "vite/module-runner";
+import type { ModuleRunner } from "vite/module-runner";
+import { vitePluginFetchModuleServer } from "../../web-worker/src/lib/fetch-module-server";
 
 const headless = !process.env["CLI_HEADED"];
 const extension = process.env["CLI_EXTENSION"] ?? "tsx";
@@ -185,25 +186,4 @@ function vitePluginBrowserRunner(): Plugin {
   };
 }
 
-// https://github.com/vitejs/vite/discussions/18191
-function vitePluginFetchModuleServer(): Plugin {
-  return {
-    name: vitePluginFetchModuleServer.name,
-    configureServer(server) {
-      server.middlewares.use(async (req, res, next) => {
-        const url = new URL(req.url ?? "/", "https://any.local");
-        if (url.pathname === "/@vite/fetchModule") {
-          const [name, ...args] = JSON.parse(url.searchParams.get("payload")!);
-          const result = await server.environments[name]!.fetchModule(
-            ...(args as Parameters<FetchFunction>),
-          );
-          res.end(JSON.stringify(result));
-          return;
-        }
-        next();
-      });
-    },
-  };
-}
-
 main();
diff --git a/examples/browser-cli/src/runner.ts b/examples/browser-cli/src/runner.ts
@@ -1,16 +1,13 @@
-import {
-  ESModulesEvaluator,
-  type FetchFunction,
-  ModuleRunner,
-} from "vite/module-runner";
+import { ESModulesEvaluator, ModuleRunner } from "vite/module-runner";
+import { fetchClientFetchModule } from "../../web-worker/src/lib/fetch-module-client";
 
 export async function start(options: { root: string }) {
   const runner = new ModuleRunner(
     {
       root: options.root,
       sourcemapInterceptor: false,
       transport: {
-        fetchModule: fetchModuleFetchClient("custom"),
+        invoke: fetchClientFetchModule("custom"),
       },
       hmr: false,
     },
@@ -19,15 +16,3 @@ export async function start(options: { root: string }) {
 
   return runner;
 }
-
-// https://github.com/vitejs/vite/discussions/18191
-function fetchModuleFetchClient(environmentName: string): FetchFunction {
-  return async (...args) => {
-    const payload = JSON.stringify([environmentName, ...args]);
-    const response = await fetch(
-      "/@vite/fetchModule?" + new URLSearchParams({ payload }),
-    );
-    const result = response.json();
-    return result as any;
-  };
-}
diff --git a/examples/web-worker-rsc/src/lib/fetch-module-client.ts b/examples/web-worker-rsc/src/lib/fetch-module-client.ts
diff --git a/examples/web-worker-rsc/src/lib/fetch-module-server.ts b/examples/web-worker-rsc/src/lib/fetch-module-server.ts
diff --git a/examples/web-worker-rsc/src/lib/runner.ts b/examples/web-worker-rsc/src/lib/runner.ts
@@ -1,5 +1,5 @@
 import { ESModulesEvaluator, ModuleRunner } from "vite/module-runner";
-import { fetchClientFetchModule } from "./fetch-module-client";
+import { fetchClientFetchModule } from "../../../web-worker/src/lib/fetch-module-client";
 
 export function createFetchRunner(options: {
   root: string;
@@ -10,7 +10,7 @@ export function createFetchRunner(options: {
       root: options.root,
       sourcemapInterceptor: false,
       transport: {
-        fetchModule: fetchClientFetchModule(options.environmentName),
+        invoke: fetchClientFetchModule(options.environmentName),
       },
       hmr: false,
     },

diff --git a/examples/web-worker-rsc/vite.config.ts b/examples/web-worker-rsc/vite.config.ts
@@ -1,7 +1,7 @@
 import react from "@vitejs/plugin-react";
 import { defineConfig } from "vite";
+import { vitePluginFetchModuleServer } from "../web-worker/src/lib/fetch-module-server";
 import { vitePluginWorkerEnvironment } from "../web-worker/vite.config";
-import { vitePluginFetchModuleServer } from "./src/lib/fetch-module-server";
 
 export default defineConfig((_env) => ({
   clearScreen: false,

diff --git a/examples/web-worker/src/lib/fetch-module-client.ts b/examples/web-worker/src/lib/fetch-module-client.ts
@@ -1,10 +1,12 @@
-import type { FetchFunction } from "vite";
+import type { ModuleRunnerTransport } from "vite/module-runner";
 
-export function fetchClientFetchModule(environmentName: string): FetchFunction {
-  return async (...args) => {
-    const payload = JSON.stringify([environmentName, ...args]);
+export function fetchClientFetchModule(
+  environmentName: string,
+): ModuleRunnerTransport["invoke"] {
+  return async (payload) => {
+    const data = JSON.stringify([environmentName, payload]);
     const response = await fetch(
-      "/@vite/fetchModule?" + new URLSearchParams({ payload }),
+      "/@vite/invoke?" + new URLSearchParams({ data }),
     );
     const result = response.json();
     return result as any;

diff --git a/examples/web-worker/src/lib/fetch-module-server.ts b/examples/web-worker/src/lib/fetch-module-server.ts
@@ -1,16 +1,15 @@
-import type { FetchFunction, Plugin } from "vite";
+import { type Plugin } from "vite";
 
 export function vitePluginFetchModuleServer(): Plugin {
   return {
     name: vitePluginFetchModuleServer.name,
     configureServer(server) {
       server.middlewares.use(async (req, res, next) => {
         const url = new URL(req.url ?? "/", "https://any.local");
-        if (url.pathname === "/@vite/fetchModule") {
-          const [name, ...args] = JSON.parse(url.searchParams.get("payload")!);
-          const result = await server.environments[name]!.fetchModule(
-            ...(args as Parameters<FetchFunction>),
-          );
+        if (url.pathname === "/@vite/invoke") {
+          const [name, payload] = JSON.parse(url.searchParams.get("data")!);
+          const devEnv = server.environments[name]!;
+          const result = await devEnv.hot.handleInvoke(payload);
           res.end(JSON.stringify(result));
           return;
         }

diff --git a/examples/web-worker/src/lib/runner.ts b/examples/web-worker/src/lib/runner.ts
@@ -10,7 +10,7 @@ export function createFetchRunner(options: {
       root: options.root,
       sourcemapInterceptor: false,
       transport: {
-        fetchModule: fetchClientFetchModule(options.environmentName),
+        invoke: fetchClientFetchModule(options.environmentName),
       },
       hmr: false,
     },

diff --git a/packages/workerd/src/plugin.ts b/packages/workerd/src/plugin.ts
@@ -123,19 +123,14 @@ export async function createWorkerdDevEnvironment(
     },
     unsafeEvalBinding: "__viteUnsafeEval",
     serviceBindings: {
-      __viteFetchModule: async (request) => {
-        const args = await request.json();
-        try {
-          const result = await devEnv.fetchModule(...(args as [any, any]));
-          return MiniflareResponse.json(result);
-        } catch (error) {
-          console.error("[fetchModule]", args, error);
-          throw error;
-        }
+      __viteInvoke: async (request) => {
+        const payload = (await request.json()) as HotPayload;
+        const result = await devEnv.hot.handleInvoke(payload);
+        return MiniflareResponse.json(result);
       },
       __viteRunnerSend: async (request) => {
         const payload = (await request.json()) as HotPayload;
-        hotListener.dispatch(payload);
+        hotListener.dispatch(payload, { send: runnerObject.__viteServerSend });
         return MiniflareResponse.json(null);
       },
     },
@@ -179,15 +174,12 @@ export async function createWorkerdDevEnvironment(
 
   // hmr channel
   const hotListener = createHotListenerManager();
-  const hot: HotChannel = {
+  const transport: HotChannel = {
     listen: () => {},
     close: () => {},
     on: hotListener.on,
     off: hotListener.off,
-    send(...args: any[]) {
-      const payload = normalizeServerSendPayload(...args);
-      runnerObject.__viteServerSend(payload);
-    },
+    send: runnerObject.__viteServerSend,
   };
 
   // TODO: move initialization code to `init`?
@@ -199,7 +191,10 @@ export async function createWorkerdDevEnvironment(
     }
   }
 
-  const devEnv = new WorkerdDevEnvironmentImpl(name, config, { hot });
+  const devEnv = new WorkerdDevEnvironmentImpl(name, config, {
+    transport,
+    hot: true,
+  });
 
   // custom environment api
   const api: WorkerdDevApi = {
@@ -231,7 +226,10 @@ export async function createWorkerdDevEnvironment(
 
 // wrapper to simplify listener management
 function createHotListenerManager(): Pick<HotChannel, "on" | "off"> & {
-  dispatch: (payload: HotPayload) => void;
+  dispatch: (
+    payload: HotPayload,
+    client: { send: (payload: HotPayload) => void },
+  ) => void;
 } {
   const listerMap: Record<string, Set<Function>> = {};
   const getListerMap = (e: string) => (listerMap[e] ??= new Set());
@@ -243,26 +241,12 @@ function createHotListenerManager(): Pick<HotChannel, "on" | "off"> & {
     off(event, listener: any) {
       getListerMap(event).delete(listener);
     },
-    dispatch(payload) {
+    dispatch(payload, client) {
       if (payload.type === "custom") {
         for (const lister of getListerMap(payload.event)) {
-          lister(payload.data);
+          lister(payload.data, client);
         }
       }
     },
   };
 }
-
-function normalizeServerSendPayload(...args: any[]) {
-  let payload: HotPayload;
-  if (typeof args[0] === "string") {
-    payload = {
-      type: "custom",
-      event: args[0],
-      data: args[1],
-    };
-  } else {
-    payload = args[0];
-  }
-  return payload;
-}
diff --git a/packages/workerd/src/shared.ts b/packages/workerd/src/shared.ts
@@ -7,7 +7,7 @@ export type RunnerEnv = {
   __viteUnsafeEval: {
     eval: (code: string, filename?: string) => any;
   };
-  __viteFetchModule: {
+  __viteInvoke: {
     fetch: (request: Request) => Promise<Response>;
   };
   __viteRunnerSend: {

diff --git a/packages/workerd/src/worker.ts b/packages/workerd/src/worker.ts
@@ -58,28 +58,25 @@ export class RunnerObject extends DurableObject implements RunnerRpc {
         root: env.__viteRoot,
         sourcemapInterceptor: "prepareStackTrace",
         transport: {
-          fetchModule: async (...args) => {
-            const response = await env.__viteFetchModule.fetch(
-              requestJson(args),
-            );
+          invoke: async (payload) => {
+            // we still need to implement fetchModule on top of service binding
+            // since websocket and rpc have tighter payload size limit
+            // (for example, rpc 1MB is not enough for large pre-bundled deps with source map)
+            const response = await env.__viteInvoke.fetch(requestJson(payload));
             tinyassert(response.ok);
             return response.json();
           },
-        },
-        hmr: {
-          connection: {
-            isReady: () => true,
-            onUpdate: (callback) => {
-              this.#viteServerSendHandler = callback;
-            },
-            send: async (payload) => {
-              const response = await env.__viteRunnerSend.fetch(
-                requestJson(payload),
-              );
-              tinyassert(response.ok);
-            },
+          connect: async (handlers) => {
+            this.#viteServerSendHandler = handlers.onMessage;
+          },
+          send: async (payload) => {
+            const response = await env.__viteRunnerSend.fetch(
+              requestJson(payload),
+            );
+            tinyassert(response.ok);
           },
         },
+        hmr: true,
       },
       {
         runInlinedModule: async (context, transformed) => {