Propagate Oracle dynamic service hash without SQL comments

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/jdbc/DBInfo.java

@@ -18,6 +18,7 @@ public final class DBInfo {
   private final String warehouse;
   private final String schema;
   private volatile String poolName;
+  private volatile String oracleServiceAction;
 
   DBInfo(
       String type,
@@ -220,6 +221,14 @@ public void setPoolName(String poolname) {
     this.poolName = poolname;
   }
 
+  public synchronized boolean markOracleServiceAction(String action) {
+    if (action.equals(oracleServiceAction)) {
+      return false;
+    }
+    oracleServiceAction = action;
+    return true;
+  }
+
   public Builder toBuilder() {
     return new Builder(
         type,

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/AbstractPreparedStatementInstrumentation.java

@@ -81,6 +81,8 @@ public static AgentScope onEnter(@Advice.This final Statement statement) {
                 connection, InstrumentationContext.get(Connection.class, DBInfo.class));
         final boolean injectTraceContext = DECORATE.shouldInjectTraceContext(dbInfo);
 
+        DECORATE.setServiceHashAction(connection, dbInfo);
+
         if (INJECT_COMMENT && injectTraceContext) {
           if (DECORATE.isSqlServer(dbInfo)) {
             // The span ID is pre-determined so that we can reference it when setting the context

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java

@@ -123,6 +123,9 @@ public static String onEnter(
       final DBInfo dbInfo =
           JDBCDecorator.parseDBInfo(
               connection, InstrumentationContext.get(Connection.class, DBInfo.class));
+      if (!DECORATE.shouldInjectSqlComment(dbInfo)) {
+        return inputSql;
+      }
       String dbService = DECORATE.getDbService(dbInfo);
       if (dbService != null) {
         dbService = traceConfig(activeSpan).getServiceMapping().getOrDefault(dbService, dbService);

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DriverInstrumentation.java

@@ -113,7 +113,8 @@ public static void addDBInfo(
           // ignore
         }
       }
-      DBInfo dbInfo = JDBCConnectionUrlParser.extractDBInfo(connectionUrl, connectionProps);
+      DBInfo dbInfo =
+          JDBCConnectionUrlParser.extractDBInfo(connectionUrl, connectionProps).toBuilder().build();
       InstrumentationContext.get(Connection.class, DBInfo.class).put(connWithContext, dbInfo);
     }
   }

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java

@@ -54,6 +54,7 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
       SpanNaming.instance().namingSchema().database().service("jdbc");
 
   public static final String DD_INSTRUMENTATION_PREFIX = "_DD_";
+  public static final String DD_ORACLE_SERVICE_HASH_PREFIX = "_DD_DDSH:";
 
   public static final String DBM_PROPAGATION_MODE = Config.get().getDbmPropagationMode();
   private static final boolean DBM_INJECT_SQL_BASE_HASH = Config.get().isDbmInjectSqlBaseHash();
@@ -65,6 +66,9 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
           || DBM_PROPAGATION_MODE.equals(DBM_PROPAGATION_MODE_DYNAMIC_SERVICE);
   private static final boolean INJECT_TRACE_CONTEXT =
       DBM_PROPAGATION_MODE.equals(DBM_PROPAGATION_MODE_FULL);
+  private static final boolean INJECT_ORACLE_SERVICE_HASH_ACTION =
+      DBM_PROPAGATION_MODE.equals(DBM_PROPAGATION_MODE_DYNAMIC_SERVICE)
+          && Config.get().isDbmPropagationOracleActionEnabled();
   public static final boolean DBM_TRACE_PREPARED_STATEMENTS =
       Config.get().isDbmTracePreparedStatements();
   public static final boolean DBM_ALWAYS_APPEND_SQL_COMMENT =
@@ -238,7 +242,7 @@ public static DBInfo parseDBInfoFromConnection(final Connection connection) {
           // getClientInfo is likely not allowed, we can still extract info from the url alone
           log.debug(LogCollector.EXCLUDE_TELEMETRY, "Could not get client info from DB", ex);
         }
-        dbInfo = JDBCConnectionUrlParser.extractDBInfo(url, clientInfo);
+        dbInfo = JDBCConnectionUrlParser.extractDBInfo(url, clientInfo).toBuilder().build();
       } else {
         dbInfo = DBInfo.DEFAULT;
       }
@@ -292,6 +296,33 @@ public boolean isSqlServer(final DBInfo dbInfo) {
     return "sqlserver".equals(dbInfo.getType());
   }
 
+  public boolean shouldInjectSqlComment(final DBInfo dbInfo) {
+    return INJECT_COMMENT && !(INJECT_ORACLE_SERVICE_HASH_ACTION && isOracle(dbInfo));
+  }
+
+  /** Sets the dynamic service hash in {@code v$session.action} once per Oracle session and hash. */
+  public void setServiceHashAction(Connection connection, DBInfo dbInfo) {
+    if (!INJECT_ORACLE_SERVICE_HASH_ACTION || !isOracle(dbInfo)) {
+      return;
+    }
+
+    final String baseHash = BaseHash.getBaseHashStr();
+    if (baseHash == null) {
+      return;
+    }
+    final String action = DD_ORACLE_SERVICE_HASH_PREFIX + baseHash;
+    if (!dbInfo.markOracleServiceAction(action)) {
+      return;
+    }
+
+    try {
+      connection.setClientInfo("OCSID.ACTION", action);
+    } catch (Throwable e) {
+      // The attempt stays recorded so unsupported drivers do not pay this cost on every query.
+      logInjectionErrorOnce("service hash action", e);
+    }
+  }
+
   /**
    * Executes `connection.setClientInfo("OCSID.ACTION", traceContext)` statement on the Oracle DB to
    * set the trace parent in `v$session.action`. This is used because it isn't possible to propagate

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java

@@ -93,6 +93,8 @@ public static AgentScope onEnter(
         final boolean isSqlServer = DECORATE.isSqlServer(dbInfo);
         final boolean isOracle = DECORATE.isOracle(dbInfo);
 
+        DECORATE.setServiceHashAction(connection, dbInfo);
+
         if (INJECT_COMMENT && injectTraceContext) {
           if (isSqlServer) {
             // The span ID is pre-determined so that we can reference it when setting the context
@@ -116,7 +118,7 @@ public static AgentScope onEnter(
         DECORATE.afterStart(span);
         DECORATE.onConnection(span, dbInfo);
         final String copy = sql;
-        if (span != null && INJECT_COMMENT) {
+        if (span != null && DECORATE.shouldInjectSqlComment(dbInfo)) {
           String traceParent = null;
 
           if (injectTraceContext) {

dd-java-agent/instrumentation/jdbc/src/test/groovy/OracleInjectionForkedTest.groovy

@@ -1,5 +1,9 @@
 import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.BaseHash
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.ProcessTags
 import datadog.trace.api.config.TraceInstrumentationConfig
+import datadog.trace.bootstrap.instrumentation.api.Tags
 import test.TestConnection
 import test.TestDatabaseMetaData
 import test.TestPreparedStatement
@@ -75,3 +79,96 @@ class OracleInjectionForkedTest extends OracleInjectionTestBase {
     serviceNameUrl | serviceNameInjection
   }
 }
+
+class OracleDynamicServiceActionInjectionForkedTest extends OracleInjectionTestBase {
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+
+    injectSysConfig(TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE, "dynamic_service")
+    injectSysConfig(
+      TraceInstrumentationConfig.DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED, "true")
+  }
+
+  def setup() {
+    ProcessTags.reset()
+    BaseHash.updateBaseHash(-6937226773133363462L)
+  }
+
+  def "Oracle dynamic service mode propagates the hash in ACTION without changing statement SQL"() {
+    setup:
+    def connection = createOracleConnection(serviceNameUrl)
+    def statement = connection.createStatement() as TestStatement
+
+    when:
+    statement.executeQuery(query)
+    statement.executeQuery(query)
+
+    then:
+    statement.sql == query
+    connection.clientInfoName == "OCSID.ACTION"
+    connection.clientInfoValue == "_DD_DDSH:-6937226773133363462"
+    connection.clientInfoSetCount == 1
+    assertTraces(2) {
+      trace(1) {
+        span {
+          spanType DDSpanTypes.SQL
+          tags(false) {
+            "$Tags.BASE_HASH" "-6937226773133363462"
+          }
+        }
+      }
+      trace(1) {
+        span {
+          spanType DDSpanTypes.SQL
+          tags(false) {
+            "$Tags.BASE_HASH" "-6937226773133363462"
+          }
+        }
+      }
+    }
+  }
+
+  def "Oracle dynamic service mode preserves prepared statement SQL"() {
+    setup:
+    def connection = createOracleConnection(sidUrl)
+
+    when:
+    def statement = connection.prepareStatement(query) as TestPreparedStatement
+    statement.execute()
+
+    then:
+    statement.sql == query
+    connection.clientInfoValue == "_DD_DDSH:-6937226773133363462"
+    connection.clientInfoSetCount == 1
+  }
+
+  def "Oracle dynamic service mode refreshes ACTION only when the hash changes"() {
+    setup:
+    def connection = createOracleConnection(serviceNameUrl)
+    def statement = connection.createStatement() as TestStatement
+
+    when:
+    statement.executeQuery(query)
+    BaseHash.updateBaseHash(123456789L)
+    statement.executeQuery(query)
+
+    then:
+    connection.clientInfoValue == "_DD_DDSH:123456789"
+    connection.clientInfoSetCount == 2
+  }
+
+  def "Oracle dynamic service mode initializes every connection with the same URL"() {
+    setup:
+    def firstConnection = createOracleConnection(serviceNameUrl)
+    def secondConnection = createOracleConnection(serviceNameUrl)
+
+    when:
+    firstConnection.createStatement().executeQuery(query)
+    secondConnection.createStatement().executeQuery(query)
+
+    then:
+    firstConnection.clientInfoSetCount == 1
+    secondConnection.clientInfoSetCount == 1
+  }
+}

dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestConnection.groovy

@@ -22,6 +22,10 @@ import java.util.concurrent.Executor
  * A JDBC connection class that optionally throws an exception in the constructor, used to test
  */
 class TestConnection implements Connection {
+  public String clientInfoName
+  public String clientInfoValue
+  public int clientInfoSetCount
+
   TestConnection(boolean throwException) {
     if (throwException) {
       throw new RuntimeException("connection exception")
@@ -232,6 +236,9 @@ class TestConnection implements Connection {
 
   @Override
   void setClientInfo(String name, String value) throws SQLClientInfoException {
+    clientInfoName = name
+    clientInfoValue = value
+    clientInfoSetCount++
   }
 
   @Override

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -73,6 +73,7 @@ public final class ConfigDefaults {
   static final boolean DEFAULT_DB_CLIENT_HOST_SPLIT_BY_INSTANCE_TYPE_SUFFIX = false;
   static final boolean DEFAULT_DB_CLIENT_HOST_SPLIT_BY_HOST = false;
   static final String DEFAULT_DB_DBM_PROPAGATION_MODE_MODE = "disabled";
+  static final boolean DEFAULT_DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED = false;
   static final boolean DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS = false;
   static final boolean DEFAULT_DB_DBM_ALWAYS_APPEND_SQL_COMMENT = false;
   // Default value is set to 0, it disables the latency trace interceptor

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -74,6 +74,8 @@ public final class TraceInstrumentationConfig {
 
   public static final String DB_DBM_INJECT_SQL_BASEHASH = "dbm.inject.sql.basehash";
   public static final String DB_DBM_PROPAGATION_MODE_MODE = "dbm.propagation.mode";
+  public static final String DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED =
+      "dbm.propagation.oracle.action.enabled";
   public static final String DB_DBM_TRACE_PREPARED_STATEMENTS = "dbm.trace_prepared_statements";
   public static final String DB_DBM_ALWAYS_APPEND_SQL_COMMENT = "dbm.always_append_sql_comment";
 

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -59,6 +59,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_CLIENT_HOST_SPLIT_BY_INSTANCE_TYPE_SUFFIX;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_ALWAYS_APPEND_SQL_COMMENT;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_PROPAGATION_MODE_MODE;
+import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DEBUGGER_EXCEPTION_CAPTURE_INTERMEDIATE_SPANS_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DEBUGGER_EXCEPTION_CAPTURE_INTERVAL_SECONDS;
@@ -583,6 +584,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_INJECT_SQL_BASEHASH;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE;
+import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_TRACE_PREPARED_STATEMENTS;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_METADATA_FETCHING_ON_CONNECT;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_METADATA_FETCHING_ON_QUERY;
@@ -1246,6 +1248,7 @@ public static String getHostName() {
 
   private final boolean dbmInjectSqlBaseHash;
   private final String dbmPropagationMode;
+  private final boolean dbmPropagationOracleActionEnabled;
   private final boolean dbmTracePreparedStatements;
   private final boolean dbmAlwaysAppendSqlComment;
   private final boolean dbMetadataFetchingOnQuery;
@@ -1840,6 +1843,11 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
         configProvider.getString(
             DB_DBM_PROPAGATION_MODE_MODE, DEFAULT_DB_DBM_PROPAGATION_MODE_MODE);
 
+    dbmPropagationOracleActionEnabled =
+        configProvider.getBoolean(
+            DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED,
+            DEFAULT_DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED);
+
     dbmTracePreparedStatements =
         configProvider.getBoolean(
             DB_DBM_TRACE_PREPARED_STATEMENTS, DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS);
@@ -6000,6 +6008,10 @@ public String getDbmPropagationMode() {
     return dbmPropagationMode;
   }
 
+  public boolean isDbmPropagationOracleActionEnabled() {
+    return dbmPropagationOracleActionEnabled;
+  }
+
   // Database monitoring propagation mode constants
   public static final String DBM_PROPAGATION_MODE_STATIC = "service";
   public static final String DBM_PROPAGATION_MODE_FULL = "full";
@@ -6579,6 +6591,8 @@ public String toString() {
         + dbmInjectSqlBaseHash
         + ", dbmPropagationMode="
         + dbmPropagationMode
+        + ", dbmPropagationOracleActionEnabled="
+        + dbmPropagationOracleActionEnabled
         + ", dbmTracePreparedStatements="
         + dbmTracePreparedStatements
         + ", splitByTags="

internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy

@@ -100,6 +100,7 @@ import static datadog.trace.api.config.RemoteConfigConfig.REMOTE_CONFIG_URL
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_HOST
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE_TYPE_SUFFIX
+import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED
 import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN
 import static datadog.trace.api.config.TraceInstrumentationConfig.RUNTIME_CONTEXT_FIELD_INJECTION
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_ENABLED
@@ -3427,6 +3428,26 @@ class ConfigTest extends DDSpecification {
     "false" | "true"  | false   // sys prop takes precedence
   }
 
+  def "Oracle DBM action propagation enabled = #configured"() {
+    setup:
+    def properties = new Properties()
+    if (configured != null) {
+      properties.setProperty(DB_DBM_PROPAGATION_ORACLE_ACTION_ENABLED, configured)
+    }
+
+    when:
+    def config = new Config(ConfigProvider.withPropertiesOverride(properties))
+
+    then:
+    config.isDbmPropagationOracleActionEnabled() == expected
+
+    where:
+    configured | expected
+    null       | false
+    "false"    | false
+    "true"     | true
+  }
+
   def "trace resource renaming activation with appsec=#appsec and explicit=#explicit"() {
     setup:
     if (appsec != null) {

metadata/supported-configurations.json

@@ -1185,6 +1185,14 @@
         "aliases": []
       }
     ],
+    "DD_DBM_PROPAGATION_ORACLE_ACTION_ENABLED": [
+      {
+        "version": "A",
+        "type": "boolean",
+        "default": "false",
+        "aliases": []
+      }
+    ],
     "DD_DBM_TRACE_PREPARED_STATEMENTS": [
       {
         "version": "A",