fix(serverless): flush configured telemetry on Vercel

.github/CODEOWNERS

@@ -61,6 +61,8 @@
 /packages/dd-trace/src/lambda/ @DataDog/serverless-aws @DataDog/apm-serverless
 /packages/dd-trace/src/azure_metadata.js @DataDog/apm-serverless
 /packages/dd-trace/src/serverless.js @DataDog/apm-serverless
+/packages/dd-trace/src/serverless/vercel.js @DataDog/apm-serverless
+/packages/dd-trace/src/flush.js @DataDog/apm-serverless @DataDog/apm-sdk-capabilities-js
 /packages/dd-trace/test/lambda/ @DataDog/serverless-aws @DataDog/apm-serverless
 /packages/dd-trace/test/azure_metadata.spec.js @DataDog/apm-serverless
 /packages/dd-trace/test/serverless.spec.js @DataDog/apm-serverless

packages/datadog-plugin-http2/test/server.spec.js

@@ -9,6 +9,7 @@ const { setImmediate } = require('node:timers/promises')
 
 const { afterEach, beforeEach, describe, it } = require('mocha')
 const sinon = require('sinon')
+const { channel } = require('dc-polyfill')
 
 const agent = require('../../dd-trace/test/plugins/agent')
 const web = require('../../dd-trace/src/plugins/util/web')
@@ -309,6 +310,19 @@ describe('Plugin', () => {
           rawExpectedSchema.server
         )
 
+        it('publishes a close response event', async () => {
+          const emit = sinon.spy()
+          const emitChannel = channel('apm:http2:server:response:emit')
+          emitChannel.subscribe(emit)
+
+          try {
+            await request(http2, `http://localhost:${port}/user`)
+            sinon.assert.calledWithMatch(emit, { eventName: 'close' })
+          } finally {
+            emitChannel.unsubscribe(emit)
+          }
+        })
+
         it('should do automatic instrumentation', done => {
           agent
             .assertFirstTraceSpan({

packages/dd-trace/src/exporters/agent/index.js

@@ -6,6 +6,7 @@ const Writer = require('./writer')
 
 class AgentExporter {
   #timer
+  #activeFlushes = new Set()
 
   constructor (config, prioritySampler) {
     this._config = config
@@ -44,10 +45,10 @@ class AgentExporter {
     const { flushInterval } = this._config
 
     if (flushInterval === 0) {
-      this._writer.flush()
+      this.#flush()
     } else if (this.#timer === undefined) {
       this.#timer = setTimeout(() => {
-        this._writer.flush()
+        this.#flush()
         this.#timer = undefined
       }, flushInterval)
       this.#timer.unref?.()
@@ -57,7 +58,38 @@ class AgentExporter {
   flush (done = () => {}) {
     clearTimeout(this.#timer)
     this.#timer = undefined
-    this._writer.flush(done)
+
+    // Snapshot before the boundary flush so a failed encoding cannot cause a
+    // Vercel lifecycle flush to abandon exports that were already in flight.
+    let activeFlushes = [...this.#activeFlushes]
+    try {
+      this.#flush()
+    } catch (error) {
+      log.error('Failed to flush traces: %s', error.message)
+    }
+    activeFlushes = [...new Set([...activeFlushes, ...this.#activeFlushes])]
+    if (activeFlushes.length === 0) return done()
+
+    let pending = activeFlushes.length
+    const complete = () => {
+      if (--pending === 0) done()
+    }
+    for (const flush of activeFlushes) flush.callbacks.push(complete)
+  }
+
+  #flush () {
+    const flush = { callbacks: [] }
+    this.#activeFlushes.add(flush)
+    const complete = () => {
+      this.#activeFlushes.delete(flush)
+      for (const callback of flush.callbacks) callback()
+    }
+    try {
+      this._writer.flush(complete)
+    } catch (error) {
+      complete()
+      throw error
+    }
   }
 }
 

packages/dd-trace/src/exporters/span-stats/index.js

@@ -3,14 +3,41 @@
 const { Writer } = require('./writer')
 
 class SpanStatsExporter {
+  #activeFlushes = new Set()
+
   constructor (config) {
     this._url = config.url
     this._writer = new Writer({ url: this._url })
   }
 
-  export (payload) {
+  export (payload, done) {
     this._writer.append(payload)
-    this._writer.flush()
+    this.#flush(done)
+  }
+
+  flush (done) {
+    const activeFlushes = [...this.#activeFlushes]
+    let pending = activeFlushes.length + 1
+    const complete = () => {
+      if (--pending === 0) done?.()
+    }
+    for (const flush of activeFlushes) flush.callbacks.push(complete)
+    this.#flush(complete)
+  }
+
+  #flush (done) {
+    const flush = { callbacks: done ? [done] : [] }
+    this.#activeFlushes.add(flush)
+    const complete = () => {
+      this.#activeFlushes.delete(flush)
+      for (const callback of flush.callbacks) callback()
+    }
+    try {
+      this._writer.flush(complete)
+    } catch (error) {
+      complete()
+      throw error
+    }
   }
 }
 

packages/dd-trace/src/flush.js

@@ -0,0 +1,87 @@
+'use strict'
+
+const log = require('./log')
+
+/**
+ * @typedef {(done: () => void) => void | Promise<void>} TelemetryFlusher
+ */
+
+/** @type {Set<TelemetryFlusher>} */
+const telemetryFlushers = new Set()
+
+/**
+ * Registers a configured telemetry pipeline so serverless lifecycle retention
+ * waits for its final export alongside trace delivery.
+ * @param {TelemetryFlusher} flusher
+ * @returns {() => void} Removes this pipeline when its provider is replaced.
+ */
+function registerTelemetryFlusher (flusher) {
+  telemetryFlushers.add(flusher)
+  // Avoid retaining a replaced provider or flushing it alongside the new one.
+  return () => telemetryFlushers.delete(flusher)
+}
+
+/**
+ * Flushes the trace exporter and every registered telemetry pipeline.
+ * @param {{
+ *   _exporter?: { flush?: TelemetryFlusher },
+ *   _processor?: { _stats?: { forceFlush?: TelemetryFlusher } }
+ * }|undefined} tracer
+ * @param {() => void} [done]
+ * @param {{ timeout?: number }} [options]
+ */
+function flushAll (tracer, done, options) {
+  const traceExporter = tracer?._exporter
+  const traceFlusher = traceExporter?.flush
+  const spanStatsFlusher = tracer?._processor?._stats?.forceFlush
+  // TODO: Include DSM after DataStreamsProcessor exposes a completion-aware flush API.
+  let pending = telemetryFlushers.size +
+    (typeof traceFlusher === 'function' ? 1 : 0) +
+    (typeof spanStatsFlusher === 'function' ? 1 : 0)
+  let completed = false
+  let timeout
+
+  const finish = () => {
+    if (completed) return
+    completed = true
+    clearTimeout(timeout)
+    done?.()
+  }
+  const complete = () => {
+    if (--pending === 0) finish()
+  }
+
+  if (pending === 0) return finish()
+  if (options?.timeout) {
+    timeout = setTimeout(() => {
+      log.warn('Timed out waiting for telemetry flush after %dms', options.timeout)
+      finish()
+    }, options.timeout)
+  }
+
+  const flush = flusher => {
+    let flushed = false
+    const onFlushed = error => {
+      if (flushed) return
+      flushed = true
+      if (error) log.error('Error flushing telemetry pipeline:', error)
+      complete()
+    }
+    try {
+      const result = flusher(onFlushed)
+      result?.then(onFlushed, error => onFlushed(error))
+    } catch (error) {
+      onFlushed(error)
+    }
+  }
+
+  if (typeof traceFlusher === 'function') {
+    flush(done => traceFlusher.call(traceExporter, done))
+  }
+  if (typeof spanStatsFlusher === 'function') {
+    flush(done => spanStatsFlusher.call(tracer._processor._stats, done))
+  }
+  for (const flusher of telemetryFlushers) flush(flusher)
+}
+
+module.exports = { flushAll, registerTelemetryFlusher }

packages/dd-trace/src/opentelemetry/logs/batch_log_processor.js

@@ -54,10 +54,34 @@ class BatchLogRecordProcessor {
 
   /**
    * Forces an immediate flush of all pending log records.
-   * @returns {undefined} Promise that resolves when flush is complete
+   * @param {Function} [done] Called after all pending log exports complete
    */
-  forceFlush () {
-    this.#export()
+  forceFlush (done) {
+    this.#clearTimer()
+    // Flush only records present at this boundary. New records belong to the
+    // later request that produced them and must not extend this lifecycle flush.
+    const logRecords = this.#logRecords
+    this.#logRecords = []
+    let pending = 2
+    const complete = () => {
+      if (--pending === 0) done?.()
+    }
+
+    // Join exports already active at this boundary before draining this snapshot.
+    if (typeof this.exporter.flush === 'function') this.exporter.flush(complete)
+    else complete()
+
+    const flushNext = () => {
+      if (logRecords.length === 0) {
+        complete()
+        return
+      }
+
+      // Drain the boundary snapshot one batch at a time.
+      const batch = logRecords.splice(0, this.#maxExportBatchSize)
+      this.exporter.export(batch, flushNext)
+    }
+    flushNext()
   }
 
   /**
@@ -79,6 +103,7 @@ class BatchLogRecordProcessor {
    * @private
    */
   #export () {
+    if (this.#logRecords.length === 0) return
     const logRecords = this.#logRecords.slice(0, this.#maxExportBatchSize)
     this.#logRecords = this.#logRecords.slice(this.#maxExportBatchSize)
 

packages/dd-trace/src/opentelemetry/logs/index.js

@@ -27,10 +27,13 @@ const os = require('os')
  * @package
  */
 
+const { registerTelemetryFlusher } = require('../../flush')
 const LoggerProvider = require('./logger_provider')
 const BatchLogRecordProcessor = require('./batch_log_processor')
 const OtlpHttpLogExporter = require('./otlp_http_log_exporter')
 
+let unregisterTelemetryFlusher
+
 /**
  * Initializes OpenTelemetry Logs support
  * @param {import('../../config/config-base')} config - Tracer configuration instance
@@ -77,8 +80,12 @@ function initializeOpenTelemetryLogs (config) {
   // Create logger provider with processor for Datadog Agent export
   const loggerProvider = new LoggerProvider({ processor })
 
-  // Register the logger provider globally with OpenTelemetry API
+  // Expose this provider to application calls through the OpenTelemetry Logs API.
   loggerProvider.register()
+  // Remove the old provider callback so lifecycle retention flushes only this global provider.
+  unregisterTelemetryFlusher?.()
+  // Include final log batches in lifecycle retention with trace delivery.
+  unregisterTelemetryFlusher = registerTelemetryFlusher(done => loggerProvider.forceFlush(done))
 }
 
 module.exports = {

packages/dd-trace/src/opentelemetry/logs/logger_provider.js

@@ -84,12 +84,15 @@ class LoggerProvider {
 
   /**
    * Forces a flush of all pending log records.
-   * @returns {undefined} Promise that resolves when flush is n ssue cncomplete
+   * @param {Function} [done] Called after all pending log exports complete
    */
-  forceFlush () {
-    if (!this.isShutdown) {
-      return this.processor?.forceFlush()
+  forceFlush (done) {
+    if (this.isShutdown || !this.processor) {
+      done?.()
+      return
     }
+
+    this.processor.forceFlush(done)
   }
 
   /**

packages/dd-trace/src/opentelemetry/metrics/index.js

@@ -6,10 +6,13 @@ const { metrics } = require('@opentelemetry/api')
 
 const { VERSION } = require('../../../../../version')
 const processTags = require('../../process-tags')
+const { registerTelemetryFlusher } = require('../../flush')
 const MeterProvider = require('./meter_provider')
 const PeriodicMetricReader = require('./periodic_metric_reader')
 const OtlpHttpMetricExporter = require('./otlp_http_metric_exporter')
 
+let unregisterTelemetryFlusher
+
 /**
  * @typedef {import('../../config')} Config
  */
@@ -76,6 +79,10 @@ function initializeOpenTelemetryMetrics (config) {
 
   const meterProvider = new MeterProvider({ reader })
   metrics.setGlobalMeterProvider(meterProvider)
+  // Remove the old provider callback so lifecycle retention flushes only this global provider.
+  unregisterTelemetryFlusher?.()
+  // Include the final metric collection and export in lifecycle retention.
+  unregisterTelemetryFlusher = registerTelemetryFlusher(done => meterProvider.forceFlush(done))
 }
 
 function buildResourceAttributes (tags, { reportHostname, otelSemanticsEnabled, service, env, serviceVersion } = {}) {

packages/dd-trace/src/opentelemetry/metrics/meter_provider.js

@@ -49,6 +49,14 @@ class MeterProvider {
     }
     return meter
   }
+
+  /**
+   * @param {Function} [done] Called after the metric export completes
+   */
+  forceFlush (done) {
+    if (this.reader) this.reader.forceFlush(done)
+    else done?.()
+  }
 }
 
 module.exports = MeterProvider

packages/dd-trace/src/opentelemetry/metrics/otlp_http_metric_exporter.js

@@ -34,10 +34,11 @@ class OtlpHttpMetricExporter extends OtlpHttpExporterBase {
    *
    * @param {Map<string, AggregatedMetric>} metrics - Map of metric data to export
    *
-   * @returns {void}
+   * @param {Function} [done] Called after the HTTP export completes
    */
-  export (metrics) {
+  export (metrics, done) {
     if (metrics.size === 0) {
+      done?.({ code: 0 })
       return
     }
 
@@ -56,6 +57,7 @@ class OtlpHttpMetricExporter extends OtlpHttpExporterBase {
       if (result.code === 0) {
         this.recordTelemetry('otel.metrics_export_successes', 1, additionalTags)
       }
+      done?.(result)
     })
   }
 }