fix(id): reset DogStatsD and exporter state on refresh

packages/dd-trace/src/dogstatsd.js

@@ -151,10 +151,11 @@ class DogStatsDClient {
    * before the swap.
    *
    * @param {string[]} tags - DogStatsD-formatted tags (e.g. `['key:value']`)
+   * @returns {boolean} True if the tag prefix actually changed (and buffered lines were dropped)
    */
   updateTags (tags) {
     const tagsPrefix = tags.length ? `|#${tags.join(',')}` : ''
-    if (tagsPrefix === this.#tagsPrefix) return
+    if (tagsPrefix === this.#tagsPrefix) return false
 
     this.#updateTagPrefixes(tagsPrefix)
     this.#metrics.queue = []
@@ -167,6 +168,8 @@ class DogStatsDClient {
       payload.message = ''
       payload.offset = 0
     }
+
+    return true
   }
 
   /**
@@ -595,11 +598,16 @@ class MetricsAggregationClient {
   }
 
   /**
-   * Recomputes the wrapped client's cached tags (e.g. after a MicroVM clone resume).
+   * Recomputes the wrapped client's cached tags (e.g. after a MicroVM clone resume). Pending
+   * counters/gauges/histograms were aggregated under the old identity, so they're reset along
+   * with the client's buffered lines — but only if the tags actually changed, so a no-op resume
+   * doesn't discard in-flight aggregation for nothing.
    * @param {string[]} tags - DogStatsD-formatted tags (e.g. `['key:value']`)
    */
   updateTags (tags) {
-    this._client.updateTags(tags)
+    if (this._client.updateTags(tags)) {
+      this.reset()
+    }
   }
 
   /**

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

@@ -3,11 +3,19 @@
 const { URL } = require('node:url')
 const os = require('node:os')
 
+const { channel } = require('dc-polyfill')
+
 const log = require('../../log')
 const { containerId } = require('../common/docker')
 const Writer = require('./writer')
 const { computeIntakeUrl } = require('./intake')
 
+const identityRefreshChannel = channel('datadog:identity:refresh')
+
+// Only one AgentlessExporter is ever live in a real process, so replacing the subscription on
+// construction is safe - it just keeps tests (which build several) from piling up listeners.
+let unsubscribeBatchReset = null
+
 /**
  * Agentless exporter for APM trace intake.
  * Sends traces directly to the Datadog intake without requiring a local agent.
@@ -52,6 +60,12 @@ class AgentlessExporter {
       metadata,
     })
 
+    // A clone resume shouldn't flush spans buffered before the snapshot under its own identity.
+    unsubscribeBatchReset?.()
+    const onIdentityRefresh = () => this._writer.resetPendingBatch()
+    identityRefreshChannel.subscribe(onIdentityRefresh)
+    unsubscribeBatchReset = () => identityRefreshChannel.unsubscribe(onIdentityRefresh)
+
     const ddTrace = globalThis[Symbol.for('dd-trace')]
     if (ddTrace?.beforeExitHandlers) {
       ddTrace.beforeExitHandlers.add(this.flush.bind(this))

packages/dd-trace/src/exporters/common/writer.js

@@ -99,6 +99,15 @@ class Writer {
   setUrl (url) {
     this._url = url
   }
+
+  /**
+   * Discards whatever's queued in the encoder. Used on a MicroVM clone resume, where anything
+   * buffered before the snapshot would otherwise flush under every clone's identity.
+   * @returns {void}
+   */
+  resetPendingBatch () {
+    this._encoder.reset()
+  }
 }
 
 module.exports = Writer

packages/dd-trace/src/span_stats.js

@@ -1,8 +1,10 @@
 'use strict'
 
 const os = require('node:os')
-const pkg = require('../../../package.json')
 
+const { channel } = require('dc-polyfill')
+
+const pkg = require('../../../package.json')
 const { LogCollapsingLowestDenseDDSketch } = require('../../../vendor/dist/@datadog/sketches-js')
 const {
   MEASURED,
@@ -29,6 +31,12 @@ const {
   DEFAULT_SERVICE_NAME,
 } = require('./encode/tags-processors')
 
+const identityRefreshChannel = channel('datadog:identity:refresh')
+
+// Only one SpanStatsProcessor is ever live in a real process, so replacing the subscription on
+// construction is safe - it just keeps tests (which build several) from piling up listeners.
+let unsubscribeBucketReset = null
+
 class SpanAggStats {
   constructor (aggKey) {
     this.aggKey = aggKey
@@ -229,6 +237,14 @@ class SpanStatsProcessor {
       this.timer = setInterval(this.onInterval.bind(this), intervalMs)
       this.timer.unref?.()
     }
+
+    // A clone resume shouldn't export buckets accumulated before the snapshot under its own identity.
+    unsubscribeBucketReset?.()
+    const onIdentityRefresh = () => {
+      this.buckets = new TimeBuckets()
+    }
+    identityRefreshChannel.subscribe(onIdentityRefresh)
+    unsubscribeBucketReset = () => identityRefreshChannel.unsubscribe(onIdentityRefresh)
   }
 
   onInterval () {

packages/dd-trace/test/dogstatsd.spec.js

@@ -1359,6 +1359,53 @@ describe('dogstatsd', () => {
 
       assert.strictEqual(udp4.send.firstCall.args[0].toString(), 'test.buffered:1|d|#runtime-id:initial-id\n')
     })
+
+    it('should drop pending aggregated counters/gauges/histograms when an identity refresh changes its tags',
+      () => {
+        const config = {
+          dogstatsd: {
+            hostname: '127.0.0.1',
+            port: 8125,
+          },
+          lookup: dns.lookup,
+          runtimeMetricsRuntimeId: true,
+          tags: { 'runtime-id': 'initial-id' },
+        }
+
+        client = new CustomMetrics(config)
+        client.increment('test.count', 10)
+        client.gauge('test.avg', 5)
+        client.histogram('test.hist', 1)
+
+        config.tags['runtime-id'] = 'refreshed-id'
+        identityRefreshChannel.publish(config)
+
+        client.flush()
+
+        sinon.assert.notCalled(udp4.send)
+      })
+
+    it('should preserve pending aggregated counters/gauges/histograms when an identity refresh does not change ' +
+      'its tags', () => {
+      const config = {
+        dogstatsd: {
+          hostname: '127.0.0.1',
+          port: 8125,
+        },
+        lookup: dns.lookup,
+        runtimeMetricsRuntimeId: true,
+        tags: { 'runtime-id': 'initial-id' },
+      }
+
+      client = new CustomMetrics(config)
+      client.increment('test.count', 10)
+
+      identityRefreshChannel.publish(config)
+      client.flush()
+
+      sinon.assert.called(udp4.send)
+      assert.strictEqual(udp4.send.firstCall.args[0].toString(), 'test.count:10|c|#runtime-id:initial-id\n')
+    })
   })
 
   describe('MetricsAggregationClient', () => {

packages/dd-trace/test/exporters/agentless/exporter.spec.js

@@ -7,11 +7,14 @@ const { inspect } = require('node:util')
 const { describe, it, beforeEach, afterEach } = require('mocha')
 const sinon = require('sinon')
 const proxyquire = require('proxyquire')
+const { channel } = require('dc-polyfill')
 
 const { assertObjectContains } = require('../../../../../integration-tests/helpers')
 
 require('../../setup/core')
 
+const identityRefreshChannel = channel('datadog:identity:refresh')
+
 describe('AgentlessExporter', () => {
   let Exporter
   let exporter
@@ -26,6 +29,7 @@ describe('AgentlessExporter', () => {
       append: sinon.stub(),
       flush: sinon.stub().callsFake((cb) => cb && cb()),
       setUrl: sinon.stub(),
+      resetPendingBatch: sinon.stub(),
     }
 
     const Writer = function () {
@@ -282,6 +286,27 @@ describe('AgentlessExporter', () => {
     })
   })
 
+  describe('identity refresh', () => {
+    it('drops the pending trace batch when the identity-refresh channel fires', () => {
+      exporter = new Exporter({})
+
+      identityRefreshChannel.publish({ tags: {} })
+
+      sinon.assert.calledOnce(writer.resetPendingBatch)
+    })
+
+    it('stops reacting once a newer exporter takes over', () => {
+      exporter = new Exporter({})
+      new Exporter({}) // eslint-disable-line no-new
+      writer.resetPendingBatch.resetHistory()
+
+      identityRefreshChannel.publish({ tags: {} })
+
+      // Only one reset, not two - the first exporter's subscription was replaced, not stacked on.
+      sinon.assert.calledOnce(writer.resetPendingBatch)
+    })
+  })
+
   describe('setUrl', () => {
     let log
 

packages/dd-trace/test/exporters/common/writer.spec.js

@@ -179,4 +179,10 @@ describe('common Writer', () => {
     assert.strictEqual(writer.append(payload), false)
     sinon.assert.calledOnceWithExactly(encoder.encode, payload)
   })
+
+  it('resetPendingBatch discards the pending encoded batch', () => {
+    writer.resetPendingBatch()
+
+    sinon.assert.calledOnce(encoder.reset)
+  })
 })

packages/dd-trace/test/span_stats.spec.js

@@ -10,6 +10,8 @@ const proxyquire = require('proxyquire')
 
 require('./setup/core')
 
+const identityRefreshChannel = channel('datadog:identity:refresh')
+
 const { LogCollapsingLowestDenseDDSketch } = require('../../../vendor/dist/@datadog/sketches-js')
 const { version } = require('../src/pkg')
 const pkg = require('../../../package.json')
@@ -776,4 +778,35 @@ describe('SpanStatsProcessor', () => {
     p.onSpanFinished(topLevelSpan)
     assert.strictEqual(p.buckets.size, 1)
   })
+
+  it('should clear pending buckets when the identity-refresh channel fires', () => {
+    const p = new SpanStatsProcessor(config)
+    clearTimeout(p.timer)
+
+    p.onSpanFinished(topLevelSpan)
+    assert.strictEqual(p.buckets.size, 1)
+
+    const previousBuckets = p.buckets
+    identityRefreshChannel.publish(config)
+
+    assert.notStrictEqual(p.buckets, previousBuckets)
+    assert.strictEqual(p.buckets.size, 0)
+  })
+
+  it('should stop reacting to identity refresh once a newer instance takes over', () => {
+    const first = new SpanStatsProcessor(config)
+    clearTimeout(first.timer)
+    const firstBuckets = first.buckets
+
+    const second = new SpanStatsProcessor(config)
+    clearTimeout(second.timer)
+    const secondBuckets = second.buckets
+
+    identityRefreshChannel.publish(config)
+
+    // Only the second (newest) instance should react - the first's subscription was replaced,
+    // not stacked on top of.
+    assert.strictEqual(first.buckets, firstBuckets)
+    assert.notStrictEqual(second.buckets, secondBuckets)
+  })
 })