precreated span_id (sibling spans)

DataDog · Jul 9, 2024 · d29a0a4 · d29a0a4
1 parent 02049b8
commit d29a0a4
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 22 deletions.
diff --git a/...ain/java/datadog/trace/instrumentation/jdbc/AbstractPreparedStatementInstrumentation.java b/...ain/java/datadog/trace/instrumentation/jdbc/AbstractPreparedStatementInstrumentation.java
@@ -17,6 +17,7 @@
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBQueryInfo;
 import java.sql.Connection;
@@ -65,30 +66,31 @@ public static AgentScope onEnter(@Advice.This final Statement statement) {
         return null;
       }
       try {
-        Connection connection = statement.getConnection();
-        DBQueryInfo queryInfo =
+        final Connection connection = statement.getConnection();
+        final DBQueryInfo queryInfo =
             InstrumentationContext.get(Statement.class, DBQueryInfo.class).get(statement);
         if (null == queryInfo) {
           logMissingQueryInfo(statement);
           return null;
         }
+        final AgentSpan span;
+        final DBInfo dbInfo =
+                JDBCDecorator.parseDBInfo(
+                        connection, InstrumentationContext.get(Connection.class, DBInfo.class));
+        final boolean isSqlServer = dbInfo.getType().equals("sqlserver");
+        final boolean injectTraceContext = DECORATE.shouldInjectTraceContext(dbInfo);
 
-        final AgentSpan span = startSpan(DATABASE_QUERY);
-        DECORATE.afterStart(span);
-        DBInfo dbInfo =
-            JDBCDecorator.parseDBInfo(
-                connection, InstrumentationContext.get(Connection.class, DBInfo.class));
-        DECORATE.onConnection(span, dbInfo);
-        DECORATE.onPreparedStatement(span, queryInfo);
-
-        final AgentScope scope = activateSpan(span);
-        boolean isSqlServer = dbInfo.getType().equals("sqlserver");
-        boolean injectTraceContext = DECORATE.shouldInjectTraceContext(dbInfo);
         if (isSqlServer && INJECT_COMMENT && injectTraceContext) {
-          DECORATE.setContextInfo(connection, span.getSpanId(), dbInfo);
+          final long spanID = DECORATE.setContextInfo(connection, dbInfo);
+          span = AgentTracer.get().buildSpan(DATABASE_QUERY).withSpanId(spanID).start();
+        } else {
+          span = startSpan(DATABASE_QUERY);
         }
+        DECORATE.onConnection(span, dbInfo);
+        DECORATE.onPreparedStatement(span, queryInfo);
+        DECORATE.afterStart(span);
 
-        return scope;
+        return activateSpan(span);
       } catch (SQLException e) {
         logSQLException(e);
         // if we can't get the connection for any reason

diff --git a/.../instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java b/.../instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java
@@ -251,7 +251,8 @@ public String traceParent(AgentSpan span, int samplingPriority) {
    * @param connection The same connection as the one that will be used for the actual statement
    * @param spanID The ID of the span covering the actual statement
    */
-  public void setContextInfo(Connection connection, long spanID, DBInfo dbInfo) {
+  public long setContextInfo(Connection connection, DBInfo dbInfo) {
+    final long spanID = Config.get().getIdGenerationStrategy().generateSpanId();
     AgentSpan instrumentationSpan = startSpan("set context_info");
     DECORATE.afterStart(instrumentationSpan);
     DECORATE.onConnection(instrumentationSpan, dbInfo);
@@ -278,6 +279,7 @@ public void setContextInfo(Connection connection, long spanID, DBInfo dbInfo) {
     } finally {
       instrumentationSpan.finish();
     }
+    return spanID;
   }
 
   @Override

diff --git a/...ation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/...ation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java
@@ -20,6 +20,7 @@
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
 import java.sql.Connection;
 import java.sql.SQLException;
@@ -86,8 +87,17 @@ public static AgentScope onEnter(
                 connection, InstrumentationContext.get(Connection.class, DBInfo.class));
         boolean injectTraceContext = DECORATE.shouldInjectTraceContext(dbInfo);
         boolean isSqlServer = dbInfo.getType().equals("sqlserver");
+        final long spanID;
+        final AgentSpan span;
+
+        if (isSqlServer && INJECT_COMMENT && injectTraceContext) {
+          spanID = DECORATE.setContextInfo(connection, dbInfo);
+          span = AgentTracer.get().buildSpan(DATABASE_QUERY).withSpanId(spanID).start();
+
+        } else {
+          span = startSpan(DATABASE_QUERY);
+        }
 
-        final AgentSpan span = startSpan(DATABASE_QUERY);
         DECORATE.afterStart(span);
         DECORATE.onConnection(span, dbInfo);
         final String copy = sql;
@@ -117,10 +127,6 @@ public static AgentScope onEnter(
         }
         DECORATE.onStatement(span, copy);
 
-        if (isSqlServer && INJECT_COMMENT && injectTraceContext) {
-          DECORATE.setContextInfo(connection, span.getSpanId(), dbInfo);
-        }
-
         return scope;
       } catch (SQLException e) {
         // if we can't get the connection for any reason

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java
@@ -1268,6 +1268,7 @@ public class CoreSpanBuilder implements AgentTracer.SpanBuilder {
     private Object builderRequestContextDataIast;
     private Object builderCiVisibilityContextData;
     private List<AgentSpanLink> links;
+    private long spanId;
 
     CoreSpanBuilder(
         final String instrumentationName, final CharSequence operationName, CoreTracer tracer) {
@@ -1423,6 +1424,13 @@ public AgentTracer.SpanBuilder withLink(AgentSpanLink link) {
       return this;
     }
 
+    @Override
+    public CoreSpanBuilder withSpanId(final long spanId) {
+      this.spanId = spanId;
+      return this;
+    }
+
+
     /**
      * Build the SpanContext, if the actual span has a parent, the following attributes must be
      * propagated: - ServiceName - Baggage - Trace (a list of all spans related) - SpanType
@@ -1431,7 +1439,7 @@ public AgentTracer.SpanBuilder withLink(AgentSpanLink link) {
      */
     private DDSpanContext buildSpanContext() {
       final DDTraceId traceId;
-      final long spanId = idGenerationStrategy.generateSpanId();
+      final long spanId;
       final long parentSpanId;
       final Map<String, String> baggage;
       final PendingTrace parentTrace;
@@ -1447,6 +1455,11 @@ private DDSpanContext buildSpanContext() {
       final PathwayContext pathwayContext;
       final PropagationTags propagationTags;
 
+      if (this.spanId == 0) {
+        spanId = idGenerationStrategy.generateSpanId();
+      } else {
+        spanId = this.spanId;
+      }
       // FIXME [API] parentContext should be an interface implemented by ExtractedContext,
       // TagContext, DDSpanContext, AgentSpan.Context
       AgentSpan.Context parentContext = parent;

diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java
@@ -321,6 +321,8 @@ public interface SpanBuilder {
     <T> SpanBuilder withRequestContextData(RequestContextSlot slot, T data);
 
     SpanBuilder withLink(AgentSpanLink link);
+
+    SpanBuilder withSpanId(long spanId);
   }
 
   static class NoopTracerAPI implements TracerAPI {