From e5cca347ae1385c3d6fc78a746fd740d78ea7b9a Mon Sep 17 00:00:00 2001 From: Adrian Lyjak Date: Wed, 24 Apr 2024 18:08:53 -0400 Subject: [PATCH] fix(cloudflare): handle wasm imports shared between ssg and ssr --- .changeset/lovely-mails-change.md | 5 + packages/cloudflare/src/index.ts | 23 ++- .../src/utils/wasm-module-loader.ts | 182 ++++++++++++------ .../fixtures/wasm/src/pages/add/[a]/[b].ts | 10 +- .../test/fixtures/wasm/src/pages/hybrid.ts | 6 +- .../test/fixtures/wasm/src/util/add.ts | 8 + 6 files changed, 165 insertions(+), 69 deletions(-) create mode 100644 .changeset/lovely-mails-change.md create mode 100644 packages/cloudflare/test/fixtures/wasm/src/util/add.ts diff --git a/.changeset/lovely-mails-change.md b/.changeset/lovely-mails-change.md new file mode 100644 index 000000000..9ab89e16f --- /dev/null +++ b/.changeset/lovely-mails-change.md @@ -0,0 +1,5 @@ +--- +'@astrojs/cloudflare': patch +--- + +Fixes build errors when wasm modules are imported from a file that is shared in both prerendered static pages and server side rendered pages diff --git a/packages/cloudflare/src/index.ts b/packages/cloudflare/src/index.ts index a575b9b0a..9d373f719 100644 --- a/packages/cloudflare/src/index.ts +++ b/packages/cloudflare/src/index.ts @@ -13,12 +13,16 @@ import { createRedirectsFromAstroRoutes } from '@astrojs/underscore-redirects'; import { AstroError } from 'astro/errors'; import { walk } from 'estree-walker'; import MagicString from 'magic-string'; +import type { PluginOption } from 'vite'; import { getPlatformProxy } from 'wrangler'; import { createRoutesFile, getParts } from './utils/generate-routes-json.js'; import { setImageConfig } from './utils/image-config.js'; import { mutateDynamicPageImportsInPlace, mutatePageMapInPlace } from './utils/index.js'; import { NonServerChunkDetector } from './utils/non-server-chunk-detector.js'; -import { wasmModuleLoader } from './utils/wasm-module-loader.js'; +import { + type CloudflareModulePluginExtra, + cloudflareModuleLoader, +} from './utils/wasm-module-loader.js'; export type { Runtime } from './entrypoints/server.js'; @@ -62,13 +66,23 @@ export type Options = { /** Configuration persistence settings. Default '.wrangler/state/v3' */ persist?: boolean | { path: string }; }; - /** Enable WebAssembly support */ + /** + * Allow bundling cloudflare worker specific file types + * https://developers.cloudflare.com/workers/wrangler/bundling/ + */ wasmModuleImports?: boolean; }; export default function createIntegration(args?: Options): AstroIntegration { let _config: AstroConfig; + const cloudflareModulePlugin: PluginOption & CloudflareModulePluginExtra = cloudflareModuleLoader( + { + wasm: args?.wasmModuleImports ?? false, + bin: false, + } + ); + // Initialize the unused chunk analyzer as a shared state between hooks. // The analyzer is used on earlier hooks to collect information about used hooks on a Vite plugin // and then later after the full build to clean up unused chunks, so it has to be shared between them. @@ -91,9 +105,7 @@ export default function createIntegration(args?: Options): AstroIntegration { vite: { // load .wasm files as WebAssembly modules plugins: [ - wasmModuleLoader({ - disabled: !args?.wasmModuleImports, - }), + cloudflareModulePlugin, chunkAnalyzer.getPlugin(), { name: 'dynamic-imports-analyzer', @@ -274,6 +286,7 @@ export default function createIntegration(args?: Options): AstroIntegration { } }, 'astro:build:done': async ({ pages, routes, dir, logger }) => { + await cloudflareModulePlugin.afterBuildCompleted(_config); const PLATFORM_FILES = ['_headers', '_redirects', '_routes.json']; if (_config.base !== '/') { for (const file of PLATFORM_FILES) { diff --git a/packages/cloudflare/src/utils/wasm-module-loader.ts b/packages/cloudflare/src/utils/wasm-module-loader.ts index dd61f501a..a734b06a2 100644 --- a/packages/cloudflare/src/utils/wasm-module-loader.ts +++ b/packages/cloudflare/src/utils/wasm-module-loader.ts @@ -1,23 +1,35 @@ -import * as fs from 'node:fs'; +import * as fs from 'node:fs/promises'; import * as path from 'node:path'; +import * as url from 'node:url'; import type { AstroConfig } from 'astro'; +import type { PluginOption } from 'vite'; +export interface CloudflareModulePluginExtra { + afterBuildCompleted(config: AstroConfig): Promise; +} /** + * Enables support for various non-standard extensions in module imports within cloudflare workers. + * + * See https://developers.cloudflare.com/workers/wrangler/bundling/ for reference + * + * This adds supports for imports in the following formats: + * - .wasm?module + * - .bin + * * Loads '*.wasm?module' imports as WebAssembly modules, which is the only way to load WASM in cloudflare workers. * Current proposal for WASM modules: https://github.com/WebAssembly/esm-integration/tree/main/proposals/esm-integration * Cloudflare worker WASM from javascript support: https://developers.cloudflare.com/workers/runtime-apis/webassembly/javascript/ - * @param disabled - if true throws a helpful error message if wasm is encountered and wasm imports are not enabled, - * otherwise it will error obscurely in the esbuild and vite builds - * @param assetsDirectory - the folder name for the assets directory in the build directory. Usually '_astro' + * @param bin - if true, will load '.bin' imports as Uint8Arrays, otherwise will throw errors when encountered to clarify that it must be enabled + * @param wasm - if true, will load '.wasm?module' imports as Uint8Arrays, otherwise will throw errors when encountered to clarify that it must be enabled * @returns Vite plugin to load WASM tagged with '?module' as a WASM modules */ -export function wasmModuleLoader({ - disabled, -}: { - disabled: boolean; -}): NonNullable[number] { - const postfix = '.wasm?module'; +export function cloudflareModuleLoader( + enabled: Record +): PluginOption & CloudflareModulePluginExtra { + const enabledAdapters = cloudflareImportAdapters.filter((x) => enabled[x.extension]); let isDev = false; + const MAGIC_STRING = '__CLOUDFLARE_ASSET__'; + const replacements: Replacement[] = []; return { name: 'vite:wasm-module-loader', @@ -28,46 +40,54 @@ export function wasmModuleLoader({ config(_, __) { // let vite know that file format and the magic import string is intentional, and will be handled in this plugin return { - assetsInclude: ['**/*.wasm?module'], + assetsInclude: enabledAdapters.map((x) => `**/*.${x.qualifiedExtension}`), build: { rollupOptions: { // mark the wasm files as external so that they are not bundled and instead are loaded from the files - external: [/^__WASM_ASSET__.+\.wasm$/i, /^__WASM_ASSET__.+\.wasm.mjs$/i], + external: enabledAdapters.map( + (x) => new RegExp(`^${MAGIC_STRING}.+\\.${x.extension}.mjs$`, 'i') + ), }, }, }; }, - load(id, _) { - if (!id.endsWith(postfix)) { + async load(id, _) { + const suffix = id.split('.').at(-1); + const importAdapter = cloudflareImportAdapters.find((x) => x.qualifiedExtension === suffix); + if (!importAdapter) { return; } - if (disabled) { + const suffixType: ImportType = importAdapter.extension; + const adapterEnabled = enabled[suffixType]; + if (!adapterEnabled) { throw new Error( - `WASM module's cannot be loaded unless you add \`wasmModuleImports: true\` to your astro config.` + `Cloudflare module loading is experimental. The ${suffix} module cannot be loaded unless you add \`wasmModuleImports: true\` to your astro config.` ); } - const filePath = id.slice(0, -1 * '?module'.length); + const filePath = id.replace(/\?module$/, ''); - const data = fs.readFileSync(filePath); + const data = await fs.readFile(filePath); const base64 = data.toString('base64'); - const base64Module = `const wasmModule = new WebAssembly.Module(Uint8Array.from(atob("${base64}"), c => c.charCodeAt(0)));export default wasmModule;`; + const inlineModule = importAdapter.asNodeModule(data); if (isDev) { // no need to wire up the assets in dev mode, just rewrite - return base64Module; + return inlineModule; } // just some shared ID const hash = hashString(base64); // emit the wasm binary as an asset file, to be picked up later by the esbuild bundle for the worker. // give it a shared deterministic name to make things easy for esbuild to switch on later - const assetName = `${path.basename(filePath).split('.')[0]}.${hash}.wasm`; + const assetName = `${path.basename(filePath).split('.')[0]}.${hash}.${ + importAdapter.extension + }`; this.emitFile({ type: 'asset', - // put it explicitly in the _astro assets directory with `fileName` rather than `name` so that - // vite doesn't give it a random id in its name. We need to be able to easily rewrite from + // emit the data explicitly as an esset with `fileName` rather than `name` so that + // vite doesn't give it a random hash-id in its name--We need to be able to easily rewrite from // the .mjs loader and the actual wasm asset later in the ESbuild for the worker fileName: assetName, source: data, @@ -77,51 +97,105 @@ export function wasmModuleLoader({ const chunkId = this.emitFile({ type: 'prebuilt-chunk', fileName: `${assetName}.mjs`, - code: base64Module, + code: inlineModule, }); - return `import wasmModule from "__WASM_ASSET__${chunkId}.wasm.mjs";export default wasmModule;`; + return `import module from "${MAGIC_STRING}${chunkId}.${importAdapter.extension}.mjs";export default module;`; }, - // output original wasm file relative to the chunk + // output original wasm file relative to the chunk now that chunking has been achieved renderChunk(code, chunk, _) { if (isDev) return; - if (!/__WASM_ASSET__/g.test(code)) return; - - const isPrerendered = Object.keys(chunk.modules).some( - (moduleId) => this.getModuleInfo(moduleId)?.meta?.astro?.pageOptions?.prerender === true - ); - - let final = code; - - // SSR - if (!isPrerendered) { - final = code.replaceAll(/__WASM_ASSET__([A-Za-z\d]+).wasm.mjs/g, (s, assetId) => { - const fileName = this.getFileName(assetId).replace(/\.mjs$/, ''); - const relativePath = path - .relative(path.dirname(chunk.fileName), fileName) - .replaceAll('\\', '/'); // fix windows paths for import - return `./${relativePath}`; - }); + if (!code.includes(MAGIC_STRING)) return; + + // SSR will need the .mjs suffix removed from the import before this works in cloudflare, but this is done as a final step + // so as to support prerendering from nodejs runtime + let replaced = code; + for (const loader of enabledAdapters) { + replaced = replaced.replaceAll( + new RegExp(`${MAGIC_STRING}([A-Za-z\\d]+)\\.${loader.extension}\\.mjs`, 'g'), + (s, assetId) => { + const fileName = this.getFileName(assetId); + const relativePath = path + .relative(path.dirname(chunk.fileName), fileName) + .replaceAll('\\', '/'); // fix windows paths for import + + // record this replacement for later, to adjust it to import the unbundled asset + replacements.push({ + fileName, + cloudflareImport: relativePath.replace(/\.mjs$/, ''), + nodejsImport: relativePath, + }); + return `./${relativePath}`; + } + ); } - - // SSG - if (isPrerendered) { - final = code.replaceAll(/__WASM_ASSET__([A-Za-z\d]+).wasm.mjs/g, (s, assetId) => { - const fileName = this.getFileName(assetId); - const relativePath = path - .relative(path.dirname(chunk.fileName), fileName) - .replaceAll('\\', '/'); // fix windows paths for import - return `./${relativePath}`; - }); + if (replaced.includes(MAGIC_STRING)) { + console.error('failed to replace', replaced); } - return { code: final }; + return { code: replaced }; + }, + + /** + * Once prerendering is complete, restore the imports in the _worker.js to cloudflare compatible ones, removing the .mjs suffix + */ + async afterBuildCompleted(config: AstroConfig) { + for (const replacement of replacements) { + const filepath = path.join( + url.fileURLToPath(config.outDir), + '_worker.js', + replacement.fileName + ); + const contents = await fs.readFile(filepath, 'utf-8'); + const newContents = contents.replaceAll( + replacement.cloudflareImport, + replacement.nodejsImport + ); + await fs.writeFile(filepath, newContents, 'utf-8'); + } }, }; } +export type ImportType = 'wasm' | 'bin'; + +interface Replacement { + // path relative to the build root (_workers.js/ in this case) + fileName: string; + // desired import for cloudflare + cloudflareImport: string; + // nodejs import that simulates a wasm/bin module + nodejsImport: string; +} + +interface ModuleImportAdapter { + extension: ImportType; + qualifiedExtension: string; + asNodeModule(fileContents: Buffer): string; +} + +const wasmImportAdapter: ModuleImportAdapter = { + extension: 'wasm', + qualifiedExtension: 'wasm?module', + asNodeModule(fileContents: Buffer) { + const base64 = fileContents.toString('base64'); + return `const wasmModule = new WebAssembly.Module(Uint8Array.from(atob("${base64}"), c => c.charCodeAt(0)));export default wasmModule;`; + }, +}; + +const binImportAdapter: ModuleImportAdapter = { + extension: 'bin', + qualifiedExtension: 'bin', + asNodeModule(fileContents: Buffer) { + const base64 = fileContents.toString('base64'); + return `const binModule = Uint8Array.from(atob("${base64}"), c => c.charCodeAt(0)).buffer;export default binModule;`; + }, +}; + +const cloudflareImportAdapters = [binImportAdapter, wasmImportAdapter]; + /** * Returns a deterministic 32 bit hash code from a string */ diff --git a/packages/cloudflare/test/fixtures/wasm/src/pages/add/[a]/[b].ts b/packages/cloudflare/test/fixtures/wasm/src/pages/add/[a]/[b].ts index 0f09faa0a..9a8283d85 100644 --- a/packages/cloudflare/test/fixtures/wasm/src/pages/add/[a]/[b].ts +++ b/packages/cloudflare/test/fixtures/wasm/src/pages/add/[a]/[b].ts @@ -1,17 +1,15 @@ import { type APIContext } from 'astro'; -// @ts-ignore -import mod from '../../../util/add.wasm?module'; +import { add } from '../../../util/add'; -const addModule: any = new WebAssembly.Instance(mod); export const prerender = false; export async function GET( context: APIContext ): Promise { - const a = Number.parseInt(context.params.a!); - const b = Number.parseInt(context.params.b!); - return new Response(JSON.stringify({ answer: addModule.exports.add(a, b) }), { + const a = Number.parseInt(context.params.a ?? "0"); + const b = Number.parseInt(context.params.b ?? "0"); + return new Response(JSON.stringify({ answer: add(a, b) }), { status: 200, headers: { 'Content-Type': 'application/json', diff --git a/packages/cloudflare/test/fixtures/wasm/src/pages/hybrid.ts b/packages/cloudflare/test/fixtures/wasm/src/pages/hybrid.ts index f1c3a0d57..408dd1517 100644 --- a/packages/cloudflare/test/fixtures/wasm/src/pages/hybrid.ts +++ b/packages/cloudflare/test/fixtures/wasm/src/pages/hybrid.ts @@ -1,13 +1,11 @@ import { type APIContext } from 'astro'; -// @ts-ignore -import mod from '../util/add.wasm?module'; -const addModule: any = new WebAssembly.Instance(mod); +import {add} from '../util/add'; export async function GET( context: APIContext ): Promise { - return new Response(JSON.stringify({ answer: addModule.exports.add(20, 1) }), { + return new Response(JSON.stringify({ answer: add(20, 1) }), { status: 200, headers: { 'Content-Type': 'application/json', diff --git a/packages/cloudflare/test/fixtures/wasm/src/util/add.ts b/packages/cloudflare/test/fixtures/wasm/src/util/add.ts new file mode 100644 index 000000000..2a865574d --- /dev/null +++ b/packages/cloudflare/test/fixtures/wasm/src/util/add.ts @@ -0,0 +1,8 @@ +import mod from './add.wasm?module'; + + +const addModule: any = new WebAssembly.Instance(mod); + +export function add(a, b) { + return addModule.exports.add(a, b); +}