@xeonr/upload-pool-sdk 1.2.0 → 1.4.0

package/src/pool.ts

@@ -12,6 +12,7 @@
 import { hostname } from "node:os";
 import { randomBytes } from "node:crypto";
 import { create } from "@bufbuild/protobuf";
+import { SpanStatusCode } from "@opentelemetry/api";
 import {
 	AcceptJobRequestSchema,
 	CompleteJobRequestSchema,
@@ -22,6 +23,13 @@ import { SseClient } from "./sse-client.js";
 import { createJobContext, type JobEnvelope } from "./job-context.js";
 import { NonRetryableError } from "./errors.js";
 import { JsonLogger, type Logger } from "./logger.js";
+import {
+	initTracing,
+	shutdownTracing,
+	recordSpanError,
+	stampJobAttributes,
+	type TracingHandle,
+} from "./tracing.js";
 import type {
 	JobContext,
 	JobHandler,
@@ -33,9 +41,13 @@ export class Pool {
 	private readonly rpc: RpcClients;
 	private readonly sse: SseClient;
 	private readonly logger: Logger;
+	private readonly tracing: TracingHandle;
 	private inFlight = 0;
 	private readonly workerId: string;
 	private readonly capabilities: string[];
+	private runningPromise: Promise<void> | null = null;
+	private runningResolve: (() => void) | null = null;
+	private keepAliveInterval: NodeJS.Timeout | null = null;
 
 	constructor(config: PoolConfig) {
 		this.config = {
@@ -47,6 +59,19 @@ export class Pool {
 		this.logger = (config.logger ?? new JsonLogger()).child({
 			workerId: this.workerId,
 		});
+
+		// Initialise tracing *before* RPC clients so the interceptor in
+		// rpc-clients picks up an active TracerProvider on its first call.
+		// Spans go to pipeline-api's /v1/traces receiver; see tracing.ts
+		// for the isolation rationale wrt the host app's OTel pipeline.
+		this.tracing = initTracing({
+			pipelineEndpoint: config.pipelineEndpoint,
+			poolToken: config.token,
+			workerId: this.workerId,
+			sdkVersion: SDK_VERSION,
+			enabled: config.tracing?.enabled,
+		});
+
 		this.rpc = createRpcClients({
 			apiEndpoint: config.apiEndpoint,
 			pipelineEndpoint: config.pipelineEndpoint,
@@ -82,14 +107,58 @@ export class Pool {
 		}
 	}
 
+	/**
+	 * Start the pool and keep the worker process alive until `stop()` is
+	 * called.
+	 *
+	 * We deliberately return a never-resolving promise (resolved only by
+	 * stop()) for two reasons:
+	 *
+	 *   1. ESM top-level `await pool.start()` is the documented worker
+	 *      entrypoint. If start() resolved as soon as the SSE client was
+	 *      kicked off, the script would end and Node would exit — the
+	 *      EventSource's underlying fetch socket isn't enough to hold the
+	 *      event loop open by itself in Node 20+ (undici unrefs idle
+	 *      connections under some conditions).
+	 *   2. A keep-alive interval lives next to the unresolved promise so
+	 *      even if every socket is unrefed, the loop stays open.
+	 *
+	 * Callers that integrate the pool into a larger app (where some other
+	 * piece holds the event loop) can ignore the returned promise.
+	 */
 	async start(): Promise<void> {
+		if (this.runningPromise) {
+			this.logger.warn("sdk.start.already_running");
+			return this.runningPromise;
+		}
 		this.logger.info("sdk.start");
 		this.sse.start();
+		this.runningPromise = new Promise<void>((resolve) => {
+			this.runningResolve = resolve;
+		});
+		// 60s heartbeat — also doubles as a "still alive" telemetry tick
+		// for log-based dashboards.
+		this.keepAliveInterval = setInterval(() => {
+			this.logger.debug("sdk.heartbeat", { inFlight: this.inFlight });
+		}, 60_000);
+		return this.runningPromise;
 	}
 
 	async stop(): Promise<void> {
 		this.logger.info("sdk.stop", { inFlight: this.inFlight });
 		this.sse.stop();
+		if (this.keepAliveInterval) {
+			clearInterval(this.keepAliveInterval);
+			this.keepAliveInterval = null;
+		}
+		// Flush any buffered spans before exit so the final job's traces
+		// land in the admin UI even on a clean SIGTERM.
+		await shutdownTracing();
+		if (this.runningResolve) {
+			this.runningResolve();
+			this.runningResolve = null;
+		}
+		this.runningPromise = null;
 	}
 
 	private async handleDispatch(envelope: JobEnvelope): Promise<void> {
@@ -110,61 +179,131 @@ export class Pool {
 			return;
 		}
 		this.inFlight++;
-		jobLogger.info("job.dispatched", {
-			filename: envelope.filename,
-			mimeType: envelope.mimeType,
-			inFlight: this.inFlight,
-		});
 
-		const handler = this.resolveHandler(envelope.contentTypeContext.urn);
-		const ctx = createJobContext(envelope, this.rpc, jobLogger);
-		const startedAt = Date.now();
+		// Root the entire job under iq.job, parented to the dispatching
+		// activity's span via the traceparent the pipeline-worker stamped
+		// onto the envelope. All ctx.* + rpc.* spans created inside the
+		// handler will nest under this via OTel's AsyncLocalStorage-backed
+		// context manager.
+		const parentCtx = this.tracing.contextFromEnvelope(
+			envelope.traceContext ?? "",
+			envelope.wfRunId ?? "",
+		);
 
-		try {
-			// AcceptJob — clears the pipeline's accept-timeout.
-			await this.rpc.integrationQueue.acceptJob(
-				create(AcceptJobRequestSchema, {
+		await this.tracing.tracer.startActiveSpan(
+			"iq.job",
+			{
+				attributes: {
+					"iq.filename": envelope.filename,
+					"iq.mime_type": envelope.mimeType,
+				},
+			},
+			parentCtx,
+			async (rootSpan) => {
+				stampJobAttributes(rootSpan, {
+					runId: envelope.wfRunId ?? "",
 					jobId: envelope.jobId,
+					uploadId: envelope.uploadId,
+					urn: envelope.contentTypeContext.urn,
 					workerId: this.workerId,
-					queueToken: this.config.token,
-				}),
-			);
-			jobLogger.info("job.accepted");
+				});
 
-			if (!handler) {
-				jobLogger.warn("job.unhandled", {
-					availableHandlers: this.capabilities,
+				jobLogger.info("job.dispatched", {
+					filename: envelope.filename,
+					mimeType: envelope.mimeType,
+					inFlight: this.inFlight,
 				});
-				await this.reportError(
-					envelope.jobId,
-					ctx,
-					jobLogger,
-					new NonRetryableError(`no handler for URN ${envelope.contentTypeContext.urn}`),
-				);
-				return;
-			}
-
-			await handler(ctx);
-
-			await this.rpc.integrationQueue.completeJob(
-				create(CompleteJobRequestSchema, {
-					jobId: envelope.jobId,
+
+				const handler = this.resolveHandler(envelope.contentTypeContext.urn);
+				const ctx = createJobContext(envelope, this.rpc, jobLogger, {
 					workerId: this.workerId,
-					queueToken: this.config.token,
-				}),
-			);
-			jobLogger.info("job.completed", {
-				durationMs: Date.now() - startedAt,
-			});
-		} catch (err) {
-			jobLogger.error("job.failed", {
-				err,
-				durationMs: Date.now() - startedAt,
-			});
-			await this.reportError(envelope.jobId, ctx, jobLogger, err as Error);
-		} finally {
-			this.inFlight--;
-		}
+				});
+				const startedAt = Date.now();
+
+				try {
+					// AcceptJob — clears the pipeline's accept-timeout.
+					await this.rpc.integrationQueue.acceptJob(
+						create(AcceptJobRequestSchema, {
+							jobId: envelope.jobId,
+							workerId: this.workerId,
+							queueToken: this.config.token,
+						}),
+					);
+					jobLogger.info("job.accepted");
+
+					if (!handler) {
+						jobLogger.warn("job.unhandled", {
+							availableHandlers: this.capabilities,
+						});
+						await this.reportError(
+							envelope.jobId,
+							ctx,
+							jobLogger,
+							new NonRetryableError(
+								`no handler for URN ${envelope.contentTypeContext.urn}`,
+							),
+						);
+						rootSpan.setStatus({
+							code: SpanStatusCode.ERROR,
+							message: "no handler",
+						});
+						return;
+					}
+
+					// iq.handler wraps the customer's handler function so its
+					// own latency is visible as a single row (excluding the
+					// accept/complete bookends). Children appear nested.
+					await this.tracing.tracer.startActiveSpan(
+						"iq.handler",
+						async (handlerSpan) => {
+							stampJobAttributes(handlerSpan, {
+								runId: envelope.wfRunId ?? "",
+								jobId: envelope.jobId,
+								uploadId: envelope.uploadId,
+								urn: envelope.contentTypeContext.urn,
+								workerId: this.workerId,
+							});
+							try {
+								await handler(ctx);
+								handlerSpan.setStatus({ code: SpanStatusCode.OK });
+							} catch (err) {
+								recordSpanError(handlerSpan, err);
+								throw err;
+							} finally {
+								handlerSpan.end();
+							}
+						},
+					);
+
+					await this.rpc.integrationQueue.completeJob(
+						create(CompleteJobRequestSchema, {
+							jobId: envelope.jobId,
+							workerId: this.workerId,
+							queueToken: this.config.token,
+						}),
+					);
+					jobLogger.info("job.completed", {
+						durationMs: Date.now() - startedAt,
+					});
+					rootSpan.setStatus({ code: SpanStatusCode.OK });
+				} catch (err) {
+					jobLogger.error("job.failed", {
+						err,
+						durationMs: Date.now() - startedAt,
+					});
+					recordSpanError(rootSpan, err);
+					await this.reportError(
+						envelope.jobId,
+						ctx,
+						jobLogger,
+						err as Error,
+					);
+				} finally {
+					rootSpan.end();
+					this.inFlight--;
+				}
+			},
+		);
 	}
 
 	private resolveHandler(urn: string): JobHandler | undefined {
@@ -197,4 +336,4 @@ export class Pool {
 	}
 }
 
-const SDK_VERSION = "1.2.0";
+const SDK_VERSION = "1.3.0";

package/src/rpc-clients.ts

@@ -13,10 +13,20 @@
  * pipeline-api) so we keep them on separate transports and don't try to
  * share a baseUrl.
  *
- * Each request is wrapped in a logging interceptor that emits one
- * `rpc.request` line on dispatch and one of `rpc.response` /
- * `rpc.error` on completion. Latency and connect-error codes are
- * captured so failed RPCs are diagnosable from worker logs alone.
+ * Each request is wrapped in a tracing-and-logging interceptor:
+ *   - emits one `iq.callback.<Method>` span per call, child of the
+ *     active job span (so the admin UI sees a row per RPC under the
+ *     dispatching parseUpload run)
+ *   - injects W3C `traceparent` on outbound headers *for the pipeline
+ *     transport only* — calls to `apiEndpoint` ship without traceparent
+ *     so upl-im-api's own OTel pipeline doesn't pick up our trace_id
+ *     and conflate the two backends
+ *   - captures sanitised request/response bodies as span events (4 KB
+ *     cap, keys matching token/secret/auth/... redacted), so debugging
+ *     a worker callback failure doesn't require correlating with
+ *     server-side logs
+ *   - logs `rpc.request` / `rpc.response` / `rpc.error` at info/error
+ *     for the legacy log-only consumer surface that predates tracing
  */
 import { createConnectTransport } from "@connectrpc/connect-node";
 import {
@@ -26,41 +36,167 @@ import {
 	ConnectError,
 	Code,
 } from "@connectrpc/connect";
+import {
+	context as otelContext,
+	SpanStatusCode,
+	trace,
+} from "@opentelemetry/api";
 import { InternalUploadsService } from "./protocol/uplim/api/v1/uploads_pb.js";
 import { IntegrationQueueService } from "./protocol/uplim/workflow/v1/integration_queue_pb.js";
 import type { Logger } from "./logger.js";
+import {
+	SPAN_ATTR,
+	getTracingHandle,
+	recordSpanError,
+	sanitizeRpcBody,
+	spanContextToTraceparent,
+} from "./tracing.js";
 
 export interface RpcClients {
 	internalUploads: Client<typeof InternalUploadsService>;
 	integrationQueue: Client<typeof IntegrationQueueService>;
 }
 
-function loggingInterceptor(logger: Logger, target: string): Interceptor {
+type RpcTarget = "api" | "pipeline";
+
+function tracingLoggingInterceptor(
+	logger: Logger,
+	target: RpcTarget,
+): Interceptor {
 	return (next) => async (req) => {
-		const startedAt = Date.now();
+		const handle = getTracingHandle();
 		const method = `${req.service.typeName}/${req.method.name}`;
+
+		// Inject W3C traceparent on the pipeline target *only*. upl-im-api
+		// runs its own OTel pipeline against a different backend; passing
+		// our trace_id over would graft our worker activity into its trees,
+		// muddying its admin dashboards. The pool token in the request
+		// body / queue token in headers remain the trust boundary either
+		// way.
+		if (target === "pipeline") {
+			const activeSpan = trace.getActiveSpan();
+			if (activeSpan) {
+				req.header.set(
+					"traceparent",
+					spanContextToTraceparent(activeSpan.spanContext()),
+				);
+			}
+		}
+
+		const startedAt = Date.now();
 		logger.debug("rpc.request", { method, target });
-		try {
-			const res = await next(req);
-			logger.debug("rpc.response", {
-				method,
-				target,
-				durationMs: Date.now() - startedAt,
-			});
-			return res;
-		} catch (err) {
-			const code =
-				err instanceof ConnectError ? Code[err.code] : undefined;
-			const message = err instanceof Error ? err.message : String(err);
-			logger.error("rpc.error", {
-				method,
-				target,
-				code,
-				message,
-				durationMs: Date.now() - startedAt,
-			});
-			throw err;
+
+		// When tracing is uninitialised (e.g. tests, dry-runs), fall through
+		// to a logging-only path so we don't crash on a null tracer.
+		if (!handle) {
+			try {
+				const res = await next(req);
+				logger.debug("rpc.response", {
+					method,
+					target,
+					durationMs: Date.now() - startedAt,
+				});
+				return res;
+			} catch (err) {
+				const code =
+					err instanceof ConnectError ? Code[err.code] : undefined;
+				const message = err instanceof Error ? err.message : String(err);
+				logger.error("rpc.error", {
+					method,
+					target,
+					code,
+					message,
+					durationMs: Date.now() - startedAt,
+				});
+				throw err;
+			}
 		}
+
+		return handle.tracer.startActiveSpan(
+			`iq.callback.${req.method.name}`,
+			async (span) => {
+				span.setAttribute(SPAN_ATTR.RPC_METHOD, method);
+				span.setAttribute(SPAN_ATTR.RPC_SERVICE, req.service.typeName);
+				span.setAttribute(SPAN_ATTR.RPC_TARGET, target);
+
+				// Inherit job-scoped attributes from the active parent span
+				// (set by Pool.handleDispatch via stampJobAttributes). Without
+				// this, the span lacks `pipeline.run_id` and gets dropped by
+				// the pipeline-api receiver. Reading them off the active span
+				// is the cheapest cross-span propagation we can do without a
+				// custom OTel context manager.
+				const parent = trace.getActiveSpan();
+				const parentAttrs = (parent as unknown as {
+					attributes?: Record<string, unknown>;
+				})?.attributes;
+				if (parentAttrs) {
+					for (const key of [
+						SPAN_ATTR.RUN_ID,
+						SPAN_ATTR.JOB_ID,
+						SPAN_ATTR.UPLOAD_ID,
+						SPAN_ATTR.URN,
+						SPAN_ATTR.WORKER_ID,
+					]) {
+						const v = parentAttrs[key];
+						if (typeof v === "string" && v.length > 0) {
+							span.setAttribute(key, v);
+						}
+					}
+				}
+
+				// Body capture: request first, response on success, error
+				// payload on failure. Sanitiser truncates to 4 KB and redacts
+				// token-shaped keys (see tracing.ts).
+				try {
+					span.addEvent("rpc.request", {
+						body: sanitizeRpcBody(req.message),
+					});
+				} catch {
+					/* never let body capture crash the RPC */
+				}
+
+				try {
+					const res = await next(req);
+					const durationMs = Date.now() - startedAt;
+					span.setAttribute(SPAN_ATTR.RPC_DURATION_MS, durationMs);
+					span.setAttribute(SPAN_ATTR.RPC_CODE, "ok");
+					try {
+						span.addEvent("rpc.response", {
+							body: sanitizeRpcBody(res.message),
+						});
+					} catch {
+						/* body capture is best-effort */
+					}
+					span.setStatus({ code: SpanStatusCode.OK });
+					logger.debug("rpc.response", {
+						method,
+						target,
+						durationMs,
+					});
+					return res;
+				} catch (err) {
+					const durationMs = Date.now() - startedAt;
+					const code =
+						err instanceof ConnectError ? Code[err.code] : "internal";
+					const message =
+						err instanceof Error ? err.message : String(err);
+					span.setAttribute(SPAN_ATTR.RPC_DURATION_MS, durationMs);
+					span.setAttribute(SPAN_ATTR.RPC_CODE, String(code));
+					span.addEvent("rpc.error", { code: String(code), message });
+					recordSpanError(span, err);
+					logger.error("rpc.error", {
+						method,
+						target,
+						code,
+						message,
+						durationMs,
+					});
+					throw err;
+				} finally {
+					span.end();
+				}
+			},
+		);
 	};
 }
 
@@ -76,15 +212,19 @@ export function createRpcClients(config: RpcClientsConfig): RpcClients {
 	const apiTransport = createConnectTransport({
 		baseUrl: config.apiEndpoint,
 		httpVersion: "1.1",
-		interceptors: [loggingInterceptor(rpcLogger, "api")],
+		interceptors: [tracingLoggingInterceptor(rpcLogger, "api")],
 	});
 	const pipelineTransport = createConnectTransport({
 		baseUrl: config.pipelineEndpoint,
 		httpVersion: "1.1",
-		interceptors: [loggingInterceptor(rpcLogger, "pipeline")],
+		interceptors: [tracingLoggingInterceptor(rpcLogger, "pipeline")],
 	});
 	return {
 		internalUploads: createClient(InternalUploadsService, apiTransport),
 		integrationQueue: createClient(IntegrationQueueService, pipelineTransport),
 	};
 }
+
+// keep otelContext import alive for future extensions that may want to
+// detach span scopes around streaming responses.
+void otelContext;