From 8c89cf4aa03c7f7a6d03c1c4c29ce954ff8a7cc3 Mon Sep 17 00:00:00 2001
From: Bryan English <bryan.english@datadoghq.com>
Date: Mon, 30 Dec 2024 22:26:30 -0500
Subject: [PATCH] refactor custom metrics to be self-contained

---
 packages/dd-trace/src/dogstatsd.js       |  8 ++++++++
 packages/dd-trace/src/proxy.js           |  9 ---------
 packages/dd-trace/test/dogstatsd.spec.js | 15 +++++++++++++++
 packages/dd-trace/test/proxy.spec.js     | 22 ----------------------
 4 files changed, 23 insertions(+), 31 deletions(-)

diff --git a/packages/dd-trace/src/dogstatsd.js b/packages/dd-trace/src/dogstatsd.js
index a396c9e98a4..61a2fe090a3 100644
--- a/packages/dd-trace/src/dogstatsd.js
+++ b/packages/dd-trace/src/dogstatsd.js
@@ -190,6 +190,14 @@ class CustomMetrics {
   constructor (config) {
     const clientConfig = DogStatsDClient.generateClientConfig(config)
     this.dogstatsd = new DogStatsDClient(clientConfig)
+
+    this._boundFlush = this.flush.bind(this)
+
+    // TODO(bengl) this magic number should be configurable
+    this._flushInterval = setInterval(this._boundFlush, 10 * 1000)
+    this._flushInterval.unref()
+
+    process.once('beforeExit', this._boundFlush)
   }
 
   increment (stat, value = 1, tags) {
diff --git a/packages/dd-trace/src/proxy.js b/packages/dd-trace/src/proxy.js
index fd814c9d6e3..2a3bf831b9d 100644
--- a/packages/dd-trace/src/proxy.js
+++ b/packages/dd-trace/src/proxy.js
@@ -67,16 +67,7 @@ class Tracer extends NoopProxy {
       telemetry.start(config, this._pluginManager)
 
       if (config.dogstatsd) {
-        // Custom Metrics
         this.dogstatsd = new dogstatsd.CustomMetrics(config)
-
-        setInterval(() => {
-          this.dogstatsd.flush()
-        }, 10 * 1000).unref()
-
-        process.once('beforeExit', () => {
-          this.dogstatsd.flush()
-        })
       }
 
       if (config.spanLeakDebug > 0) {
diff --git a/packages/dd-trace/test/dogstatsd.spec.js b/packages/dd-trace/test/dogstatsd.spec.js
index 18a39d6d304..1e961d87e7d 100644
--- a/packages/dd-trace/test/dogstatsd.spec.js
+++ b/packages/dd-trace/test/dogstatsd.spec.js
@@ -431,5 +431,20 @@ describe('dogstatsd', () => {
       expect(udp4.send).to.have.been.called
       expect(udp4.send.firstCall.args[0].toString()).to.equal('test.histogram:10|h\n')
     })
+
+    it('should flush via interval', () => {
+      const clock = sinon.useFakeTimers()
+
+      client = new CustomMetrics({ dogstatsd: {} })
+
+      client.gauge('test.avg', 10, { foo: 'bar' })
+
+      expect(udp4.send).not.to.have.been.called
+
+      clock.tick(10 * 1000)
+
+      expect(udp4.send).to.have.been.called
+      expect(udp4.send.firstCall.args[0].toString()).to.equal('test.avg:10|g|#foo:bar\n')
+    })
   })
 })
diff --git a/packages/dd-trace/test/proxy.spec.js b/packages/dd-trace/test/proxy.spec.js
index dd145390245..24a67c63671 100644
--- a/packages/dd-trace/test/proxy.spec.js
+++ b/packages/dd-trace/test/proxy.spec.js
@@ -401,28 +401,6 @@ describe('TracerProxy', () => {
         proxy.dogstatsd.increment('foo')
       })
 
-      it('should call custom metrics flush via interval', () => {
-        const clock = sinon.useFakeTimers()
-
-        config.dogstatsd = {
-          hostname: 'localhost',
-          port: 9876
-        }
-        config.tags = {
-          service: 'photos',
-          env: 'prod',
-          version: '1.2.3'
-        }
-
-        proxy.init()
-
-        expect(dogStatsD._flushes()).to.equal(0)
-
-        clock.tick(10000)
-
-        expect(dogStatsD._flushes()).to.equal(1)
-      })
-
       it('should expose real metrics methods after init when configured', () => {
         config.dogstatsd = {
           hostname: 'localhost',