Add per-isolate SSE + OTel memory counters to cloud worker (#1255)

* Add per-isolate SSE + OTel memory counters to cloud worker * Address review: abort-free byte counts, SSE-only wrap guard, no drop/export double-count

Rhys Sullivan committed 4 months ago

Commit fd840191 parent 1c481823 files · +373 −14Browse files

Unified · Split

3 files

apps/cloud/src/mcp/agent-handler.ts +3 −1
apps/cloud/src/observability/memory-metrics.ts +329 −0
apps/cloud/src/observability/telemetry.ts +41 −13

apps/cloud/src/mcp/agent-handler.ts+3−1

 } from "@executor-js/cloudflare/mcp/do-headers";
 import type { McpSessionProps } from "@executor-js/cloudflare/mcp/agent-durable-object";
+import { wrapMcpSseResponse } from "../observability/memory-metrics";
 import { cloudMcpAuth } from "./auth-provider";
 import { McpSessionDOSqlite } from "./session-durable-object";
       },
       resource,
     );
-    return serve.fetch(forwarded, env, ctx);
+    const response = await serve.fetch(forwarded, env, ctx);
+    return wrapMcpSseResponse(request, env, response);
   };
 };

apps/cloud/src/observability/memory-metrics.ts+329−0

@@ -0,0 +1,329 @@
	1	import type { Context } from "@opentelemetry/api";
	2	import type {
	3	ReadableSpan,
	4	Span,
	5	SpanExporter,
	6	SpanProcessor,
	7	} from "@opentelemetry/sdk-trace-base";
	8
	9	const LOG_PREFIX = "[executor:mem-metrics]";
	10	const PERIODIC_EMIT_INTERVAL_MS = 60_000;
	11	const STALLED_WRITE_MS = 30_000;
	12	const EXPORT_SUCCESS_CODE = 0;
	13	export const OTEL_MAX_SPAN_QUEUE_SIZE = 2_048;
	14
	15	type CloseReason = "normal" \| "cancel" \| "error";
	16
	17	type CfRequestMetadata = {
	18	readonly colo?: string;
	19	};
	20
	21	type VersionMetadataEnv = {
	22	readonly VERSION_METADATA?: WorkerVersionMetadata;
	23	readonly CF_VERSION_METADATA?: WorkerVersionMetadata;
	24	readonly SCRIPT_VERSION?: string;
	25	};
	26
	27	type SseConnection = {
	28	readonly id: number;
	29	readonly startedAt: number;
	30	lastWriteAt: number;
	31	chunks: number;
	32	bytes: number;
	33	closed: boolean;
	34	colo: string;
	35	scriptVersion: string;
	36	};
	37
	38	type AgeBuckets = {
	39	readonly lt1m: number;
	40	readonly m1to5: number;
	41	readonly m5to30: number;
	42	readonly m30to60: number;
	43	readonly gt60m: number;
	44	};
	45
	46	type OtelMetricsSnapshot = {
	47	readonly spansEnded: number;
	48	readonly spansExported: number;
	49	readonly spansDropped: number;
	50	readonly spanExportFailures: number;
	51	readonly forceFlushCalls: number;
	52	readonly forceFlushFailures: number;
	53	readonly forceFlushDurationMsTotal: number;
	54	readonly forceFlushDurationMsLast: number;
	55	readonly maxQueueSize: number;
	56	};
	57
	58	const activeSse = new Map<number, SseConnection>();
	59
	60	let nextConnectionId = 1;
	61	let totalBytesForwarded = 0;
	62	let lastSnapshotEmittedAt = 0;
	63
	64	const otelMetrics = {
	65	spansEnded: 0,
	66	spansExported: 0,
	67	spansDropped: 0,
	68	spanExportFailures: 0,
	69	forceFlushCalls: 0,
	70	forceFlushFailures: 0,
	71	forceFlushDurationMsTotal: 0,
	72	forceFlushDurationMsLast: 0,
	73	};
	74
	75	const requestWithCf = (request: Request): Request & { readonly cf?: CfRequestMetadata } =>
	76	request as Request & { readonly cf?: CfRequestMetadata };
	77
	78	const scriptVersionFromEnv = (env: Env): string => {
	79	const metadataEnv = env as Env & VersionMetadataEnv;
	80	return (
	81	metadataEnv.VERSION_METADATA?.id ??
	82	metadataEnv.VERSION_METADATA?.tag ??
	83	metadataEnv.CF_VERSION_METADATA?.id ??
	84	metadataEnv.CF_VERSION_METADATA?.tag ??
	85	metadataEnv.SCRIPT_VERSION ??
	86	""
	87	);
	88	};
	89
	90	const utf8ByteCounter = new TextEncoder();
	91
	92	const byteLengthOf = (chunk: unknown): number => {
	93	if (typeof chunk === "string") return utf8ByteCounter.encode(chunk).length;
	94	if (chunk instanceof ArrayBuffer) return chunk.byteLength;
	95	if (ArrayBuffer.isView(chunk)) return chunk.byteLength;
	96	return 0;
	97	};
	98
	99	const emptyAgeBuckets = (): AgeBuckets => ({
	100	lt1m: 0,
	101	m1to5: 0,
	102	m5to30: 0,
	103	m30to60: 0,
	104	gt60m: 0,
	105	});
	106
	107	const bucketAge = (buckets: AgeBuckets, ageMs: number): AgeBuckets => {
	108	const next = { ...buckets };
	109	if (ageMs < 60_000) next.lt1m += 1;
	110	else if (ageMs < 5 * 60_000) next.m1to5 += 1;
	111	else if (ageMs < 30 * 60_000) next.m5to30 += 1;
	112	else if (ageMs < 60 * 60_000) next.m30to60 += 1;
	113	else next.gt60m += 1;
	114	return next;
	115	};
	116
	117	const otelSnapshot = (maxQueueSize: number): OtelMetricsSnapshot => ({
	118	...otelMetrics,
	119	maxQueueSize,
	120	});
	121
	122	const snapshot = (now: number, maxQueueSize: number) => {
	123	let oldestConnectionAgeMs = 0;
	124	let stalledConnections = 0;
	125	let ageBuckets = emptyAgeBuckets();
	126	const colos = new Set<string>();
	127	const scriptVersions = new Set<string>();
	128
	129	for (const connection of activeSse.values()) {
	130	const ageMs = now - connection.startedAt;
	131	oldestConnectionAgeMs = Math.max(oldestConnectionAgeMs, ageMs);
	132	if (now - connection.lastWriteAt > STALLED_WRITE_MS) stalledConnections += 1;
	133	ageBuckets = bucketAge(ageBuckets, ageMs);
	134	if (connection.colo) colos.add(connection.colo);
	135	if (connection.scriptVersion) scriptVersions.add(connection.scriptVersion);
	136	}
	137	const sortedColos = Array.from(colos).sort();
	138	const sortedScriptVersions = Array.from(scriptVersions).sort();
	139
	140	return {
	141	activeSseConnections: activeSse.size,
	142	ageBuckets,
	143	oldestConnectionAgeMs,
	144	totalBytesForwarded,
	145	stalledConnections,
	146	stalledWriteThresholdMs: STALLED_WRITE_MS,
	147	colo: sortedColos[0] ?? "",
	148	colos: sortedColos,
	149	scriptVersion: sortedScriptVersions[0] ?? "",
	150	scriptVersions: sortedScriptVersions,
	151	otel: otelSnapshot(maxQueueSize),
	152	};
	153	};
	154
	155	const emitSnapshot = (event: string, now: number, maxQueueSize: number): void => {
	156	lastSnapshotEmittedAt = now;
	157	console.log(
	158	`${LOG_PREFIX} ${JSON.stringify({
	159	event,
	160	...snapshot(now, maxQueueSize),
	161	})}`,
	162	);
	163	};
	164
	165	const maybeEmitPeriodicSnapshot = (now: number, maxQueueSize: number): void => {
	166	if (now - lastSnapshotEmittedAt < PERIODIC_EMIT_INTERVAL_MS) return;
	167	emitSnapshot("snapshot", now, maxQueueSize);
	168	};
	169
	170	const closeConnection = (
	171	connection: SseConnection,
	172	reason: CloseReason,
	173	maxQueueSize: number,
	174	): void => {
	175	if (connection.closed) return;
	176	connection.closed = true;
	177	activeSse.delete(connection.id);
	178	const now = Date.now();
	179	console.log(
	180	`${LOG_PREFIX} ${JSON.stringify({
	181	event: "sse_close",
	182	connectionId: connection.id,
	183	reason,
	184	ageMs: now - connection.startedAt,
	185	bytesForwarded: connection.bytes,
	186	chunksForwarded: connection.chunks,
	187	lastWriteAgeMs: now - connection.lastWriteAt,
	188	colo: connection.colo,
	189	scriptVersion: connection.scriptVersion,
	190	})}`,
	191	);
	192	maybeEmitPeriodicSnapshot(now, maxQueueSize);
	193	};
	194
	195	const sseHeaders = (headers: Headers): Headers => {
	196	const next = new Headers(headers);
	197	next.delete("content-length");
	198	return next;
	199	};
	200
	201	export const wrapMcpSseResponse = (request: Request, env: Env, response: Response): Response => {
	202	if (
	203	request.method !== "GET" \|\|
	204	response.body === null \|\|
	205	!(response.headers.get("content-type") ?? "").includes("text/event-stream")
	206	)
	207	return response;
	208
	209	const now = Date.now();
	210	const connection: SseConnection = {
	211	id: nextConnectionId++,
	212	startedAt: now,
	213	lastWriteAt: now,
	214	chunks: 0,
	215	bytes: 0,
	216	closed: false,
	217	colo: requestWithCf(request).cf?.colo ?? "",
	218	scriptVersion: scriptVersionFromEnv(env),
	219	};
	220	activeSse.set(connection.id, connection);
	221	emitSnapshot("sse_open", now, OTEL_MAX_SPAN_QUEUE_SIZE);
	222
	223	const counting = new TransformStream<unknown, unknown>({
	224	transform(chunk, controller) {
	225	const written = byteLengthOf(chunk);
	226	const writeAt = Date.now();
	227	connection.chunks += 1;
	228	connection.bytes += written;
	229	connection.lastWriteAt = writeAt;
	230	totalBytesForwarded += written;
	231	maybeEmitPeriodicSnapshot(writeAt, OTEL_MAX_SPAN_QUEUE_SIZE);
	232	controller.enqueue(chunk);
	233	},
	234	flush() {
	235	closeConnection(connection, "normal", OTEL_MAX_SPAN_QUEUE_SIZE);
	236	},
	237	cancel(reason) {
	238	closeConnection(
	239	connection,
	240	reason === undefined ? "cancel" : "error",
	241	OTEL_MAX_SPAN_QUEUE_SIZE,
	242	);
	243	},
	244	});
	245
	246	const body = response.body.pipeThrough(counting);
	247	return new Response(body, {
	248	status: response.status,
	249	statusText: response.statusText,
	250	headers: sseHeaders(response.headers),
	251	});
	252	};
	253
	254	export const recordForceFlush = (durationMs: number, failed: boolean): void => {
	255	otelMetrics.forceFlushCalls += 1;
	256	otelMetrics.forceFlushDurationMsLast = durationMs;
	257	otelMetrics.forceFlushDurationMsTotal += durationMs;
	258	if (failed) otelMetrics.forceFlushFailures += 1;
	259	};
	260
	261	export class CountingSpanExporter implements SpanExporter {
	262	constructor(
	263	private readonly inner: SpanExporter,
	264	private readonly onExportAttempt: (spans: number) => void,
	265	) {}
	266
	267	export(spans: ReadableSpan[], resultCallback: Parameters<SpanExporter["export"]>[1]): void {
	268	this.onExportAttempt(spans.length);
	269	this.inner.export(spans, (result) => {
	270	if (result.code === EXPORT_SUCCESS_CODE) {
	271	otelMetrics.spansExported += spans.length;
	272	} else {
	273	otelMetrics.spansDropped += spans.length;
	274	otelMetrics.spanExportFailures += 1;
	275	}
	276	resultCallback(result);
	277	});
	278	}
	279
	280	shutdown(): Promise<void> {
	281	return this.inner.shutdown();
	282	}
	283
	284	forceFlush(): Promise<void> {
	285	return this.inner.forceFlush?.() ?? Promise.resolve();
	286	}
	287	}
	288
	289	export class CountingSpanProcessor implements SpanProcessor {
	290	private estimatedQueuedSpans = 0;
	291
	292	constructor(
	293	private readonly inner: SpanProcessor,
	294	private readonly maxQueueSize: number,
	295	) {}
	296
	297	forceFlush(): Promise<void> {
	298	return this.inner.forceFlush();
	299	}
	300
	301	onStart(span: Span, parentContext: Context): void {
	302	this.inner.onStart(span, parentContext);
	303	}
	304
	305	onEnding(span: Span): void {
	306	this.inner.onEnding?.(span);
	307	}
	308
	309	onEnd(span: ReadableSpan): void {
	310	otelMetrics.spansEnded += 1;
	311	if (this.estimatedQueuedSpans >= this.maxQueueSize) {
	312	// At the bound the inner processor would drop this span anyway; drop it
	313	// here instead of forwarding, so a dropped span is never also counted as
	314	// exported (which would double-count and inflate the drop signal).
	315	otelMetrics.spansDropped += 1;
	316	return;
	317	}
	318	this.estimatedQueuedSpans += 1;
	319	this.inner.onEnd(span);
	320	}
	321
	322	shutdown(): Promise<void> {
	323	return this.inner.shutdown();
	324	}
	325
	326	recordExportAttempt(spans: number): void {
	327	this.estimatedQueuedSpans = Math.max(0, this.estimatedQueuedSpans - spans);
	328	}
	329	}

apps/cloud/src/observability/telemetry.ts+41−13

 import { env } from "cloudflare:workers";
 import { Effect, Layer } from "effect";
+import {
+  CountingSpanExporter,
+  CountingSpanProcessor,
+  OTEL_MAX_SPAN_QUEUE_SIZE,
+  recordForceFlush,
+} from "./memory-metrics";
 const SERVICE_NAME = "executor-cloud";
 const SERVICE_VERSION = "1.0.0";
       [ATTR_SERVICE_NAME]: SERVICE_NAME,
       [ATTR_SERVICE_VERSION]: SERVICE_VERSION,
     }),
-    spanProcessors: [
+    spanProcessors: (() => {
-      // Batch, not Simple: SimpleSpanProcessor issues one synchronous fetch
+      let countingProcessor: CountingSpanProcessor;
-      // per span END — a busy MCP request emits dozens of spans, and the
+      const exporter = new CountingSpanExporter(
-      // serialized export fetches add seconds of latency inside the request
-      // (enough to flip pause/resume timing in e2e). Batch buffers and
-      // ships on a timer; `ctx.waitUntil(flushTracerProvider())` at the end
-      // of each request still drains the buffer before the isolate exits.
-      new BatchSpanProcessor(
         new OTLPTraceExporter({
           url: env.AXIOM_TRACES_URL ?? "https://api.axiom.co/v1/traces",
           headers: {
             "X-Axiom-Dataset": env.AXIOM_DATASET ?? "executor-cloud",
           },
         }),
-        { scheduledDelayMillis: 1_000, maxExportBatchSize: 512 },
+        (spans) => countingProcessor.recordExportAttempt(spans),
-      ),
+      );
-    ],
+      // Batch, not Simple: SimpleSpanProcessor issues one synchronous fetch
+      // per span END — a busy MCP request emits dozens of spans, and the
+      // serialized export fetches add seconds of latency inside the request
+      // (enough to flip pause/resume timing in e2e). Batch buffers and
+      // ships on a timer; `ctx.waitUntil(flushTracerProvider())` at the end
+      // of each request still drains the buffer before the isolate exits.
+      countingProcessor = new CountingSpanProcessor(
+        new BatchSpanProcessor(exporter, {
+          scheduledDelayMillis: 1_000,
+          maxExportBatchSize: 512,
+          maxQueueSize: OTEL_MAX_SPAN_QUEUE_SIZE,
+        }),
+        OTEL_MAX_SPAN_QUEUE_SIZE,
+      );
+      return [countingProcessor];
+    })(),
   });
   // Skip `provider.register()` — its StackContextManager / W3C propagator
   // setup wires the global OTel context API, but Effect's tracer goes
 // `ctx.waitUntil(flushTracerProvider())` so SimpleSpanProcessor exports
 // survive request termination.
 export const installTracerProvider = (): boolean => ensureGlobalTracerProvider();
-export const flushTracerProvider = (): Promise<void> =>
+export const flushTracerProvider = async (): Promise<void> => {
-  provider ? provider.forceFlush() : Promise.resolve();
+  if (!provider) return;
+  const startedAt = Date.now();
+  // oxlint-disable-next-line executor/no-try-catch-or-throw -- observability boundary; preserve forceFlush failure behavior while counting it
+  try {
+    await provider.forceFlush();
+    recordForceFlush(Date.now() - startedAt, false);
+  } catch (error) {
+    recordForceFlush(Date.now() - startedAt, true);
+    // oxlint-disable-next-line executor/no-try-catch-or-throw -- preserve original rejection for waitUntil diagnostics
+    throw error;
+  }
+};
 const makeTelemetryLive = (): Layer.Layer<never> =>
   Layer.unwrap(