Reduce per-query allocations in SQLCommenter.inject()

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/dbm/SharedDBCommenter.java

@@ -9,6 +9,7 @@
 import java.io.UnsupportedEncodingException;
 import java.net.URLEncoder;
 import java.nio.charset.StandardCharsets;
+import java.util.concurrent.ConcurrentHashMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,17 +33,63 @@ public class SharedDBCommenter {
   private static final String TRACEPARENT = encode("traceparent");
   private static final String DD_SERVICE_HASH = encode("ddsh");
 
+  /**
+   * Cache for static comment strings (those without traceParent or peerService). The key combines
+   * dbService, hostname, dbName, and Config identity to ensure correctness if Config is replaced.
+   */
+  private static final ConcurrentHashMap<StaticCommentKey, String> staticCommentCache =
+      new ConcurrentHashMap<>();
+
+  // Pre-computed marker strings for trace comment detection
+  private static final String PARENT_SERVICE_EQ = PARENT_SERVICE + "=";
+  private static final String DATABASE_SERVICE_EQ = DATABASE_SERVICE + "=";
+  private static final String DD_HOSTNAME_EQ = DD_HOSTNAME + "=";
+  private static final String DD_DB_NAME_EQ = DD_DB_NAME + "=";
+  private static final String DD_PEER_SERVICE_EQ = DD_PEER_SERVICE + "=";
+  private static final String DD_ENV_EQ = DD_ENV + "=";
+  private static final String DD_VERSION_EQ = DD_VERSION + "=";
+  private static final String TRACEPARENT_EQ = TRACEPARENT + "=";
+  private static final String DD_SERVICE_HASH_EQ = DD_SERVICE_HASH + "=";
+
   // Used by SQLCommenter and MongoCommentInjector to avoid duplicate comment injection
   public static boolean containsTraceComment(String commentContent) {
-    return commentContent.contains(PARENT_SERVICE + "=")
-        || commentContent.contains(DATABASE_SERVICE + "=")
-        || commentContent.contains(DD_HOSTNAME + "=")
-        || commentContent.contains(DD_DB_NAME + "=")
-        || commentContent.contains(DD_PEER_SERVICE + "=")
-        || commentContent.contains(DD_ENV + "=")
-        || commentContent.contains(DD_VERSION + "=")
-        || commentContent.contains(TRACEPARENT + "=")
-        || commentContent.contains(DD_SERVICE_HASH + "=");
+    return commentContent.contains(PARENT_SERVICE_EQ)
+        || commentContent.contains(DATABASE_SERVICE_EQ)
+        || commentContent.contains(DD_HOSTNAME_EQ)
+        || commentContent.contains(DD_DB_NAME_EQ)
+        || commentContent.contains(DD_PEER_SERVICE_EQ)
+        || commentContent.contains(DD_ENV_EQ)
+        || commentContent.contains(DD_VERSION_EQ)
+        || commentContent.contains(TRACEPARENT_EQ)
+        || commentContent.contains(DD_SERVICE_HASH_EQ);
+  }
+
+  /**
+   * Checks for trace comment markers within a range of the given string, without allocating a
+   * substring. Searches within [fromIndex, toIndex) of the source string.
+   */
+  public static boolean containsTraceComment(String sql, int fromIndex, int toIndex) {
+    return containsInRange(sql, PARENT_SERVICE_EQ, fromIndex, toIndex)
+        || containsInRange(sql, DATABASE_SERVICE_EQ, fromIndex, toIndex)
+        || containsInRange(sql, DD_HOSTNAME_EQ, fromIndex, toIndex)
+        || containsInRange(sql, DD_DB_NAME_EQ, fromIndex, toIndex)
+        || containsInRange(sql, DD_PEER_SERVICE_EQ, fromIndex, toIndex)
+        || containsInRange(sql, DD_ENV_EQ, fromIndex, toIndex)
+        || containsInRange(sql, DD_VERSION_EQ, fromIndex, toIndex)
+        || containsInRange(sql, TRACEPARENT_EQ, fromIndex, toIndex)
+        || containsInRange(sql, DD_SERVICE_HASH_EQ, fromIndex, toIndex);
+  }
+
+  /** Checks if {@code target} appears within the range [fromIndex, toIndex) of {@code source}. */
+  private static boolean containsInRange(String source, String target, int fromIndex, int toIndex) {
+    int targetLen = target.length();
+    int limit = toIndex - targetLen;
+    for (int i = fromIndex; i <= limit; i++) {
+      if (source.regionMatches(i, target, 0, targetLen)) {
+        return true;
+      }
+    }
+    return false;
   }
 
   // Build database comment content without comment delimiters such as /* */
@@ -69,6 +116,60 @@ public static String buildComment(
     return sb.length() > 0 ? sb.toString() : null;
   }
 
+  /**
+   * Builds the static portion of a database comment that does not change per-span. This includes
+   * parentService, databaseService, hostname, dbName, env, version, and serviceHash. The dynamic
+   * parts (peerService, traceParent) are excluded and must be appended separately.
+   *
+   * <p>Results are cached per (dbService, hostname, dbName, Config) combination to avoid redundant
+   * URLEncoder.encode() calls and StringBuilder work on every query execution.
+   *
+   * @return the static comment prefix, or null if no static fields are set
+   */
+  public static String buildStaticComment(String dbService, String hostname, String dbName) {
+    Config config = Config.get();
+    StaticCommentKey key = new StaticCommentKey(dbService, hostname, dbName, config);
+    String cached = staticCommentCache.get(key);
+    if (cached != null) {
+      return cached;
+    }
+
+    StringBuilder sb = new StringBuilder();
+
+    int initSize = 0;
+    append(sb, PARENT_SERVICE, config.getServiceName(), initSize);
+    append(sb, DATABASE_SERVICE, dbService, initSize);
+    append(sb, DD_HOSTNAME, hostname, initSize);
+    append(sb, DD_DB_NAME, dbName, initSize);
+    // peerService is per-span, skip here
+    append(sb, DD_ENV, config.getEnv(), initSize);
+    append(sb, DD_VERSION, config.getVersion(), initSize);
+    // traceParent is per-span, skip here
+
+    if (config.isDbmInjectSqlBaseHash() && config.isExperimentalPropagateProcessTagsEnabled()) {
+      append(sb, DD_SERVICE_HASH, BaseHash.getBaseHashStr(), initSize);
+    }
+
+    String result = sb.length() > 0 ? sb.toString() : null;
+    if (result != null) {
+      staticCommentCache.putIfAbsent(key, result);
+    }
+    return result;
+  }
+
+  /** Returns true if the current active span has a non-null, non-empty peer service tag. */
+  public static boolean hasPeerService() {
+    AgentSpan span = activeSpan();
+    if (span != null) {
+      Object peerService = span.getTag(Tags.PEER_SERVICE);
+      if (peerService != null) {
+        String str = peerService.toString();
+        return str != null && !str.isEmpty();
+      }
+    }
+    return false;
+  }
+
   private static String getPeerService() {
     AgentSpan span = activeSpan();
     Object peerService = null;
@@ -104,4 +205,50 @@ private static void append(StringBuilder sb, String key, String value, int initS
     }
     sb.append(key).append(EQUALS).append(QUOTE).append(encodedValue).append(QUOTE);
   }
+
+  /**
+   * Cache key for static comment lookup. Uses Config identity (rather than individual field values)
+   * to ensure the cache is automatically invalidated if Config is replaced (as happens in tests).
+   * In production, Config is created once, so identity comparison is both correct and cheap.
+   */
+  private static final class StaticCommentKey {
+    private final String dbService;
+    private final String hostname;
+    private final String dbName;
+    private final Config config;
+    private final int hashCode;
+
+    StaticCommentKey(String dbService, String hostname, String dbName, Config config) {
+      this.dbService = dbService;
+      this.hostname = hostname;
+      this.dbName = dbName;
+      this.config = config;
+      int h = 17;
+      h = 31 * h + (dbService != null ? dbService.hashCode() : 0);
+      h = 31 * h + (hostname != null ? hostname.hashCode() : 0);
+      h = 31 * h + (dbName != null ? dbName.hashCode() : 0);
+      h = 31 * h + System.identityHashCode(config);
+      this.hashCode = h;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (!(o instanceof StaticCommentKey)) return false;
+      StaticCommentKey that = (StaticCommentKey) o;
+      return config == that.config
+          && eq(dbService, that.dbService)
+          && eq(hostname, that.hostname)
+          && eq(dbName, that.dbName);
+    }
+
+    private static boolean eq(String a, String b) {
+      return a == null ? b == null : a.equals(b);
+    }
+
+    @Override
+    public int hashCode() {
+      return hashCode;
+    }
+  }
 }

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

@@ -15,6 +15,11 @@ public class SQLCommenter {
   private static final int BUFFER_EXTRA = 4;
   private static final int SQL_COMMENT_OVERHEAD = SPACE_CHARS + COMMENT_DELIMITERS + BUFFER_EXTRA;
 
+  /**
+   * Returns the first non-whitespace word in the SQL string. The result is a substring allocation,
+   * so callers that only need to check prefix/equality should use {@link #firstWordStartsWith} or
+   * {@link #firstWordEqualsIgnoreCase} instead.
+   */
   protected static String getFirstWord(String sql) {
     int beginIndex = 0;
     while (beginIndex < sql.length() && Character.isWhitespace(sql.charAt(beginIndex))) {
@@ -27,6 +32,38 @@ protected static String getFirstWord(String sql) {
     return sql.substring(beginIndex, endIndex);
   }
 
+  /**
+   * Checks if the first non-whitespace word in sql starts with the given prefix, without allocating
+   * a substring.
+   */
+  private static boolean firstWordStartsWith(String sql, String prefix) {
+    int beginIndex = 0;
+    int len = sql.length();
+    while (beginIndex < len && Character.isWhitespace(sql.charAt(beginIndex))) {
+      beginIndex++;
+    }
+    return sql.regionMatches(beginIndex, prefix, 0, prefix.length());
+  }
+
+  /**
+   * Checks if the first non-whitespace word in sql equals the given target (case-insensitive),
+   * without allocating a substring.
+   */
+  private static boolean firstWordEqualsIgnoreCase(String sql, String target) {
+    int beginIndex = 0;
+    int len = sql.length();
+    while (beginIndex < len && Character.isWhitespace(sql.charAt(beginIndex))) {
+      beginIndex++;
+    }
+    int endIndex = beginIndex;
+    while (endIndex < len && !Character.isWhitespace(sql.charAt(endIndex))) {
+      endIndex++;
+    }
+    int wordLen = endIndex - beginIndex;
+    return wordLen == target.length()
+        && sql.regionMatches(true, beginIndex, target, 0, target.length());
+  }
+
   public static String inject(
       String sql,
       String dbService,
@@ -40,25 +77,25 @@ public static String inject(
     }
     boolean appendComment = preferAppend;
     if (dbType != null) {
-      final String firstWord = getFirstWord(sql);
+      boolean startsWithBrace = firstWordStartsWith(sql, "{");
 
       // The Postgres JDBC parser doesn't allow SQL comments anywhere in a JDBC
       // callable statements
       // https://github.com/pgjdbc/pgjdbc/blob/master/pgjdbc/src/main/java/org/postgresql/core/Parser.java#L1038
       // TODO: Could we inject the comment after the JDBC has been converted to
       // standard SQL?
-      if (firstWord.startsWith("{") && dbType.startsWith("postgres")) {
+      if (startsWithBrace && dbType.startsWith("postgres")) {
         return sql;
       }
 
       // Append the comment for mysql JDBC callable statements
-      if (firstWord.startsWith("{") && "mysql".equals(dbType)) {
+      if (startsWithBrace && "mysql".equals(dbType)) {
         appendComment = true;
       }
 
       // Both Postgres and MySQL are unhappy with anything before CALL in a stored
       // procedure invocation, but they seem ok with it after so we force append mode
-      if (firstWord.equalsIgnoreCase("call")) {
+      if (firstWordEqualsIgnoreCase(sql, "call")) {
         appendComment = true;
       }
 
@@ -71,8 +108,18 @@ public static String inject(
       return sql;
     }
 
-    String commentContent =
-        SharedDBCommenter.buildComment(dbService, dbType, hostname, dbName, traceParent);
+    // When there are no per-span dynamic fields (traceParent, peerService), use the cached
+    // static comment to avoid redundant URLEncoder.encode() calls and StringBuilder allocations.
+    // This is the common path in DBM "static" propagation mode.
+    boolean hasDynamic =
+        (traceParent != null && !traceParent.isEmpty()) || SharedDBCommenter.hasPeerService();
+    String commentContent;
+    if (!hasDynamic) {
+      commentContent = SharedDBCommenter.buildStaticComment(dbService, hostname, dbName);
+    } else {
+      commentContent =
+          SharedDBCommenter.buildComment(dbService, dbType, hostname, dbName, traceParent);
+    }
 
     if (commentContent == null) {
       return sql;
@@ -112,11 +159,6 @@ private static boolean hasDDComment(String sql, boolean appendComment) {
       return false;
     }
 
-    String commentContent = extractCommentContent(sql, appendComment);
-    return SharedDBCommenter.containsTraceComment(commentContent);
-  }
-
-  private static String extractCommentContent(String sql, boolean appendComment) {
     int startIdx;
     int endIdx;
     if (appendComment) {
@@ -127,9 +169,11 @@ private static String extractCommentContent(String sql, boolean appendComment) {
       endIdx = sql.indexOf(CLOSE_COMMENT);
     }
     if (startIdx != -1 && endIdx != -1 && endIdx > startIdx) {
-      return sql.substring(startIdx + OPEN_COMMENT_LEN, endIdx);
+      // Check for trace comment markers directly in the SQL without allocating a substring.
+      // We search within the bounds [startIdx + OPEN_COMMENT_LEN, endIdx) of the original string.
+      return SharedDBCommenter.containsTraceComment(sql, startIdx + OPEN_COMMENT_LEN, endIdx);
     }
-    return "";
+    return false;
   }
 
   /**