fix(serverless): flush configured telemetry on Vercel

.github/CODEOWNERS

@@ -62,6 +62,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/dogstatsd.js

@@ -67,23 +67,23 @@ class DogStatsDClient {
     this._add(stat, value, TYPE_HISTOGRAM, tags)
   }
 
-  flush () {
+  flush (done) {
     const queue = this._enqueue()
 
-    if (queue.length === 0) return
+    if (queue.length === 0) return done?.()
 
     log.debug('Flushing %s metrics via %s', queue.length, this._httpOptions ? 'HTTP' : 'UDP')
 
     this._queue = []
 
     if (this._httpOptions) {
-      this._sendHttp(queue)
+      this._sendHttp(queue, done)
     } else {
-      this._sendUdp(queue)
+      this._sendUdp(queue, done)
     }
   }
 
-  _sendHttp (queue) {
+  _sendHttp (queue, done) {
     const buffer = Buffer.concat(queue)
     request(buffer, this._httpOptions, (err) => {
       if (err) {
@@ -95,32 +95,46 @@ class DogStatsDClient {
           // options. Either way, we can give UDP a try.
           this._httpOptions = undefined
         }
-        this._sendUdp(queue)
+        this._sendUdp(queue, done)
+      } else {
+        done?.()
       }
     })
   }
 
-  _sendUdp (queue) {
+  _sendUdp (queue, done) {
     // dgram resolves the local address via the instrumented dns.lookup when it
     // binds on first send; the noop store keeps that self-traffic off the trace.
     legacyStorage.run({ noop: true }, () => {
       if (this._family === 0) {
         this.#lookup(this._host, (error, address, family) => {
-          if (error) return log.error('DogStatsDClient: Host not found', error)
-          this._sendUdpFromQueue(queue, address, family)
+          if (error) {
+            log.error('DogStatsDClient: Host not found', error)
+            return done?.()
+          }
+          this._sendUdpFromQueue(queue, address, family, done)
         })
       } else {
-        this._sendUdpFromQueue(queue, this._host, this._family)
+        this._sendUdpFromQueue(queue, this._host, this._family, done)
       }
     })
   }
 
-  _sendUdpFromQueue (queue, address, family) {
+  _sendUdpFromQueue (queue, address, family, done) {
     const socket = family === 6 ? this._udp6 : this._udp4
+    let pending = queue.length
+    const complete = () => {
+      if (--pending === 0) done?.()
+    }
 
     for (const buffer of queue) {
       log.debug('Sending to DogStatsD: %s', buffer)
-      socket.send(buffer, 0, buffer.length, this._port, address)
+      try {
+        socket.send(buffer, 0, buffer.length, this._port, address, complete)
+      } catch (error) {
+        log.error('DogStatsDClient: UDP error sending metrics', error)
+        complete()
+      }
     }
   }
 
@@ -212,12 +226,12 @@ class MetricsAggregationClient {
     this.reset()
   }
 
-  flush () {
+  flush (done) {
     this._captureCounters()
     this._captureGauges()
     this._captureHistograms()
 
-    this._client.flush()
+    this._client.flush(done)
   }
 
   reset () {

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,56 @@
 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) {
+    if (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._writer.append(payload)
+      try {
+        this.#flush(complete)
+      } catch {
+        // `#flush` has notified the boundary request; keep waiting for prior exports.
+      }
+      return
+    }
     this._writer.append(payload)
-    this._writer.flush()
+    this.#flush()
+  }
+
+  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)