feat(cloudflare): remove unnecessary code from the server output (#222)

2 files changed, 157 insertions, 1 deletions

diff --git a/packages/integrations/cloudflare/src/index.ts b/packages/integrations/cloudflare/src/index.ts
index deed35655..d7a68c3e4 100644
--- a/packages/integrations/cloudflare/src/index.ts
+++ b/packages/integrations/cloudflare/src/index.ts
@@ -1,7 +1,7 @@
 import type { AstroConfig, AstroIntegration, RouteData } from 'astro';
 
 import { createReadStream } from 'node:fs';
-import { appendFile, rename, stat } from 'node:fs/promises';
+import { appendFile, rename, stat, unlink } from 'node:fs/promises';
 import { createInterface } from 'node:readline/promises';
 import {
 	appendForwardSlash,
@@ -13,6 +13,7 @@ import { AstroError } from 'astro/errors';
 import { getPlatformProxy } from 'wrangler';
 import { createRoutesFile, getParts } from './utils/generate-routes-json.js';
 import { setImageConfig } from './utils/image-config.js';
+import { NonServerChunkDetector } from './utils/non-server-chunk-detector.js';
 import { wasmModuleLoader } from './utils/wasm-module-loader.js';
 
 export type { Runtime } from './entrypoints/server.js';
@@ -64,6 +65,13 @@ export type Options = {
 export default function createIntegration(args?: Options): AstroIntegration {
 	let _config: AstroConfig;
 
+	// 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.
+	const chunkAnalyzer = new NonServerChunkDetector();
+
+	const prerenderImports: string[][] = [];
+
 	return {
 		name: '@astrojs/cloudflare',
 		hooks: {
@@ -84,6 +92,63 @@ export default function createIntegration(args?: Options): AstroIntegration {
 							wasmModuleLoader({
 								disabled: !args?.wasmModuleImports,
 							}),
+							chunkAnalyzer.getPlugin(),
+							{
+								name: 'dynamic-imports-analyzer',
+								enforce: 'post',
+								generateBundle(_, bundle) {
+									// Find all pages (ignore the ssr entrypoint) which are prerendered based on the dynamic imports of the prerender chunk
+									for (const chunk of Object.values(bundle)) {
+										if (chunk.type !== 'chunk') continue;
+
+										const isPrerendered = chunk.dynamicImports.some(
+											(entry) =>
+												entry.includes('prerender') && chunk.name !== '_@astrojs-ssr-virtual-entry'
+										);
+										if (isPrerendered) {
+											prerenderImports.push([chunk.facadeModuleId ?? '', chunk.fileName]);
+										}
+									}
+
+									const entryChunk = bundle['index.js'];
+									if (
+										entryChunk &&
+										entryChunk.type === 'chunk' &&
+										entryChunk.name === '_@astrojs-ssr-virtual-entry'
+									) {
+										// Update dynamicImports information, so that there are no imports listed which we remove later
+										entryChunk.dynamicImports = entryChunk.dynamicImports.filter(
+											(entry) => !prerenderImports.map((e) => e[1]).includes(entry)
+										);
+
+										// Clean the ssr entry file from prerendered imporst, since Astro adds them, which it shouldn't. But this is a current limitation in core, because the prerender meta info gets added later in the chain
+										for (const page of prerenderImports) {
+											// Find the dynamic import inside of the ssr entry file, which get generated by Astro: https://github.com/withastro/astro/blob/08cdd0919d3249a762822e4bba9e0c5d3966916c/packages/astro/src/core/build/plugins/plugin-ssr.ts#L56
+											const importRegex = new RegExp(
+												`^const (_page\\d) = \\(\\) => import\\('.\\/${page[1]}'\\);$\\n`,
+												'gm'
+											);
+
+											let pageId: string | undefined;
+											const matches = entryChunk.code.matchAll(importRegex);
+											for (const match of matches) {
+												if (match[1]) {
+													pageId = match[1];
+												}
+											}
+											const pageSource = page[0].split(':')[1].replace('@_@', '.');
+											entryChunk.code = entryChunk.code.replace(importRegex, '');
+											if (pageId) {
+												// Find the page in the pageMap of the ssr entry file, which get generated by Astro: https://github.com/withastro/astro/blob/08cdd0919d3249a762822e4bba9e0c5d3966916c/packages/astro/src/core/build/plugins/plugin-ssr.ts#L65
+												const arrayRegex = new RegExp(`\\["${pageSource}", ?${pageId}\\],?`, 'gm');
+												entryChunk.code = entryChunk.code.replace(arrayRegex, '');
+											}
+										}
+									} else {
+										// We don't want to handle this case, since it will always occur for the client build.
+									}
+								},
+							},
 						],
 					},
 					image: setImageConfig(args?.imageService ?? 'DEFAULT', config.image, command, logger),
@@ -303,6 +368,14 @@ export default function createIntegration(args?: Options): AstroIntegration {
 						logger.error('Failed to write _redirects file');
 					}
 				}
+
+				// Get chunks from the bundle that are not needed on the server and delete them
+				// Those modules are build only for prerendering routes.
+				const chunksToDelete = chunkAnalyzer.getNonServerChunks();
+				for (const chunk of chunksToDelete) {
+					// Chunks are located on `./_worker.js` directory inside of the output directory
+					await unlink(new URL(`./_worker.js/${chunk}`, _config.outDir));
+				}
 			},
 		},
 	};
diff --git a/packages/integrations/cloudflare/src/utils/non-server-chunk-detector.ts b/packages/integrations/cloudflare/src/utils/non-server-chunk-detector.ts
new file mode 100644
index 000000000..c7e766802
--- /dev/null
+++ b/packages/integrations/cloudflare/src/utils/non-server-chunk-detector.ts
@@ -0,0 +1,83 @@
+import type { OutputBundle } from 'rollup';
+import type { Plugin } from 'vite';
+
+/**
+ * A Vite bundle analyzer that identifies chunks that are not used for server rendering.
+ *
+ * The chunks injected by Astro for prerendering are flagged as non-server chunks.
+ * Any chunks that is only used by a non-server chunk are also flagged as non-server chunks.
+ * This continues transitively until all non-server chunks are found.
+ */
+export class NonServerChunkDetector {
+	private nonServerChunks?: string[];
+
+	public getPlugin(): Plugin {
+		return {
+			name: 'non-server-chunk-detector',
+			generateBundle: (_, bundle) => {
+				this.processBundle(bundle);
+			},
+		};
+	}
+
+	private processBundle(bundle: OutputBundle) {
+		const chunkNamesToFiles = new Map<string, string>();
+
+		const entryChunks: string[] = [];
+		const chunkToDependencies = new Map<string, string[]>();
+
+		for (const chunk of Object.values(bundle)) {
+			if (chunk.type !== 'chunk') continue;
+
+			// Construct a mapping from a chunk name to its file name
+			chunkNamesToFiles.set(chunk.name, chunk.fileName);
+			// Construct a mapping from a chunk file to all the modules it imports
+			chunkToDependencies.set(chunk.fileName, [...chunk.imports, ...chunk.dynamicImports]);
+
+			if (chunk.isEntry) {
+				// Entry chunks should always be kept around since they are to be imported by the runtime
+				entryChunks.push(chunk.fileName);
+			}
+		}
+
+		const chunkDecisions = new Map<string, boolean>();
+
+		for (const entry of entryChunks) {
+			// Entry chunks are used on the server
+			chunkDecisions.set(entry, true);
+		}
+
+		for (const chunk of ['prerender', 'prerender@_@astro']) {
+			// Prerender chunks are not used on the server
+			const fileName = chunkNamesToFiles.get(chunk);
+			if (fileName) {
+				chunkDecisions.set(fileName, false);
+			}
+		}
+
+		// Start a stack of chunks that are used on the server
+		const chunksToWalk = [...entryChunks];
+
+		// Iterate over the chunks, traversing the transitive dependencies of the chunks used on the server
+		for (let chunk = chunksToWalk.pop(); chunk; chunk = chunksToWalk.pop()) {
+			for (const dep of chunkToDependencies.get(chunk) ?? []) {
+				// Skip dependencies already flagged, dependencies may be repeated and/or circular
+				if (chunkDecisions.has(dep)) continue;
+
+				// A dependency of a module used on the server is also used on the server
+				chunkDecisions.set(dep, true);
+				// Add the dependency to the stack so its own dependencies are also flagged
+				chunksToWalk.push(dep);
+			}
+		}
+
+		// Any chunk not flagged as used on the server is a non-server chunk
+		this.nonServerChunks = Array.from(chunkToDependencies.keys()).filter(
+			(chunk) => !chunkDecisions.get(chunk)
+		);
+	}
+
+	public getNonServerChunks(): string[] {
+		return this.nonServerChunks ?? [];
+	}
+}