Swith to 128bits for the traceID and prepare the span switch as well …

…(currently still in 64b)
DataDog · Jun 28, 2024 · a9ba6fa · a9ba6fa
1 parent 04cda74
commit a9ba6fa
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 28 deletions.
diff --git a/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/ErpcTls.java b/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/ErpcTls.java
@@ -10,7 +10,9 @@
 import com.sun.jna.Pointer;
 import datadog.cws.erpc.Erpc;
 import datadog.cws.erpc.Request;
+import datadog.trace.api.DD128bTraceId;
 import datadog.trace.api.DDTraceId;
+import java.math.BigInteger;
 
 /**
  * This class is as a thread local storage.
@@ -21,7 +23,8 @@
 public class ErpcTls implements Tls {
   public static final byte REGISTER_SPAN_TLS_OP = 6;
   public static final long TLS_FORMAT = 0;
-  static final long ENTRY_SIZE = Native.LONG_SIZE * 2;
+  static final int ID_SIZE = 16; // 128 bits
+  static final long ENTRY_SIZE = ID_SIZE * 2; // 2 x 128 bits
 
   // Thread local storage
   private Pointer tls;
@@ -134,33 +137,36 @@ private long getSpanIdOffset(int threadId) {
   }
 
   private long getTraceIdOffset(int threadId) {
-    return getSpanIdOffset(threadId) + Native.LONG_SIZE;
+    return getSpanIdOffset(threadId) + ID_SIZE;
   }
 
-  public void registerSpan(int threadId, DDTraceId traceId, long spanId) {
+  public void registerSpan(int threadId, DDTraceId traceId, BigInteger spanId) {
     long spanIdOffset = getSpanIdOffset(threadId);
     long traceIdOffset = getTraceIdOffset(threadId);
 
-    tls.setLong(spanIdOffset, spanId);
-    tls.setLong(traceIdOffset, traceId.toLong());
+    tls.write(spanIdOffset, spanId.toByteArray(), 0, ID_SIZE);
+    tls.setLong(traceIdOffset, traceId.toLong()); // low bits
+    tls.setLong(traceIdOffset + 8, traceId.toHighOrderLong()); // high bits
   }
 
-  public void registerSpan(DDTraceId traceId, long spanId) {
+  public void registerSpan(DDTraceId traceId, BigInteger spanId) {
     registerSpan(getTID(), traceId, spanId);
   }
 
-  public long getSpanId(int threadId) {
+  public BigInteger getSpanId(int threadId) {
     long offset = getSpanIdOffset(threadId);
-    return tls.getLong(offset);
+    return new BigInteger(tls.getByteArray(offset, ID_SIZE));
   }
 
-  public long getSpanId() {
+  public BigInteger getSpanId() {
     return getSpanId(getTID());
   }
 
   public DDTraceId getTraceId(int threadId) {
     long offset = getTraceIdOffset(threadId);
-    return DDTraceId.from(tls.getLong(offset));
+    long lowBits = tls.getLong(offset);
+    long highBits = tls.getLong(offset + 8);
+    return DD128bTraceId.from(highBits, lowBits);
   }
 
   public DDTraceId getTraceId() {

diff --git a/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/NoTls.java b/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/NoTls.java
@@ -1,17 +1,18 @@
 package datadog.cws.tls;
 
-import datadog.trace.api.DDSpanId;
+import datadog.trace.api.DD128bTraceId;
 import datadog.trace.api.DDTraceId;
+import java.math.BigInteger;
 
 public class NoTls implements Tls {
 
-  public void registerSpan(DDTraceId traceId, long spanId) {}
+  public void registerSpan(DDTraceId traceId, BigInteger spanId) {}
 
-  public long getSpanId() {
-    return DDSpanId.ZERO;
+  public BigInteger getSpanId() {
+    return BigInteger.ZERO;
   }
 
   public DDTraceId getTraceId() {
-    return DDTraceId.ZERO;
+    return DD128bTraceId.ZERO;
   }
 }
diff --git a/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/Span.java b/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/Span.java
@@ -1,13 +1,14 @@
 package datadog.cws.tls;
 
 import datadog.trace.api.DDTraceId;
+import java.math.BigInteger;
 
 final class Span {
 
   private final DDTraceId traceId;
-  private final long spanId;
+  private final BigInteger spanId;
 
-  Span(DDTraceId traceId, long spanId) {
+  Span(DDTraceId traceId, BigInteger spanId) {
     this.traceId = traceId;
     this.spanId = spanId;
   }
@@ -16,7 +17,7 @@ DDTraceId getTraceId() {
     return traceId;
   }
 
-  long getSpanId() {
+  BigInteger getSpanId() {
     return spanId;
   }
 }
diff --git a/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/Tls.java b/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/Tls.java
@@ -1,12 +1,13 @@
 package datadog.cws.tls;
 
 import datadog.trace.api.DDTraceId;
+import java.math.BigInteger;
 
 public interface Tls {
 
-  public void registerSpan(DDTraceId traceId, long spanId);
+  public void registerSpan(DDTraceId traceId, BigInteger spanId);
 
-  public long getSpanId();
+  public BigInteger getSpanId();
 
   public DDTraceId getTraceId();
 }
diff --git a/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/TlsScopeListener.java b/dd-java-agent/cws-tls/src/main/java/datadog/cws/tls/TlsScopeListener.java
@@ -1,9 +1,11 @@
 package datadog.cws.tls;
 
+import datadog.trace.api.DD128bTraceId;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
 import datadog.trace.api.scopemanager.ExtendedScopeListener;
+import java.math.BigInteger;
 import java.util.ArrayDeque;
 import java.util.Deque;
 
@@ -23,15 +25,18 @@ public TlsScopeListener(Tls tls) {
 
   void push(DDTraceId traceId, long spanId) {
     Deque<Span> stack = spanStack.get();
+    BigInteger spanId128b = BigInteger.valueOf(spanId);
 
     Span top = stack.peek();
-    if (top == null || top.getTraceId().toLong() != traceId.toLong() || top.getSpanId() != spanId) {
-
-      Span span = new Span(traceId, spanId);
+    if (top == null
+        || top.getTraceId().toLong() != traceId.toLong()
+        || top.getTraceId().toHighOrderLong() != traceId.toHighOrderLong()
+        || top.getSpanId().longValue() != spanId) {
+      Span span = new Span(traceId, spanId128b);
       stack.push(span);
     }
 
-    tls.registerSpan(traceId, spanId);
+    tls.registerSpan(traceId, spanId128b);
   }
 
   void poll() {
@@ -45,7 +50,7 @@ void poll() {
         return;
       }
     }
-    tls.registerSpan(DDTraceId.ZERO, DDSpanId.ZERO);
+    tls.registerSpan(DD128bTraceId.ZERO, BigInteger.ZERO);
   }
 
   @Override

diff --git a/dd-java-agent/cws-tls/src/test/java/datadog/cws/tls/DummyTls.java b/dd-java-agent/cws-tls/src/test/java/datadog/cws/tls/DummyTls.java
@@ -1,20 +1,21 @@
 package datadog.cws.tls;
 
 import datadog.trace.api.DDTraceId;
+import java.math.BigInteger;
 
 class DummyTls implements Tls {
 
   private DDTraceId traceId;
-  private long spanId;
+  private BigInteger spanId;
 
   @Override
-  public void registerSpan(DDTraceId traceId, long spanId) {
+  public void registerSpan(DDTraceId traceId, BigInteger spanId) {
     this.traceId = traceId;
     this.spanId = spanId;
   }
 
   @Override
-  public long getSpanId() {
+  public BigInteger getSpanId() {
     return spanId;
   }