From 666f2afac1ce3e009e38669007a259d3b37867c5 Mon Sep 17 00:00:00 2001
From: Mikayla Toffler <46911781+mtoffl01@users.noreply.github.com>
Date: Thu, 17 Oct 2024 12:57:15 -0400
Subject: [PATCH] Support DD_TRACE_<INTEGRATION>_ENABLED (#7718)

---
 .../datadog/trace/api/InstrumenterConfig.java | 28 ++++++++-
 .../trace/api/InstrumenterConfigTest.groovy   | 63 +++++++++++++++++++
 2 files changed, 90 insertions(+), 1 deletion(-)

diff --git a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java
index 78152292a48..8585ddcf54f 100644
--- a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java
+++ b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java
@@ -286,9 +286,35 @@ public boolean isIntegrationsEnabled() {
     return integrationsEnabled;
   }
 
+  /**
+   * isIntegrationEnabled determines whether an integration under the specified name(s) is enabled
+   * according to the following list of configurations, from highest to lowest precedence:
+   * trace.name.enabled, trace.integration.name.enabled, integration.name.enabled. If none of these
+   * configurations is set, the defaultEnabled value is used. All system properties take precedence
+   * over all env vars.
+   *
+   * @param integrationNames the name(s) that represent(s) the integration
+   * @param defaultEnabled true if enabled by default, else false
+   * @return boolean on whether the integration is enabled
+   */
   public boolean isIntegrationEnabled(
       final Iterable<String> integrationNames, final boolean defaultEnabled) {
-    return configProvider.isEnabled(integrationNames, "integration.", ".enabled", defaultEnabled);
+    // If default is enabled, we want to disable individually.
+    // If default is disabled, we want to enable individually.
+    boolean anyEnabled = defaultEnabled;
+    for (final String name : integrationNames) {
+      final String primaryKey = "trace." + name + ".enabled";
+      final String[] aliases = {
+        "trace.integration." + name + ".enabled", "integration." + name + ".enabled"
+      }; // listed in order of precedence
+      final boolean configEnabled = configProvider.getBoolean(primaryKey, defaultEnabled, aliases);
+      if (defaultEnabled) {
+        anyEnabled &= configEnabled;
+      } else {
+        anyEnabled |= configEnabled;
+      }
+    }
+    return anyEnabled;
   }
 
   public boolean isIntegrationShortcutMatchingEnabled(
diff --git a/internal-api/src/test/groovy/datadog/trace/api/InstrumenterConfigTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/InstrumenterConfigTest.groovy
index 0a00703bff7..b7f429e2509 100644
--- a/internal-api/src/test/groovy/datadog/trace/api/InstrumenterConfigTest.groovy
+++ b/internal-api/src/test/groovy/datadog/trace/api/InstrumenterConfigTest.groovy
@@ -8,6 +8,7 @@ class InstrumenterConfigTest extends DDSpecification {
     setup:
     environmentVariables.set("DD_INTEGRATION_ORDER_ENABLED", "false")
     environmentVariables.set("DD_INTEGRATION_TEST_ENV_ENABLED", "true")
+    environmentVariables.set("DD_TRACE_NEW_ENV_ENABLED", "false")
     environmentVariables.set("DD_INTEGRATION_DISABLED_ENV_ENABLED", "false")
 
     System.setProperty("dd.integration.order.enabled", "true")
@@ -16,6 +17,7 @@ class InstrumenterConfigTest extends DDSpecification {
 
     environmentVariables.set("DD_INTEGRATION_ORDER_MATCHING_SHORTCUT_ENABLED", "false")
     environmentVariables.set("DD_INTEGRATION_TEST_ENV_MATCHING_SHORTCUT_ENABLED", "true")
+    environmentVariables.set("DD_INTEGRATION_NEW_ENV_MATCHING_SHORTCUT_ENABLED", "false")
     environmentVariables.set("DD_INTEGRATION_DISABLED_ENV_MATCHING_SHORTCUT_ENABLED", "false")
 
     System.setProperty("dd.integration.order.matching.shortcut.enabled", "true")
@@ -44,11 +46,72 @@ class InstrumenterConfigTest extends DDSpecification {
     ["disabled-env", "test-env"]   | false          | true
     ["test-prop", "disabled-prop"] | true           | false
     ["disabled-env", "test-env"]   | true           | false
+    ["new-env"] | true | false
     // spotless:on
 
     integrationNames = new TreeSet<>(names)
   }
 
+  def setEnv(String key, String value) {
+    environmentVariables.set(key, value)
+  }
+
+  def setSysProp(String key, String value) {
+    System.setProperty(key, value)
+  }
+
+  def randomIntegrationEnabled() {
+    return InstrumenterConfig.get().isIntegrationEnabled(["random"], true)
+  }
+
+  def "verify integration enabled hierarchy"() {
+    when:
+    // the below should have no effect
+    setEnv("DD_RANDOM_ENABLED", "false")
+    setSysProp("dd.random.enabled", "false")
+
+    then:
+    randomIntegrationEnabled() == true
+
+    when:
+    setEnv("DD_INTEGRATION_RANDOM_ENABLED", "false")
+
+    then:
+    randomIntegrationEnabled() == false
+
+    when:
+    setEnv("DD_TRACE_INTEGRATION_RANDOM_ENABLED", "true")
+
+    then:
+    randomIntegrationEnabled() == true
+
+    when:
+    setEnv("DD_TRACE_RANDOM_ENABLED", "false")
+
+    then:
+    randomIntegrationEnabled() == false
+
+    // assert all system properties take precedence over all env vars
+    when:
+    setSysProp("dd.integration.random.enabled", "true")
+
+    then:
+    randomIntegrationEnabled() == true
+
+    when:
+    setSysProp("dd.trace.integration.random.enabled", "false")
+
+    then:
+    randomIntegrationEnabled() == false
+
+    when:
+    setSysProp("dd.trace.random.enabled", "true")
+
+    then:
+    randomIntegrationEnabled() == true
+
+  }
+
   def "valid resolver presets"() {
     setup:
     injectSysConfig("resolver.cache.config", preset)