fix(tracing): upgrade to otel 2.0

Author: MichaelDoyle

js/core/package.json

@@ -26,12 +26,12 @@
   "author": "genkit",
   "license": "Apache-2.0",
   "dependencies": {
-    "@opentelemetry/api": "^1.9.0",
-    "@opentelemetry/context-async-hooks": "~1.25.0",
-    "@opentelemetry/core": "~1.25.0",
-    "@opentelemetry/sdk-metrics": "~1.25.0",
-    "@opentelemetry/sdk-node": "^0.52.0",
-    "@opentelemetry/sdk-trace-base": "~1.25.0",
+    "@opentelemetry/api": "^1.9.1",
+    "@opentelemetry/context-async-hooks": "^2.6.1",
+    "@opentelemetry/core": "^2.6.1",
+    "@opentelemetry/sdk-metrics": "^2.6.1",
+    "@opentelemetry/sdk-node": "^0.214.0",
+    "@opentelemetry/sdk-trace-base": "^2.6.1",
     "@types/json-schema": "^7.0.15",
     "ajv": "^8.12.0",
     "ajv-formats": "^3.0.1",

js/core/src/tracing.ts

@@ -100,9 +100,10 @@ export async function enableTelemetry(
   if (isOTelInitializationDisabled()) {
     return;
   }
-  global[instrumentationKey] =
-    telemetryConfig instanceof Promise ? telemetryConfig : Promise.resolve();
-  return getTelemetryProvider().enableTelemetry(telemetryConfig);
+  const instrumentationPromise =
+    getTelemetryProvider().enableTelemetry(telemetryConfig);
+  global[instrumentationKey] = instrumentationPromise;
+  return instrumentationPromise;
 }
 
 /**

js/core/src/tracing/exporter.ts

@@ -73,7 +73,7 @@ export class TraceServerExporter implements SpanExporter {
       displayName: span.name,
       links: span.links,
       spanKind: SpanKind[span.kind],
-      parentSpanId: span.parentSpanId,
+      parentSpanId: span.parentSpanContext?.spanId,
       sameProcessAsParentSpan: { value: !span.spanContext().isRemote },
       status: span.status,
       timeEvents: {
@@ -86,17 +86,17 @@ export class TraceServerExporter implements SpanExporter {
         })),
       },
     };
-    if (span.instrumentationLibrary !== undefined) {
+    if (span.instrumentationScope !== undefined) {
       spanData.instrumentationLibrary = {
-        name: span.instrumentationLibrary.name,
+        name: span.instrumentationScope.name,
       };
-      if (span.instrumentationLibrary.schemaUrl !== undefined) {
+      if (span.instrumentationScope.schemaUrl !== undefined) {
         spanData.instrumentationLibrary.schemaUrl =
-          span.instrumentationLibrary.schemaUrl;
+          span.instrumentationScope.schemaUrl;
       }
-      if (span.instrumentationLibrary.version !== undefined) {
+      if (span.instrumentationScope.version !== undefined) {
         spanData.instrumentationLibrary.version =
-          span.instrumentationLibrary.version;
+          span.instrumentationScope.version;
       }
     }
     deleteUndefinedProps(spanData);
@@ -127,19 +127,19 @@ export class TraceServerExporter implements SpanExporter {
         await this.save(traceId, traces[traceId]);
       } catch (e) {
         error = true;
-        logger.error(`Failed to save trace ${traceId}`, e);
-      }
-      if (done) {
-        return done({
-          code: error ? ExportResultCode.FAILED : ExportResultCode.SUCCESS,
-        });
+        logger.defaultLogger.error(`Failed to save trace ${traceId}`, e);
       }
     }
+    if (done) {
+      return done({
+        code: error ? ExportResultCode.FAILED : ExportResultCode.SUCCESS,
+      });
+    }
   }
 
-  private async save(traceId, spans: ReadableSpan[]): Promise<void> {
+  private async save(traceId: string, spans: ReadableSpan[]): Promise<void> {
     if (!telemetryServerUrl) {
-      logger.debug(
+      logger.defaultLogger.debug(
         `Telemetry server is not configured, trace ${traceId} not saved!`
       );
       return;

js/core/src/tracing/node-telemetry-provider.ts

@@ -14,10 +14,14 @@
  * limitations under the License.
  */
 
+import { Context } from '@opentelemetry/api';
+import { MetricReader } from '@opentelemetry/sdk-metrics';
 import { NodeSDK } from '@opentelemetry/sdk-node';
 import {
   BatchSpanProcessor,
   SimpleSpanProcessor,
+  type ReadableSpan,
+  type Span,
   type SpanProcessor,
 } from '@opentelemetry/sdk-trace-base';
 import { logger } from '../logging.js';
@@ -29,6 +33,7 @@ import { RealtimeSpanProcessor } from './realtime-span-processor.js';
 
 let telemetrySDK: NodeSDK | null = null;
 let nodeOtelConfig: TelemetryConfig | null = null;
+let isSigtermHandlerRegistered = false;
 
 export function initNodeTelemetryProvider() {
   setTelemetryProvider({
@@ -37,6 +42,72 @@ export function initNodeTelemetryProvider() {
   });
 }
 
+/**
+ * MultiSpanProcessor is a multiplexer that allows Genkit to register multiple
+ * span processors reliably.
+ *
+ * We provide it as the sole entry in the `spanProcessors` array to the
+ * OpenTelemetry NodeSDK because the SDK's internal logic for merging
+ * `traceExporter`, `spanProcessor` (singular), and `spanProcessors` (plural)
+ * varies across versions and can lead to exporters being silently overwritten
+ * or ignored.
+ *
+ * By wrapping all processors into this single delegate, we guarantee that
+ * Genkit's internal telemetry and any user-provided exporters both receive
+ * every start and end event.
+ */
+export class MultiSpanProcessor implements SpanProcessor {
+  constructor(private processors: SpanProcessor[]) {}
+  onStart(span: Span, parentContext: Context) {
+    this.processors.forEach((p) => {
+      try {
+        p.onStart(span, parentContext);
+      } catch (e) {
+        logger.defaultLogger.error(
+          `Error in span processor (${p.constructor.name}) onStart: ${e}`
+        );
+      }
+    });
+  }
+  onEnd(span: ReadableSpan) {
+    this.processors.forEach((p) => {
+      try {
+        p.onEnd(span);
+      } catch (e) {
+        logger.defaultLogger.error(
+          `Error in span processor (${p.constructor.name}) onEnd: ${e}`
+        );
+      }
+    });
+  }
+  async forceFlush() {
+    await Promise.all(
+      this.processors.map(async (p) => {
+        try {
+          await p.forceFlush();
+        } catch (e) {
+          logger.defaultLogger.error(
+            `Error in span processor (${p.constructor.name}) forceFlush: ${e}`
+          );
+        }
+      })
+    );
+  }
+  async shutdown() {
+    await Promise.all(
+      this.processors.map(async (p) => {
+        try {
+          await p.shutdown();
+        } catch (e) {
+          logger.defaultLogger.error(
+            `Error in span processor (${p.constructor.name}) shutdown: ${e}`
+          );
+        }
+      })
+    );
+  }
+}
+
 /**
  * Enables tracing and metrics open telemetry configuration.
  */
@@ -52,23 +123,47 @@ async function enableTelemetry(
       ? await telemetryConfig
       : telemetryConfig;
 
-  nodeOtelConfig = telemetryConfig || {};
+  if (telemetrySDK) {
+    await cleanUpTracing();
+  }
+
+  nodeOtelConfig = { ...telemetryConfig };
 
   const processors: SpanProcessor[] = [createTelemetryServerProcessor()];
-  if (nodeOtelConfig.traceExporter) {
-    throw new Error('Please specify spanProcessors instead.');
-  }
   if (nodeOtelConfig.spanProcessors) {
     processors.push(...nodeOtelConfig.spanProcessors);
   }
   if (nodeOtelConfig.spanProcessor) {
     processors.push(nodeOtelConfig.spanProcessor);
-    delete nodeOtelConfig.spanProcessor;
   }
-  nodeOtelConfig.spanProcessors = processors;
+  if (nodeOtelConfig.traceExporter) {
+    processors.push(new BatchSpanProcessor(nodeOtelConfig.traceExporter));
+  }
+
+  if (processors.length > 1) {
+    nodeOtelConfig.spanProcessors = [new MultiSpanProcessor(processors)];
+  } else {
+    nodeOtelConfig.spanProcessors = processors;
+  }
+  delete nodeOtelConfig.spanProcessor;
+  delete nodeOtelConfig.traceExporter;
+
   telemetrySDK = new NodeSDK(nodeOtelConfig);
   telemetrySDK.start();
-  process.on('SIGTERM', async () => await cleanUpTracing());
+
+  if (!isSigtermHandlerRegistered) {
+    let isShuttingDown = false;
+    const shutdownHandler = (signal: string) => () => {
+      if (isShuttingDown) return; // Prevent SIGTERM + SIGINT race
+      isShuttingDown = true;
+      cleanUpTracing().finally(() => {
+        process.kill(process.pid, signal);
+      });
+    };
+    process.once('SIGTERM', shutdownHandler('SIGTERM'));
+    process.once('SIGINT', shutdownHandler('SIGINT'));
+    isSigtermHandlerRegistered = true;
+  }
 }
 
 async function cleanUpTracing(): Promise<void> {
@@ -79,10 +174,20 @@ async function cleanUpTracing(): Promise<void> {
   // Metrics are not flushed as part of the shutdown operation. If metrics
   // are enabled, we need to manually flush them *before* the reader
   // receives shutdown order.
-  await maybeFlushMetrics();
-  await telemetrySDK.shutdown();
-  logger.debug('OpenTelemetry SDK shut down.');
-  telemetrySDK = null;
+  try {
+    await maybeFlushMetrics();
+  } catch (e) {
+    logger.defaultLogger.error(`Error flushing metrics during shutdown: ${e}`);
+  }
+
+  try {
+    await telemetrySDK.shutdown();
+    logger.debug('OpenTelemetry SDK shut down.');
+  } catch (e) {
+    logger.defaultLogger.error(`Error shutting down OpenTelemetry SDK: ${e}`);
+  } finally {
+    telemetrySDK = null;
+  }
 }
 
 /**
@@ -102,11 +207,23 @@ function createTelemetryServerProcessor(): SpanProcessor {
 }
 
 /** Flush metrics if present. */
-function maybeFlushMetrics(): Promise<void> {
+async function maybeFlushMetrics(): Promise<void> {
+  const readers: MetricReader[] = [];
   if (nodeOtelConfig?.metricReader) {
-    return nodeOtelConfig.metricReader.forceFlush();
+    readers.push(nodeOtelConfig.metricReader as MetricReader);
+  }
+  if (nodeOtelConfig?.metricReaders) {
+    readers.push(...(nodeOtelConfig.metricReaders as MetricReader[]));
   }
-  return Promise.resolve();
+  await Promise.all(
+    readers.map(async (r) => {
+      try {
+        await r.forceFlush();
+      } catch (e) {
+        logger.defaultLogger.error(`Error flushing metrics: ${e}`);
+      }
+    })
+  );
 }
 
 /**

js/core/tests/action_test.ts

@@ -21,7 +21,7 @@ import { z } from 'zod';
 import { action, defineAction } from '../src/action.js';
 import { initNodeFeatures } from '../src/node.js';
 import { Registry } from '../src/registry.js';
-import { enableTelemetry } from '../src/tracing.js';
+import { enableTelemetry, flushTracing } from '../src/tracing.js';
 import { TestSpanExporter } from './utils.js';
 
 initNodeFeatures();
@@ -33,8 +33,9 @@ enableTelemetry({
 
 describe('action', () => {
   var registry: Registry;
-  beforeEach(() => {
+  beforeEach(async () => {
     registry = new Registry();
+    await flushTracing();
     spanExporter.exportedSpans = [];
   });
 
@@ -289,6 +290,7 @@ describe('action', () => {
     act.__action.key = 'some-custom-key';
 
     await act();
+    await flushTracing();
 
     assert.strictEqual(spanExporter.exportedSpans.length, 1);
     assert.strictEqual(

js/core/tests/flow_test.ts

@@ -21,7 +21,7 @@ import { defineFlow, run } from '../src/flow.js';
 import { defineAction, getContext, z } from '../src/index.js';
 import { initNodeFeatures } from '../src/node.js';
 import { Registry } from '../src/registry.js';
-import { enableTelemetry } from '../src/tracing.js';
+import { enableTelemetry, flushTracing } from '../src/tracing.js';
 import { TestSpanExporter } from './utils.js';
 
 initNodeFeatures();
@@ -48,10 +48,12 @@ function createTestFlow(registry: Registry) {
 describe('flow', () => {
   let registry: Registry;
 
-  beforeEach(() => {
+  beforeEach(async () => {
     // Skips starting reflection server.
     delete process.env.GENKIT_ENV;
     registry = new Registry();
+    await flushTracing();
+    spanExporter.exportedSpans = [];
   });
 
   describe('runFlow', () => {
@@ -370,6 +372,8 @@ describe('flow', () => {
         telemetryLabels: { custom: 'label' },
       });
 
+      await flushTracing();
+
       assert.equal(result, 'bar foo');
       assert.strictEqual(spanExporter.exportedSpans.length, 1);
       assert.strictEqual(spanExporter.exportedSpans[0].displayName, 'testFlow');
@@ -394,6 +398,8 @@ describe('flow', () => {
       });
       const result = await output;
 
+      await flushTracing();
+
       assert.equal(result, 'bar foo');
       assert.strictEqual(spanExporter.exportedSpans.length, 1);
       assert.strictEqual(spanExporter.exportedSpans[0].displayName, 'testFlow');
@@ -445,6 +451,8 @@ describe('flow', () => {
         context: { user: 'pavel' },
       });
 
+      await flushTracing();
+
       assert.equal(result, 'foo bar');
       assert.strictEqual(spanExporter.exportedSpans.length, 3);