Add output for tail being in sampling mode (#3146)

cloudflare · May 8, 2023 · 5b234cf · 5b234cf
1 parent 133c042
commit 5b234cf
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 7 deletions.
diff --git a/.changeset/witty-pumpkins-swim.md b/.changeset/witty-pumpkins-swim.md
@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+Added output for tail being in "sampling mode"
diff --git a/packages/wrangler/src/__tests__/pages-deployment-tail.test.ts b/packages/wrangler/src/__tests__/pages-deployment-tail.test.ts
@@ -13,6 +13,7 @@ import type {
 	ScheduledEvent,
 	AlarmEvent,
 	EmailEvent,
+	TailInfo,
 } from "../tail/createTail";
 import type { RequestInit } from "undici";
 import type WebSocket from "ws";
@@ -655,6 +656,7 @@ function isRequest(
 		| RequestEvent
 		| AlarmEvent
 		| EmailEvent
+		| TailInfo
 		| undefined
 		| null
 ): event is RequestEvent {

diff --git a/packages/wrangler/src/__tests__/tail.test.ts b/packages/wrangler/src/__tests__/tail.test.ts
@@ -14,6 +14,7 @@ import type {
 	ScheduledEvent,
 	AlarmEvent,
 	EmailEvent,
+	TailInfo,
 } from "../tail/createTail";
 import type { RequestInit } from "undici";
 import type WebSocket from "ws";
@@ -56,10 +57,10 @@ describe("tail", () => {
 			await runWrangler("tail durable-object--websocket--response");
 			expect(std.out).toMatchInlineSnapshot(`""`);
 			expect(std.warn).toMatchInlineSnapshot(`
-"[33m▲ [43;33m[[43;30mWARNING[43;33m][0m [1mBeginning log collection requires restarting the Durable Objects associated with durable-object--websocket--response. Any WebSocket connections or other non-persisted state will be lost as part of this restart.[0m
+			"[33m▲ [43;33m[[43;30mWARNING[43;33m][0m [1mBeginning log collection requires restarting the Durable Objects associated with durable-object--websocket--response. Any WebSocket connections or other non-persisted state will be lost as part of this restart.[0m
 
-"
-`);
+			"
+		`);
 			expect(std.err).toMatchInlineSnapshot(`""`);
 		});
 		it("creates and then delete tails", async () => {
@@ -503,10 +504,31 @@ describe("tail", () => {
 					)
 					.replace(mockTailExpiration.toISOString(), "[mock expiration date]")
 			).toMatchInlineSnapshot(`
-			        "Successfully created tail, expires at [mock expiration date]
-			        Connected to test-worker, waiting for logs...
-			        Email from:${mockEmailEventFrom} to:${mockEmailEventTo} size:${mockEmailEventSize} @ [mock event timestamp] - Ok"
-		      `);
+			"Successfully created tail, expires at [mock expiration date]
+			Connected to test-worker, waiting for logs...
+			Email from:[email protected] to:[email protected] size:45416 @ [mock event timestamp] - Ok"
+		`);
+		});
+
+		it("logs tail overload message", async () => {
+			const api = mockWebsocketAPIs();
+			await runWrangler("tail test-worker --format pretty");
+
+			const event = generateTailInfo();
+			const message = generateMockEventMessage({ event });
+			const serializedMessage = serialize(message);
+
+			api.ws.send(serializedMessage);
+			expect(
+				std.out.replace(
+					mockTailExpiration.toISOString(),
+					"[mock expiration date]"
+				)
+			).toMatchInlineSnapshot(`
+			"Successfully created tail, expires at [mock expiration date]
+			Connected to test-worker, waiting for logs...
+			Tail is currently in sampling mode due to the high volume of messages. To prevent messages from being dropped consider adding filters."
+		`);
 		});
 
 		it("should not crash when the tail message has a void event", async () => {
@@ -675,6 +697,7 @@ function isRequest(
 		| RequestEvent
 		| AlarmEvent
 		| EmailEvent
+		| TailInfo
 		| undefined
 		| null
 ): event is RequestEvent {
@@ -956,3 +979,11 @@ function generateMockEmailEvent(opts?: Partial<EmailEvent>): EmailEvent {
 		rawSize: opts?.rawSize || mockEmailEventSize,
 	};
 }
+
+function generateTailInfo(): TailInfo {
+	return {
+		message:
+			"Tail is currently in sampling mode due to the high volume of messages. To prevent messages from being dropped consider adding filters.",
+		type: "overload",
+	};
+}
diff --git a/packages/wrangler/src/tail/createTail.ts b/packages/wrangler/src/tail/createTail.ts
@@ -252,6 +252,7 @@ export type TailEventMessage = {
 		| ScheduledEvent
 		| AlarmEvent
 		| EmailEvent
+		| TailInfo
 		| undefined
 		| null;
 };
@@ -404,3 +405,11 @@ export type EmailEvent = {
 	 */
 	rawSize: number;
 };
+
+/**
+ * Message from tail with information about the tail itself
+ */
+export type TailInfo = {
+	message: string;
+	type: string;
+};
diff --git a/packages/wrangler/src/tail/printing.ts b/packages/wrangler/src/tail/printing.ts
@@ -1,9 +1,11 @@
+import chalk from "chalk";
 import { logger } from "../logger";
 import type {
 	AlarmEvent,
 	EmailEvent,
 	RequestEvent,
 	ScheduledEvent,
+	TailInfo,
 	TailEventMessage,
 } from "./createTail";
 import type { Outcome } from "./filters";
@@ -48,6 +50,10 @@ export function prettyPrintLogs(data: WebSocket.RawData): void {
 		).toLocaleString();
 
 		logger.log(`Alarm @ ${datetime} - ${outcome}`);
+	} else if (isTailInfo(eventMessage.event)) {
+		if (eventMessage.event.type === "overload") {
+			logger.log(`${chalk.red.bold(eventMessage.event.message)}`);
+		}
 	} else {
 		// Unknown event type
 		const outcome = prettifyOutcome(eventMessage.outcome);
@@ -103,6 +109,10 @@ function isAlarmEvent(event: TailEventMessage["event"]): event is AlarmEvent {
 	return Boolean(event && "scheduledTime" in event && !("cron" in event));
 }
 
+function isTailInfo(event: TailEventMessage["event"]): event is TailInfo {
+	return Boolean(event && "message" in event && "type" in event);
+}
+
 function prettifyOutcome(outcome: Outcome): string {
 	switch (outcome) {
 		case "ok":