Support user CPU time limit (#4204)

* Support user CPU time limit User limits provided via script metadata on upload Example configuration: ``` [limits] cpu_ms = 20000 ``` * Warn when attempting to set limits and standard is not enabled
cloudflare · Oct 19, 2023 · 38fdbe9 · 38fdbe9
1 parent ff3c3dc
commit 38fdbe9
Show file tree

Hide file tree

Showing 12 changed files with 127 additions and 0 deletions.
diff --git a/.changeset/pink-jokes-double.md b/.changeset/pink-jokes-double.md
@@ -0,0 +1,14 @@
+---
+"wrangler": minor
+---
+
+Support user limits for CPU time
+
+User limits provided via script metadata on upload
+
+Example configuration:
+
+```
+[limits]
+cpu_ms = 20000
+```
diff --git a/packages/wrangler/src/__tests__/deploy.test.ts b/packages/wrangler/src/__tests__/deploy.test.ts
@@ -5072,6 +5072,29 @@ addEventListener('fetch', event => {});`
 		});
 	});
 
+	describe("user limits", () => {
+		it("should allow specifying a cpu millisecond limit", async () => {
+			writeWranglerToml({
+				limits: { cpu_ms: 15_000 },
+			});
+
+			await fs.promises.writeFile("index.js", `export default {};`);
+			mockSubDomainRequest();
+			mockUploadWorkerRequest({
+				expectedLimits: { cpu_ms: 15_000 },
+			});
+
+			await runWrangler("publish index.js");
+			expect(std.out).toMatchInlineSnapshot(`
+			"Total Upload: xx KiB / gzip: xx KiB
+			Uploaded test-name (TIMINGS)
+			Published test-name (TIMINGS)
+			  https://test-name.test-sub-domain.workers.dev
+			Current Deployment ID: Galaxy-Class"
+		`);
+		});
+	});
+
 	describe("bindings", () => {
 		it("should allow bindings with different names", async () => {
 			writeWranglerToml({
@@ -8637,6 +8660,7 @@ export function mockUploadWorkerRequest(
 		expectedTailConsumers?: CfWorkerInit["tail_consumers"];
 		expectedUnsafeMetaData?: Record<string, string>;
 		expectedCapnpSchema?: string;
+		expectedLimits?: CfWorkerInit["limits"];
 		env?: string;
 		legacyEnv?: boolean;
 		keepVars?: boolean;
@@ -8658,6 +8682,7 @@ export function mockUploadWorkerRequest(
 		expectedTailConsumers,
 		expectedUnsafeMetaData,
 		expectedCapnpSchema,
+		expectedLimits,
 		keepVars,
 	} = options;
 	if (env && !legacyEnv) {
@@ -8734,6 +8759,9 @@ export function mockUploadWorkerRequest(
 				expectedCapnpSchema
 			);
 		}
+		if ("expectedLimits" in options) {
+			expect(metadata.limits).toEqual(expectedLimits);
+		}
 		if (expectedUnsafeMetaData !== undefined) {
 			Object.keys(expectedUnsafeMetaData).forEach((key) => {
 				expect(metadata[key]).toEqual(expectedUnsafeMetaData[key]);

diff --git a/packages/wrangler/src/__tests__/standard-pricing.test.ts b/packages/wrangler/src/__tests__/standard-pricing.test.ts
@@ -85,6 +85,34 @@ describe("standard-pricing", () => {
 		}
 	`);
 	});
+	it("should warn user about limits set if not enabled", async () => {
+		msw.use(...mswSuccessDeploymentScriptMetadata);
+		writeWranglerToml({ limits: { cpu_ms: 20_000 } });
+		writeWorkerSource();
+		mockSubDomainRequest();
+		mockUploadWorkerRequest();
+
+		mockStandardEnabled(false);
+
+		await runWrangler("deploy ./index");
+
+		expect(std).toMatchInlineSnapshot(`
+		Object {
+		  "debug": "",
+		  "err": "",
+		  "info": "",
+		  "out": "🚧 New Workers Standard pricing is now available. Please visit the dashboard to view details and opt-in to new pricing: https://dash.cloudflare.com/some-account-id/workers/standard/opt-in.
+		Total Upload: xx KiB / gzip: xx KiB
+		Uploaded test-name (TIMINGS)
+		Published test-name (TIMINGS)
+		  https://test-name.test-sub-domain.workers.dev
+		Current Deployment ID: Galaxy-Class",
+		  "warn": "[33m▲ [43;33m[[43;30mWARNING[43;33m][0m [1mThe \`limits\` defined in wrangler.toml can only be applied to scripts opted into Workers Standard pricing. Agree to the new pricing details to set limits for your script.[0m
+
+		",
+		}
+	`);
+	});
 	it("should not notify user about new pricing if enterprise", async () => {
 		msw.use(...mswSuccessDeploymentScriptMetadata);
 		writeWranglerToml();

diff --git a/packages/wrangler/src/api/pages/create-worker-bundle-contents.ts b/packages/wrangler/src/api/pages/create-worker-bundle-contents.ts
@@ -78,6 +78,7 @@ function createWorkerBundleFormData(workerBundle: BundleResult): FormData {
 		logpush: undefined,
 		placement: undefined,
 		tail_consumers: undefined,
+		limits: undefined,
 	};
 
 	return createWorkerUploadForm(worker);

diff --git a/packages/wrangler/src/config/environment.ts b/packages/wrangler/src/config/environment.ts
@@ -170,6 +170,14 @@ interface EnvironmentInheritable {
 	 */
 	usage_model: "bundled" | "unbound" | undefined;
 
+	/**
+	 * Specify limits for runtime behavior.
+	 * Only supported for the "standard" Usage Model
+	 *
+	 * @inheritable
+	 */
+	limits: UserLimits | undefined;
+
 	/**
 	 * An ordered list of rules that define which modules to import,
 	 * and what type to import them as. You will need to specify rules
@@ -736,3 +744,8 @@ export interface DispatchNamespaceOutbound {
 	/** (Optional) List of parameter names, for sending context from your dispatch worker to the outbound handler */
 	parameters?: string[];
 }
+
+export interface UserLimits {
+	/** Maximum allowed CPU time for a worker's invocation in milliseconds */
+	cpu_ms: number;
+}
diff --git a/packages/wrangler/src/config/validation.ts b/packages/wrangler/src/config/validation.ts
@@ -1155,6 +1155,7 @@ function normalizeAndValidateEnvironment(
 			isOneOf("bundled", "unbound"),
 			undefined
 		),
+		limits: normalizeAndValidateLimits(diagnostics, topLevelEnv, rawEnv),
 		placement: normalizeAndValidatePlacement(diagnostics, topLevelEnv, rawEnv),
 		build,
 		workers_dev,
@@ -2801,3 +2802,28 @@ const validateConsumer: ValidatorFn = (diagnostics, field, value, _config) => {
 
 	return isValid;
 };
+
+function normalizeAndValidateLimits(
+	diagnostics: Diagnostics,
+	topLevelEnv: Environment | undefined,
+	rawEnv: RawEnvironment
+): Config["limits"] {
+	if (rawEnv.limits) {
+		validateRequiredProperty(
+			diagnostics,
+			"limits",
+			"cpu_ms",
+			rawEnv.limits.cpu_ms,
+			"number"
+		);
+	}
+
+	return inheritable(
+		diagnostics,
+		topLevelEnv,
+		rawEnv,
+		"limits",
+		() => true,
+		undefined
+	);
+}
diff --git a/packages/wrangler/src/deploy/deploy.ts b/packages/wrangler/src/deploy/deploy.ts
@@ -608,6 +608,7 @@ See https://developers.cloudflare.com/workers/platform/compatibility-dates for m
 			logpush: props.logpush !== undefined ? props.logpush : config.logpush,
 			placement,
 			tail_consumers: config.tail_consumers,
+			limits: config.limits,
 		};
 
 		// As this is not deterministic for testing, we detect if in a jest environment and run asynchronously

diff --git a/packages/wrangler/src/deploy/index.ts b/packages/wrangler/src/deploy/index.ts
@@ -37,6 +37,11 @@ async function standardPricingWarning(
 					`🚧 New Workers Standard pricing is now available. Please visit the dashboard to view details and opt-in to new pricing: https://dash.cloudflare.com/${accountId}/workers/standard/opt-in.`
 				)
 			);
+			if (config.limits?.cpu_ms !== undefined) {
+				logger.warn(
+					"The `limits` defined in wrangler.toml can only be applied to scripts opted into Workers Standard pricing. Agree to the new pricing details to set limits for your script."
+				);
+			}
 			return;
 		}
 		if (standard && config.usage_model !== undefined) {

diff --git a/packages/wrangler/src/deployment-bundle/create-worker-upload-form.ts b/packages/wrangler/src/deployment-bundle/create-worker-upload-form.ts
@@ -9,6 +9,7 @@ import type {
 	CfDurableObjectMigrations,
 	CfPlacement,
 	CfTailConsumer,
+	CfUserLimits,
 } from "./worker.js";
 import type { Json } from "miniflare";
 
@@ -105,6 +106,7 @@ export interface WorkerMetadata {
 	logpush?: boolean;
 	placement?: CfPlacement;
 	tail_consumers?: CfTailConsumer[];
+	limits?: CfUserLimits;
 	// Allow unsafe.metadata to add arbitrary properties at runtime
 	[key: string]: unknown;
 }
@@ -125,6 +127,7 @@ export function createWorkerUploadForm(worker: CfWorkerInit): FormData {
 		logpush,
 		placement,
 		tail_consumers,
+		limits,
 	} = worker;
 
 	let { modules } = worker;
@@ -462,6 +465,7 @@ export function createWorkerUploadForm(worker: CfWorkerInit): FormData {
 		...(logpush !== undefined && { logpush }),
 		...(placement && { placement }),
 		...(tail_consumers && { tail_consumers }),
+		...(limits && { limits }),
 	};
 
 	if (bindings.unsafe?.metadata !== undefined) {

diff --git a/packages/wrangler/src/deployment-bundle/worker.ts b/packages/wrangler/src/deployment-bundle/worker.ts
@@ -248,6 +248,10 @@ export interface CfTailConsumer {
 	environment?: string;
 }
 
+export interface CfUserLimits {
+	cpu_ms?: number;
+}
+
 /**
  * Options for creating a `CfWorker`.
  */
@@ -298,6 +302,7 @@ export interface CfWorkerInit {
 	logpush: boolean | undefined;
 	placement: CfPlacement | undefined;
 	tail_consumers: CfTailConsumer[] | undefined;
+	limits: CfUserLimits | undefined;
 }
 
 export interface CfWorkerContext {

diff --git a/packages/wrangler/src/dev/remote.tsx b/packages/wrangler/src/dev/remote.tsx
@@ -585,6 +585,7 @@ async function createRemoteWorkerInit(props: {
 		logpush: false,
 		placement: undefined, // no placement in dev
 		tail_consumers: undefined, // no tail consumers in dev - TODO revisit?
+		limits: undefined, // no limits in preview - not supported yet but can be added
 	};
 
 	return init;

diff --git a/packages/wrangler/src/secret/index.ts b/packages/wrangler/src/secret/index.ts
@@ -90,6 +90,7 @@ async function createDraftWorker({
 				logpush: false,
 				placement: undefined,
 				tail_consumers: undefined,
+				limits: undefined,
 			}),
 		}
 	);