dd-trace 6.12.0 → 6.14.0

package/packages/dd-trace/src/dogstatsd.js

@@ -2,49 +2,186 @@
 
 const dgram = require('dgram')
 const isIP = require('net').isIP
+const { performance } = require('node:perf_hooks')
 
+const { channel } = require('dc-polyfill')
+
+const tracerVersion = require('../../../package.json').version
 const { storage } = require('../../datadog-core')
 const request = require('./exporters/common/request')
 const log = require('./log')
 const Histogram = require('./histogram')
 const { entityId } = require('./exporters/common/docker')
+const { registerTelemetryFlusher } = require('./flush')
+const { createServerlessDeliveryTracker } = require('./serverless')
 
 const legacyStorage = storage('legacy')
 
 const MAX_BUFFER_SIZE = 1024 // limit from the agent
+const TELEMETRY_INTERVAL = 10_000
 
 const TYPE_COUNTER = 'c'
 const TYPE_GAUGE = 'g'
 const TYPE_DISTRIBUTION = 'd'
 const TYPE_HISTOGRAM = 'h'
 
+const TYPE_COUNTER_INDEX = 0
+const TYPE_GAUGE_INDEX = 1
+const TYPE_DISTRIBUTION_INDEX = 2
+const TYPE_HISTOGRAM_INDEX = 3
+const TYPE_LABELS = ['count', 'gauge', 'distribution', 'histogram']
+
+/**
+ * @typedef {'c'|'g'|'d'|'h'} MetricType
+ */
+
+/**
+ * @typedef {object} DogStatsDBufferState
+ * @property {string} message
+ * @property {number} offset
+ * @property {Buffer[]} queue
+ */
+
+/**
+ * @typedef {object} DogStatsDClientOptions
+ * @property {string} host
+ * @property {number} port
+ * @property {string[]} tags
+ * @property {typeof import('node:dns').lookup} lookup
+ * @property {URL|string} [metricsProxyUrl]
+ */
+
+/**
+ * @typedef {object} DogStatsDTelemetryState
+ * @property {number[]} aggregatedContextsByType
+ * @property {number} bytesDropped
+ * @property {number} bytesSent
+ * @property {number[]} metricsByType
+ * @property {number} nextFlush
+ * @property {number} packetsDropped
+ * @property {number} packetsSent
+ * @property {DogStatsDBufferState} payload
+ */
+
+/**
+ * @typedef {object} MetricNode
+ * @property {Map<string, MetricNode>} nodes
+ * @property {number} submissions
+ * @property {number|Histogram|null} value
+ */
+
+/**
+ * @callback CaptureMetric
+ * @param {MetricNode} node
+ * @param {string} name
+ * @param {string[]} tags
+ * @returns {void}
+ */
+
+const identityRefreshChannel = channel('datadog:identity:refresh')
+
 /**
  * @import { DogStatsD } from "../../../index.d.ts"
  * @implements {DogStatsD}
  */
 class DogStatsDClient {
+  #family
+  #host
+  #httpOptions
   #lookup
+  #metrics = { message: '', offset: 0, queue: [] }
+  #port
+  #serverlessDeliveryTracker
   #tagsPrefix
-  constructor (options) {
+  #telemetryHttpTagsPrefix
+  #telemetryUdpTagsPrefix
+  #udp4
+  #udp6
+
+  /** @type {DogStatsDTelemetryState|undefined} */
+  telemetry
+
+  /**
+   * @param {DogStatsDClientOptions} options - DogStatsD transport options
+   * @param {boolean} [telemetryEnabled] - Whether to collect client telemetry
+   */
+  constructor (options, telemetryEnabled = false) {
     this.#lookup = options.lookup
     if (options.metricsProxyUrl) {
-      this._httpOptions = {
+      this.#httpOptions = {
         method: 'POST',
+        retry: false,
         url: options.metricsProxyUrl.toString(),
         path: '/dogstatsd/v2/proxy',
       }
     }
 
-    this._host = options.host
-    this._family = isIP(this._host)
-    this._port = options.port
-    this._tags = options.tags
-    this.#tagsPrefix = this._tags?.length ? `|#${this._tags.join(',')}` : ''
-    this._queue = []
-    this._buffer = ''
-    this._offset = 0
-    this._udp4 = this._socket('udp4')
-    this._udp6 = this._socket('udp6')
+    this.#family = isIP(options.host)
+    this.#host = options.host
+    this.#port = options.port
+    this.#serverlessDeliveryTracker = createServerlessDeliveryTracker()
+
+    if (telemetryEnabled) {
+      this.telemetry = {
+        aggregatedContextsByType: [0, 0, 0, 0],
+        bytesDropped: 0,
+        bytesSent: 0,
+        metricsByType: [0, 0, 0, 0],
+        nextFlush: performance.now() + TELEMETRY_INTERVAL,
+        packetsDropped: 0,
+        packetsSent: 0,
+        payload: { message: '', offset: 0, queue: [] },
+      }
+    }
+
+    this.#updateTagPrefixes(options.tags.length ? `|#${options.tags.join(',')}` : '')
+
+    this.#udp4 = this._socket('udp4')
+    this.#udp6 = this._socket('udp6')
+  }
+
+  /**
+   * Recomputes the cached tags and tag-prefix (mirrors the constructor) after a `config.tags`
+   * change, e.g. a MicroVM clone resume.
+   *
+   * Buffered lines have the old prefix baked in, and on a clone resume they were produced during
+   * the image build, so every clone holds the same bytes — flushing them would submit one identical
+   * copy per clone. Dropping is right here for that reason only: for a tag change on a live process
+   * the buffer holds unique data whose old tags are still correct, so that case wants a flush
+   * before the swap.
+   *
+   * @param {string[]} tags - DogStatsD-formatted tags (e.g. `['key:value']`)
+   */
+  updateTags (tags) {
+    const tagsPrefix = tags.length ? `|#${tags.join(',')}` : ''
+    if (tagsPrefix === this.#tagsPrefix) return
+
+    this.#updateTagPrefixes(tagsPrefix)
+    this.#metrics.queue = []
+    this.#metrics.message = ''
+    this.#metrics.offset = 0
+
+    const payload = this.telemetry?.payload
+    if (payload) {
+      payload.queue = []
+      payload.message = ''
+      payload.offset = 0
+    }
+  }
+
+  /**
+   * Updates the prefixes used for application metrics and client telemetry.
+   * @param {string} tagsPrefix - Serialized global DogStatsD tags
+   */
+  #updateTagPrefixes (tagsPrefix) {
+    this.#tagsPrefix = tagsPrefix
+
+    if (!this.telemetry) return
+
+    const separator = tagsPrefix ? ',' : '|#'
+    const prefix = `${tagsPrefix}${separator}client:nodejs,client_version:${tracerVersion},client_transport:`
+    this.#telemetryHttpTagsPrefix = `${prefix}http`
+    this.#telemetryUdpTagsPrefix = `${prefix}udp`
   }
 
   increment (stat, value, tags) {
@@ -67,100 +204,325 @@ class DogStatsDClient {
     this._add(stat, value, TYPE_HISTOGRAM, tags)
   }
 
-  flush () {
-    const queue = this._enqueue()
+  /**
+   * @param {boolean|(() => void)} [forceTelemetry] - Whether to ignore the telemetry interval, or completion callback
+   * @param {() => void} [done] - Called after serverless deliveries complete
+   * @returns {void}
+   */
+  flush (forceTelemetry = false, done) {
+    if (typeof forceTelemetry === 'function') {
+      done = forceTelemetry
+      forceTelemetry = false
+    }
+
+    let complete
+    let track
+    if (done && !this.#serverlessDeliveryTracker) {
+      let pending = 1
+      complete = () => {
+        if (--pending === 0) done()
+      }
+      track = send => {
+        pending++
+        send(complete)
+      }
+    }
+
+    this.#flush(this.#metrics, true, track)
+    if (this.telemetry) this.#flushTelemetry(forceTelemetry, track)
+
+    if (this.#serverlessDeliveryTracker) return this.#serverlessDeliveryTracker.waitForIdle(done)
+    complete?.()
+  }
+
+  /**
+   * @param {boolean} force - Whether to ignore the telemetry interval
+   * @param {((send: (done?: () => void) => void) => void)} [track] - Tracks a non-serverless delivery
+   * @returns {void}
+   */
+  #flushTelemetry (force, track) {
+    const telemetry = this.telemetry
+
+    const now = performance.now()
+
+    if (!force && now < telemetry.nextFlush) return
+
+    telemetry.nextFlush = now + TELEMETRY_INTERVAL
+
+    let aggregatedContexts = 0
+    let metrics = 0
+    for (let index = 0; index < TYPE_LABELS.length; index++) {
+      aggregatedContexts += telemetry.aggregatedContextsByType[index]
+      metrics += telemetry.metricsByType[index]
+    }
+
+    this.#addTelemetry('datadog.dogstatsd.client.metrics', metrics)
+    this.#addTelemetry('datadog.dogstatsd.client.aggregated_context', aggregatedContexts)
+    for (let index = 0; index < TYPE_LABELS.length; index++) {
+      this.#addTelemetry(
+        'datadog.dogstatsd.client.metrics_by_type',
+        telemetry.metricsByType[index],
+        TYPE_LABELS[index]
+      )
+      this.#addTelemetry(
+        'datadog.dogstatsd.client.aggregated_context_by_type',
+        telemetry.aggregatedContextsByType[index],
+        TYPE_LABELS[index]
+      )
+    }
+    this.#addTelemetry('datadog.dogstatsd.client.bytes_sent', telemetry.bytesSent)
+    this.#addTelemetry('datadog.dogstatsd.client.bytes_dropped', telemetry.bytesDropped)
+    this.#addTelemetry('datadog.dogstatsd.client.packets_sent', telemetry.packetsSent)
+    this.#addTelemetry('datadog.dogstatsd.client.packets_dropped', telemetry.packetsDropped)
+
+    telemetry.bytesSent = 0
+    telemetry.bytesDropped = 0
+    telemetry.packetsSent = 0
+    telemetry.packetsDropped = 0
+    for (let index = 0; index < TYPE_LABELS.length; index++) {
+      telemetry.aggregatedContextsByType[index] = 0
+      telemetry.metricsByType[index] = 0
+    }
+
+    this.#flush(telemetry.payload, false, track)
+  }
+
+  /**
+   * @param {DogStatsDBufferState} state - Payload state to flush
+   * @param {boolean} recordTelemetry - Whether to record the transport outcome
+   * @param {((send: (done?: () => void) => void) => void)} [track] - Tracks a non-serverless delivery
+   * @returns {void}
+   */
+  #flush (state, recordTelemetry, track) {
+    const queue = this._enqueue(state)
 
     if (queue.length === 0) return
 
-    log.debug('Flushing %s metrics via %s', queue.length, this._httpOptions ? 'HTTP' : 'UDP')
+    log.debug('Flushing %s metrics via %s', queue.length, this.#httpOptions ? 'HTTP' : 'UDP')
 
-    this._queue = []
+    state.queue = []
 
-    if (this._httpOptions) {
-      this._sendHttp(queue)
-    } else {
-      this._sendUdp(queue)
+    const send = complete => {
+      if (this.#httpOptions) this._sendHttp(queue, recordTelemetry, complete)
+      else this._sendUdp(queue, recordTelemetry, complete)
     }
+
+    if (this.#serverlessDeliveryTracker) return this.#serverlessDeliveryTracker.track(send)
+    if (track) return track(send)
+    send()
   }
 
-  _sendHttp (queue) {
+  /**
+   * @param {number} bytes - Number of bytes sent
+   * @returns {void}
+   */
+  #recordSent (bytes) {
+    const telemetry = this.telemetry
+
+    telemetry.bytesSent += bytes
+    telemetry.packetsSent++
+  }
+
+  /**
+   * @param {number} bytes - Number of bytes dropped
+   * @param {number} [packets] - Number of packets dropped
+   * @returns {void}
+   */
+  #recordDropped (bytes, packets = 1) {
+    const telemetry = this.telemetry
+
+    telemetry.bytesDropped += bytes
+    telemetry.packetsDropped += packets
+  }
+
+  /**
+   * Send metrics to the agent via HTTP
+   *
+   * @param {Buffer[]} queue - The metrics to send
+   * @param {boolean} recordTelemetry - Whether to record the transport outcome
+   * @param {() => void} [done] - Called after delivery completes
+   * @returns {void}
+   * @memberof DogStatsDClient
+   */
+  _sendHttp (queue, recordTelemetry, done) {
     const buffer = Buffer.concat(queue)
-    request(buffer, this._httpOptions, (err) => {
-      if (err) {
-        log.error('DogStatsDClient: HTTP error from agent: %s', err.message, err)
-        if (err.status === 404) {
+    request(buffer, this.#httpOptions, (error, _result, _statusCode, _headers, dropped) => {
+      if (dropped) {
+        if (recordTelemetry && this.telemetry) {
+          this.#recordDropped(buffer.length)
+        }
+        done?.()
+      } else if (error) {
+        log.error('DogStatsDClient: HTTP error from agent: %s', error.message, error)
+        if (error.status === 404) {
           // Inside this if-block, we have connectivity to the agent, but
           // we're not getting a 200 from the proxy endpoint. If it's a 404,
           // then we know we'll never have the endpoint, so just clear out the
           // options. Either way, we can give UDP a try.
-          this._httpOptions = undefined
+          this.#httpOptions = undefined
         }
-        this._sendUdp(queue)
+        this._sendUdp(queue, recordTelemetry, done)
+      } else {
+        if (recordTelemetry && this.telemetry) this.#recordSent(buffer.length)
+        done?.()
       }
     })
   }
 
-  _sendUdp (queue) {
+  /**
+   * Send metrics to the agent via UDP
+   *
+   * @param {Buffer[]} queue - The metrics to send
+   * @param {boolean} recordTelemetry - Whether to record the transport outcome
+   * @param {() => void} [done] - Called after delivery completes
+   * @returns {void}
+   * @memberof DogStatsDClient
+   */
+  _sendUdp (queue, recordTelemetry, 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 (this.#family === 0) {
+        this.#lookup(this.#host, (error, address, family) => {
+          if (error) {
+            if (recordTelemetry && this.telemetry) {
+              let bytes = 0
+              for (const buffer of queue) {
+                bytes += buffer.length
+              }
+              this.#recordDropped(bytes, queue.length)
+            }
+            log.error('DogStatsDClient: Host not found', error)
+            return done?.()
+          }
+          this._sendUdpFromQueue(queue, address, family, recordTelemetry, done)
         })
       } else {
-        this._sendUdpFromQueue(queue, this._host, this._family)
+        this._sendUdpFromQueue(queue, this.#host, this.#family, recordTelemetry, done)
       }
     })
   }
 
-  _sendUdpFromQueue (queue, address, family) {
-    const socket = family === 6 ? this._udp6 : this._udp4
+  /**
+   * Send metrics to the agent via UDP from queue
+   *
+   * @param {Buffer[]} queue - The metrics to send
+   * @param {string} address - The address to send the metrics to
+   * @param {number} family - The family of the address
+   * @param {boolean} recordTelemetry - Whether to record the transport outcome
+   * @param {() => void} [done] - Called after every packet completes
+   * @returns {void}
+   * @memberof DogStatsDClient
+   */
+  _sendUdpFromQueue (queue, address, family, recordTelemetry, 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)
+
+      if (!this.telemetry && !done) {
+        socket.send(buffer, 0, buffer.length, this.#port, address)
+        continue
+      }
+
+      try {
+        socket.send(buffer, 0, buffer.length, this.#port, address, (error) => {
+          if (error) {
+            if (recordTelemetry && this.telemetry) {
+              this.#recordDropped(buffer.length)
+            }
+            log.error('DogStatsDClient: UDP error', error)
+          } else if (recordTelemetry && this.telemetry) {
+            this.#recordSent(buffer.length)
+          }
+          complete()
+        })
+      } catch (error) {
+        log.error('DogStatsDClient: UDP error sending metrics', error)
+        complete()
+      }
     }
   }
 
-  _add (stat, value, type, tags) {
+  /**
+   * Add a metric to the queue
+   *
+   * @param {string} stat - The metric name
+   * @param {number} value - The metric value
+   * @param {string} type - The metric type
+   * @param {string[]} tags - The metric tags
+   * @param {string} [tagsPrefix] - Serialized global tags
+   * @param {DogStatsDBufferState} [state] - Payload state receiving the metric
+   * @returns {void}
+   * @memberof DogStatsDClient
+   */
+  _add (stat, value, type, tags, tagsPrefix = this.#tagsPrefix, state = this.#metrics) {
     let message = `${stat}:${value}|${type}`
 
     if (tags?.length) {
-      message += this.#tagsPrefix
-        ? `${this.#tagsPrefix},${tags.join(',')}`
+      message += tagsPrefix
+        ? `${tagsPrefix},${tags.join(',')}`
         : `|#${tags.join(',')}`
     } else {
-      message += this.#tagsPrefix
+      message += tagsPrefix
     }
 
     if (entityId) {
       message += `|c:${entityId}`
     }
 
-    this._write(`${message}\n`)
+    this._write(`${message}\n`, state)
   }
 
-  _write (message) {
-    const offset = Buffer.byteLength(message)
+  /**
+   * @param {string} stat - Telemetry metric name
+   * @param {number} value - Telemetry metric value
+   * @param {string} [typeLabel] - Submitted or aggregated metric type
+   * @returns {void}
+   */
+  #addTelemetry (stat, value, typeLabel) {
+    const tags = typeLabel === undefined ? undefined : [`metrics_type:${typeLabel}`]
+    const tagsPrefix = this.#httpOptions ? this.#telemetryHttpTagsPrefix : this.#telemetryUdpTagsPrefix
 
-    if (this._offset + offset > MAX_BUFFER_SIZE) {
-      this._enqueue()
-    }
+    this._add(stat, value, TYPE_COUNTER, tags, tagsPrefix, this.telemetry.payload)
+  }
 
-    this._offset += offset
-    this._buffer += message
+  /**
+   * Write a message to the queue
+   *
+   * @param {string} message - The message to write
+   * @param {DogStatsDBufferState} [state] - Payload state receiving the message
+   * @returns {void}
+   * @memberof DogStatsDClient
+   */
+  _write (message, state = this.#metrics) {
+    const offset = Buffer.byteLength(message)
+    if (state.offset + offset > MAX_BUFFER_SIZE) {
+      this._enqueue(state)
+    }
+    state.offset += offset
+    state.message += message
   }
 
-  _enqueue () {
-    if (this._offset > 0) {
-      this._queue.push(Buffer.from(this._buffer))
-      this._buffer = ''
-      this._offset = 0
+  /**
+   * Enqueue a message to the queue
+   *
+   * @param {DogStatsDBufferState} [state] - Payload state to enqueue
+   * @returns {Buffer[]} Queued payloads
+   * @memberof DogStatsDClient
+   */
+  _enqueue (state = this.#metrics) {
+    if (state.offset > 0) {
+      state.queue.push(Buffer.from(state.message))
+      state.message = ''
+      state.offset = 0
     }
 
-    return this._queue
+    return state.queue
   }
 
   _socket (type) {
@@ -205,19 +567,65 @@ class DogStatsDClient {
   }
 }
 
+/**
+ * @param {DogStatsDClientOptions} options - DogStatsD transport options
+ * @returns {MetricsAggregationClient} Aggregating client with shared telemetry
+ */
+function createMetricsAggregationClient (options) {
+  const client = new DogStatsDClient(options, true)
+
+  return new MetricsAggregationClient(client)
+}
+
 class MetricsAggregationClient {
+  #metricsByType
+
+  /** @type {DogStatsDTelemetryState|undefined} */
+  #telemetry
+
+  /**
+   * @param {DogStatsDClient} client - DogStatsD transport client
+   */
   constructor (client) {
     this._client = client
+    this.#telemetry = client.telemetry
+    this.#metricsByType = this.#telemetry?.metricsByType ?? [0, 0, 0, 0]
 
     this.reset()
   }
 
-  flush () {
-    this._captureCounters()
-    this._captureGauges()
-    this._captureHistograms()
+  /**
+   * Recomputes the wrapped client's cached tags (e.g. after a MicroVM clone resume).
+   * @param {string[]} tags - DogStatsD-formatted tags (e.g. `['key:value']`)
+   */
+  updateTags (tags) {
+    this._client.updateTags(tags)
+  }
+
+  /**
+   * @param {boolean|(() => void)} [forceTelemetry] - Whether to ignore the telemetry interval, or completion callback
+   * @param {() => void} [done] - Called after serverless deliveries complete
+   * @returns {void}
+   */
+  flush (forceTelemetry = false, done) {
+    if (typeof forceTelemetry === 'function') {
+      done = forceTelemetry
+      forceTelemetry = false
+    }
+
+    const counters = this._captureCounters()
+    const gauges = this._captureGauges()
+    const histograms = this._captureHistograms()
+    const telemetry = this.#telemetry
 
-    this._client.flush()
+    if (telemetry) {
+      telemetry.aggregatedContextsByType[TYPE_COUNTER_INDEX] += counters
+      telemetry.aggregatedContextsByType[TYPE_GAUGE_INDEX] += gauges
+      telemetry.aggregatedContextsByType[TYPE_HISTOGRAM_INDEX] += histograms
+    }
+
+    if (forceTelemetry) this._client.flush(true, done)
+    else this._client.flush(done)
   }
 
   reset () {
@@ -229,6 +637,7 @@ class MetricsAggregationClient {
   // TODO: Aggregate with a histogram and send the buckets to the client.
   distribution (name, value, tags) {
     this._client.distribution(name, value, tags)
+    this.#metricsByType[TYPE_DISTRIBUTION_INDEX]++
   }
 
   boolean (name, value, tags) {
@@ -271,24 +680,45 @@ class MetricsAggregationClient {
     this.count(name, -count, tags)
   }
 
+  /**
+   * @returns {number} Number of gauge contexts flushed
+   */
   _captureGauges () {
-    this._captureTree(this._gauges, (node, name, tags) => {
+    let metrics = 0
+    const contexts = this._captureTree(this._gauges, (node, name, tags) => {
+      metrics += node.submissions
       this._client.gauge(name, node.value, tags)
     })
 
+    this.#metricsByType[TYPE_GAUGE_INDEX] += metrics
     this._gauges.clear()
+
+    return contexts
   }
 
+  /**
+   * @returns {number} Number of counter contexts flushed
+   */
   _captureCounters () {
-    this._captureTree(this._counters, (node, name, tags) => {
+    let metrics = 0
+    const contexts = this._captureTree(this._counters, (node, name, tags) => {
+      metrics += node.submissions
       this._client.increment(name, node.value, tags)
     })
 
+    this.#metricsByType[TYPE_COUNTER_INDEX] += metrics
     this._counters.clear()
+
+    return contexts
   }
 
+  /**
+   * @returns {number} Number of histogram contexts flushed
+   */
   _captureHistograms () {
-    this._captureTree(this._histograms, (node, name, tags) => {
+    let metrics = 0
+    const contexts = this._captureTree(this._histograms, (node, name, tags) => {
+      metrics += node.submissions
       const stats = node.value
 
       this._client.gauge(`${name}.min`, stats.min, tags)
@@ -301,25 +731,49 @@ class MetricsAggregationClient {
       this._client.gauge(`${name}.95percentile`, stats.p95, tags)
     })
 
+    this.#metricsByType[TYPE_HISTOGRAM_INDEX] += metrics
     this._histograms.clear()
+
+    return contexts
   }
 
+  /**
+   * @param {Map<string, MetricNode>} tree - Metric context tree
+   * @param {CaptureMetric} fn - Called for every context with submissions
+   * @returns {number} Number of contexts with submissions
+   */
   _captureTree (tree, fn) {
+    let contexts = 0
+
     for (const [name, root] of tree) {
-      this._captureNode(root, name, [], fn)
+      contexts += this._captureNode(root, name, [], fn)
     }
+
+    return contexts
   }
 
+  /**
+   * @param {MetricNode} node - Current metric context node
+   * @param {string} name - Metric name
+   * @param {string[]} tags - Current metric tags
+   * @param {CaptureMetric} fn - Called for every context with submissions
+   * @returns {number} Number of contexts with submissions
+   */
   _captureNode (node, name, tags, fn) {
-    if (node.touched) {
+    let contexts = 0
+
+    if (node.submissions !== 0) {
       fn(node, name, tags)
+      contexts++
     }
 
     for (const [tag, next] of node.nodes) {
       tags.push(tag)
-      this._captureNode(next, name, tags, fn)
+      contexts += this._captureNode(next, name, tags, fn)
       tags.pop()
     }
+
+    return contexts
   }
 
   _ensureTree (tree, name, tags = [], value) {
@@ -333,7 +787,7 @@ class MetricsAggregationClient {
       node = this._ensureNode(node.nodes, tag, value)
     }
 
-    node.touched = true
+    node.submissions++
 
     return node
   }
@@ -342,7 +796,7 @@ class MetricsAggregationClient {
     let node = container.get(key)
 
     if (!node) {
-      node = { nodes: new Map(), touched: false, value }
+      node = { nodes: new Map(), submissions: 0, value }
 
       if (typeof key === 'string') {
         container.set(key, node)
@@ -362,14 +816,20 @@ class CustomMetrics {
   #client
   constructor (config) {
     const clientConfig = DogStatsDClient.generateClientConfig(config)
-    this.#client = new MetricsAggregationClient(new DogStatsDClient(clientConfig))
+    this.#client = createMetricsAggregationClient(clientConfig)
+
+    // CustomMetrics has process-lifetime flush handlers and no stop hook, so this shares that lifetime.
+    identityRefreshChannel.subscribe(() => {
+      this.#client.updateTags(DogStatsDClient.generateClientConfig(config).tags)
+    })
 
     const flush = this.flush.bind(this)
 
     // TODO(bengl) this magic number should be configurable
     setInterval(flush, 10 * 1000).unref?.()
 
-    globalThis[Symbol.for('dd-trace')].beforeExitHandlers.add(flush)
+    globalThis[Symbol.for('dd-trace')].beforeExitHandlers.add(() => this.#client.flush(true))
+    registerTelemetryFlusher(done => this.flush(done))
   }
 
   increment (stat, value = 1, tags) {
@@ -392,8 +852,8 @@ class CustomMetrics {
     this.#client.histogram(stat, value, CustomMetrics.tagTranslator(tags))
   }
 
-  flush () {
-    return this.#client.flush()
+  flush (done) {
+    return this.#client.flush(false, done)
   }
 
   /**
@@ -419,4 +879,5 @@ module.exports = {
   DogStatsDClient,
   CustomMetrics,
   MetricsAggregationClient,
+  createMetricsAggregationClient,
 }