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

12 files changed, 313 insertions, 7 deletions

diff --git a/packages/integrations/cloudflare/package.json b/packages/integrations/cloudflare/package.json
index 44b07ac84..433d3c4b3 100644
--- a/packages/integrations/cloudflare/package.json
+++ b/packages/integrations/cloudflare/package.json
@@ -27,11 +27,11 @@
     "test": "astro-scripts test \"test/**/*.test.js\""
   },
   "dependencies": {
-    "@astrojs/underscore-redirects": "^0.3.3",
     "@astrojs/internal-helpers": "0.3.0",
+    "@astrojs/underscore-redirects": "^0.3.3",
     "@cloudflare/workers-types": "^4.20240320.1",
-    "miniflare": "^3.20240320.0",
     "esbuild": "^0.19.5",
+    "miniflare": "^3.20240320.0",
     "tiny-glob": "^0.2.9",
     "wrangler": "^3.39.0"
   },
@@ -39,13 +39,15 @@
     "astro": "^4.2.0"
   },
   "devDependencies": {
+    "@astrojs/test-utils": "workspace:*",
+    "astro": "^4.5.8",
+    "astro-scripts": "workspace:*",
+    "cheerio": "1.0.0-rc.12",
     "execa": "^8.0.1",
     "fast-glob": "^3.3.2",
+    "rollup": "^4.14.0",
     "strip-ansi": "^7.1.0",
-    "astro": "^4.5.8",
-    "cheerio": "1.0.0-rc.12",
-    "@astrojs/test-utils": "workspace:*",
-    "astro-scripts": "workspace:*"
+    "vite": "^5.2.6"
   },
   "publishConfig": {
     "provenance": true
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 ?? [];
+	}
+}
diff --git a/packages/integrations/cloudflare/test/exclude-prerender-only-chunks.test.js b/packages/integrations/cloudflare/test/exclude-prerender-only-chunks.test.js
new file mode 100644
index 000000000..5a142f429
--- /dev/null
+++ b/packages/integrations/cloudflare/test/exclude-prerender-only-chunks.test.js
@@ -0,0 +1,86 @@
+import * as assert from 'node:assert/strict';
+import { readFile, readdir } from 'node:fs/promises';
+import { join, relative } from 'node:path';
+import { before, describe, it } from 'node:test';
+import { fileURLToPath } from 'node:url';
+import { astroCli } from './_test-utils.js';
+
+const root = new URL('./fixtures/prerender-optimizations/', import.meta.url);
+
+async function lookForCodeInServerBundle(code) {
+	const serverBundleRoot = fileURLToPath(new URL('./dist/_worker.js/', root));
+
+	const entries = await readdir(serverBundleRoot, {
+		withFileTypes: true,
+		recursive: true,
+	}).catch((err) => {
+		console.log('Failed to read server bundle directory:', err);
+
+		throw err;
+	});
+
+	for (const entry of entries) {
+		if (!entry.isFile()) continue;
+
+		const filePath = join(entry.path, entry.name);
+		const fileContent = await readFile(filePath, 'utf-8').catch((err) => {
+			console.log(`Failed to read file ${filePath}:`, err);
+
+			throw err;
+		});
+
+		if (fileContent.includes(code)) {
+			return relative(serverBundleRoot, filePath);
+		}
+	}
+
+	return null;
+}
+
+describe('worker.js cleanup after pre-rendering', () => {
+	before(async () => {
+		const res = await astroCli(fileURLToPath(root), 'build');
+	});
+
+	it('should not include code from pre-rendered pages in the server bundle', async () => {
+		assert.equal(
+			await lookForCodeInServerBundle('frontmatter of prerendered page'),
+			null,
+			'Code from pre-rendered pages should not be included in the server bundle.'
+		);
+
+		assert.equal(
+			await lookForCodeInServerBundle('Body of Prerendered Page'),
+			null,
+			'Code from pre-rendered pages should not be included in the server bundle.'
+		);
+	});
+
+	it('should not include markdown content used only in pre-rendered pages in the server bundle', async () => {
+		assert.equal(
+			await lookForCodeInServerBundle('Sample Post Title'),
+			null,
+			'Markdown frontmatter used only on pre-rendered pages should not be included in the server bundle.'
+		);
+
+		assert.equal(
+			await lookForCodeInServerBundle('Sample Post Content'),
+			null,
+			'Markdown content used only on pre-rendered pages should not be included in the server bundle.'
+		);
+	});
+
+	it('should include code for on-demand pages in the server bundle', async () => {
+		assert.notEqual(
+			await lookForCodeInServerBundle('frontmatter of SSR page'),
+			null,
+			'Code from pre-rendered pages should not be included in the server bundle.'
+		);
+
+		assert.notEqual(
+			await lookForCodeInServerBundle('Body of SSR Page'),
+			null,
+			'Code from pre-rendered pages should not be included in the server bundle.'
+		);
+	});
+});
diff --git a/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/.gitignore b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/.gitignore
new file mode 100644
index 000000000..97e663b17
--- /dev/null
+++ b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/.gitignore
@@ -0,0 +1 @@
+.astro/
diff --git a/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/astro.config.ts b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/astro.config.ts
new file mode 100644
index 000000000..339f0e2a4
--- /dev/null
+++ b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/astro.config.ts
@@ -0,0 +1,7 @@
+import cloudflare from '@astrojs/cloudflare';
+import { defineConfig } from 'astro/config';
+
+export default defineConfig({
+	adapter: cloudflare(),
+	output: 'server',
+});
diff --git a/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/package.json b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/package.json
new file mode 100644
index 000000000..ad583d9d5
--- /dev/null
+++ b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/package.json
@@ -0,0 +1,9 @@
+{
+  "name": "@test/astro-cloudflare-prerender-optimizations",
+  "version": "0.0.0",
+  "private": true,
+  "dependencies": {
+    "@astrojs/cloudflare": "workspace:*",
+    "astro": "^4.3.5"
+  }
+}
diff --git a/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/content/config.ts b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/content/config.ts
new file mode 100644
index 000000000..83ef12586
--- /dev/null
+++ b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/content/config.ts
@@ -0,0 +1,8 @@
+import { defineCollection } from 'astro:content';
+import { z } from 'astro/zod';
+
+export const collections = {
+  posts: defineCollection({
+    schema: z.any(),
+  }),
+}
diff --git a/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/content/posts/sample.md b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/content/posts/sample.md
new file mode 100644
index 000000000..060253707
--- /dev/null
+++ b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/content/posts/sample.md
@@ -0,0 +1,5 @@
+---
+title: Sample Post Title
+---
+
+Sample Post Content
diff --git a/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/env.d.ts b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/env.d.ts
new file mode 100644
index 000000000..c13bd73c7
--- /dev/null
+++ b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/env.d.ts
@@ -0,0 +1,2 @@
+/// <reference path="../.astro/types.d.ts" />
+/// <reference types="astro/client" />
\ No newline at end of file
diff --git a/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/pages/on-demand.astro b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/pages/on-demand.astro
new file mode 100644
index 000000000..490295b7b
--- /dev/null
+++ b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/pages/on-demand.astro
@@ -0,0 +1,12 @@
+---
+export const prerender = false;
+
+console.log('frontmatter of SSR page');
+---
+
+<html>
+	<head></head>
+	<body>
+		<h1>Body of SSR Page</h1>
+	</body>
+</html>
diff --git a/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/pages/prerender.astro b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/pages/prerender.astro
new file mode 100644
index 000000000..08d0318e6
--- /dev/null
+++ b/packages/integrations/cloudflare/test/fixtures/prerender-optimizations/src/pages/prerender.astro
@@ -0,0 +1,18 @@
+---
+import { getEntry } from 'astro:content';
+
+export const prerender = true;
+
+console.log('frontmatter of prerendered page');
+
+const samplePost = await getEntry('posts', 'sample');
+const { Content } = await samplePost.render();
+---
+
+<html>
+	<head></head>
+	<body>
+		<h1>Body of Prerendered Page</h1>
+		<Content />
+	</body>
+</html>