History

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
  1. apps/cloud/src/mcp/agent-handler.ts +3 −1
  2. apps/cloud/src/observability/memory-metrics.ts +329 −0
  3. apps/cloud/src/observability/telemetry.ts +41 −13
apps/cloud/src/mcp/agent-handler.ts+3−1
@@ -14,6 +14,7 @@ import { · 13 unmodified lines
1414 } from "@executor-js/cloudflare/mcp/do-headers";
1515 import type { McpSessionProps } from "@executor-js/cloudflare/mcp/agent-durable-object";
1616
17+import { wrapMcpSseResponse } from "../observability/memory-metrics";
1718 import { cloudMcpAuth } from "./auth-provider";
1819 import { McpSessionDOSqlite } from "./session-durable-object";
1920
@@ -158,6 +159,7 @@ export const makeCloudMcpAgentHandler = () => { · 138 unmodified lines
158159 },
159160 resource,
160161 );
161− return serve.fetch(forwarded, env, ctx);
162+ const response = await serve.fetch(forwarded, env, ctx);
163+ return wrapMcpSseResponse(request, env, response);
162164 };
163165 };
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
@@ -45,6 +45,13 @@ import { ATTR_SERVICE_NAME, ATTR_SERVICE_VERSION } from "@opentelemetry/semantic · 44 unmodified lines
4545 import { env } from "cloudflare:workers";
4646 import { Effect, Layer } from "effect";
4747
48+import {
49+ CountingSpanExporter,
50+ CountingSpanProcessor,
51+ OTEL_MAX_SPAN_QUEUE_SIZE,
52+ recordForceFlush,
53+} from "./memory-metrics";
54+
4855 const SERVICE_NAME = "executor-cloud";
4956 const SERVICE_VERSION = "1.0.0";
5057
@@ -60,14 +67,9 @@ const ensureGlobalTracerProvider = (): boolean => { · 9 unmodified lines
6067 [ATTR_SERVICE_NAME]: SERVICE_NAME,
6168 [ATTR_SERVICE_VERSION]: SERVICE_VERSION,
6269 }),
63− spanProcessors: [
64− // Batch, not Simple: SimpleSpanProcessor issues one synchronous fetch
65− // per span END — a busy MCP request emits dozens of spans, and the
66− // serialized export fetches add seconds of latency inside the request
67− // (enough to flip pause/resume timing in e2e). Batch buffers and
68− // ships on a timer; `ctx.waitUntil(flushTracerProvider())` at the end
69− // of each request still drains the buffer before the isolate exits.
70− new BatchSpanProcessor(
70+ spanProcessors: (() => {
71+ let countingProcessor: CountingSpanProcessor;
72+ const exporter = new CountingSpanExporter(
7173 new OTLPTraceExporter({
7274 url: env.AXIOM_TRACES_URL ?? "https://api.axiom.co/v1/traces",
7375 headers: {
@@ -75,9 +77,24 @@ const ensureGlobalTracerProvider = (): boolean => { · 1 unmodified line
7577 "X-Axiom-Dataset": env.AXIOM_DATASET ?? "executor-cloud",
7678 },
7779 }),
78− { scheduledDelayMillis: 1_000, maxExportBatchSize: 512 },
79− ),
80− ],
80+ (spans) => countingProcessor.recordExportAttempt(spans),
81+ );
82+ // Batch, not Simple: SimpleSpanProcessor issues one synchronous fetch
83+ // per span END — a busy MCP request emits dozens of spans, and the
84+ // serialized export fetches add seconds of latency inside the request
85+ // (enough to flip pause/resume timing in e2e). Batch buffers and
86+ // ships on a timer; `ctx.waitUntil(flushTracerProvider())` at the end
87+ // of each request still drains the buffer before the isolate exits.
88+ countingProcessor = new CountingSpanProcessor(
89+ new BatchSpanProcessor(exporter, {
90+ scheduledDelayMillis: 1_000,
91+ maxExportBatchSize: 512,
92+ maxQueueSize: OTEL_MAX_SPAN_QUEUE_SIZE,
93+ }),
94+ OTEL_MAX_SPAN_QUEUE_SIZE,
95+ );
96+ return [countingProcessor];
97+ })(),
8198 });
8299 // Skip `provider.register()` — its StackContextManager / W3C propagator
83100 // setup wires the global OTel context API, but Effect's tracer goes
@@ -92,8 +109,19 @@ const ensureGlobalTracerProvider = (): boolean => { · 8 unmodified lines
92109 // `ctx.waitUntil(flushTracerProvider())` so SimpleSpanProcessor exports
93110 // survive request termination.
94111 export const installTracerProvider = (): boolean => ensureGlobalTracerProvider();
95−export const flushTracerProvider = (): Promise<void> =>
96− provider ? provider.forceFlush() : Promise.resolve();
112+export const flushTracerProvider = async (): Promise<void> => {
113+ if (!provider) return;
114+ const startedAt = Date.now();
115+ // oxlint-disable-next-line executor/no-try-catch-or-throw -- observability boundary; preserve forceFlush failure behavior while counting it
116+ try {
117+ await provider.forceFlush();
118+ recordForceFlush(Date.now() - startedAt, false);
119+ } catch (error) {
120+ recordForceFlush(Date.now() - startedAt, true);
121+ // oxlint-disable-next-line executor/no-try-catch-or-throw -- preserve original rejection for waitUntil diagnostics
122+ throw error;
123+ }
124+};
97125
98126 const makeTelemetryLive = (): Layer.Layer<never> =>
99127 Layer.unwrap(