From dd677136dcd81905fa7dc65fd5b002d16e810d90 Mon Sep 17 00:00:00 2001
From: Attila Szegedi <attila.szegedi@datadoghq.com>
Date: Wed, 15 May 2024 15:49:20 +0200
Subject: [PATCH] Logic for heuristic-triggered profiler start. Also make sure
 mock profiler event emitter runs until heuristics trigger.

---
 packages/dd-trace/src/profiler.js                |  3 ++-
 packages/dd-trace/src/profiling/config.js        |  1 +
 packages/dd-trace/src/profiling/profiler.js      |  2 +-
 packages/dd-trace/src/profiling/ssi-telemetry.js | 10 ++++++++--
 packages/dd-trace/src/proxy.js                   | 13 +++++++++++--
 5 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/packages/dd-trace/src/profiler.js b/packages/dd-trace/src/profiler.js
index 349f0438d7c..ce08a3d0f6f 100644
--- a/packages/dd-trace/src/profiler.js
+++ b/packages/dd-trace/src/profiler.js
@@ -9,7 +9,7 @@ process.once('beforeExit', () => { profiler.stop() })
 module.exports = {
   start: config => {
     const { service, version, env, url, hostname, port, tags, repositoryUrl, commitSHA } = config
-    const { enabled, sourceMap, exporters } = config.profiling
+    const { enabled, sourceMap, exporters, heuristicsEnabled } = config.profiling
     const logger = {
       debug: (message) => log.debug(message),
       info: (message) => log.info(message),
@@ -19,6 +19,7 @@ module.exports = {
 
     return profiler.start({
       enabled,
+      heuristicsEnabled,
       service,
       version,
       env,
diff --git a/packages/dd-trace/src/profiling/config.js b/packages/dd-trace/src/profiling/config.js
index 615a728c968..18b85f7c447 100644
--- a/packages/dd-trace/src/profiling/config.js
+++ b/packages/dd-trace/src/profiling/config.js
@@ -65,6 +65,7 @@ class Config {
       DD_PROFILING_PPROF_PREFIX, '')
 
     this.enabled = enabled
+    this.heuristicsEnabled = options.heuristicsEnabled
     this.service = service
     this.env = env
     this.host = host
diff --git a/packages/dd-trace/src/profiling/profiler.js b/packages/dd-trace/src/profiling/profiler.js
index 2c0ebcd5584..17422c5f993 100644
--- a/packages/dd-trace/src/profiling/profiler.js
+++ b/packages/dd-trace/src/profiling/profiler.js
@@ -55,7 +55,7 @@ class Profiler extends EventEmitter {
     if (this._enabled) return true
 
     const config = this._config = new Config(options)
-    if (!config.enabled) return false
+    if (!config.enabled && !config.heuristicsEnabled) return false
 
     this._logger = config.logger
     this._enabled = true
diff --git a/packages/dd-trace/src/profiling/ssi-telemetry.js b/packages/dd-trace/src/profiling/ssi-telemetry.js
index 1b8b559efb3..0316585c47d 100644
--- a/packages/dd-trace/src/profiling/ssi-telemetry.js
+++ b/packages/dd-trace/src/profiling/ssi-telemetry.js
@@ -158,8 +158,14 @@ class SSITelemetry {
     this._profileCount = profilersNamespace.count('ssi_heuristic.number_of_profiles', tags)
     this._runtimeIdCount = profilersNamespace.count('ssi_heuristic.number_of_runtime_id', tags)
 
-    if (!this._emittedRuntimeId && decision[0] === 'triggered') {
-      // Tags won't change anymore, so we can emit the runtime ID metric now
+    if (
+      !this._emittedRuntimeId &&
+      decision[0] === 'triggered' &&
+      // When enablement choice is SSI_ENABLED, hasSentProfiles can transition from false to true when the
+      // profiler gets started and the first profile is submitted, so we have to wait for it.
+      (this.enablementChoice !== EnablementChoice.SSI_ENABLED || this.hasSentProfiles)
+    ) {
+      // Tags won't change anymore, so we can emit the runtime ID metric now.
       this._emittedRuntimeId = true
       this._runtimeIdCount.inc()
     }
diff --git a/packages/dd-trace/src/proxy.js b/packages/dd-trace/src/proxy.js
index 329c16e00ba..addce75b611 100644
--- a/packages/dd-trace/src/proxy.js
+++ b/packages/dd-trace/src/proxy.js
@@ -100,8 +100,17 @@ class Tracer extends NoopProxy {
       ssiTelemetry.start()
       if (config.profiling.enabled) {
         this._profilerStarted = this._startProfiler(config)
-      } else if (ssiTelemetry.enabled()) {
-        require('./profiling/ssi-telemetry-mock-profiler').start(config)
+      } else if (config.profiling.ssi) {
+        const mockProfiler = require('./profiling/ssi-telemetry-mock-profiler')
+        mockProfiler.start(config)
+
+        if (config.profiling.heuristicsEnabled) {
+          ssiTelemetry.onHeuristicsTriggered(() => {
+            mockProfiler.stop()
+            this._startProfiler(config)
+            ssiTelemetry.onHeuristicsTriggered()
+          })
+        }
       }
       if (!this._profilerStarted) {
         this._profilerStarted = Promise.resolve(false)