feat: add setTraceId and clearTraceId to ProfilerApi (#313)

* feat: add setTraceId and clearTraceId to ProfilerApi

* refactor: route setTraceId through the native AsyncProfiler API

* chore: bump async-profiler to 4.4.0.0 (adds native setTraceId)

* perf: parse trace_id hex without allocating substrings

* test: cover trace_id hex parser cases

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

* fix: validate trace_id length up front and fail with a clear exception

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: marcsanmi

Makefile

@@ -1,4 +1,4 @@
-GRAFANA_PYROSCOPE_VERSION := 4.3.0.1
+GRAFANA_PYROSCOPE_VERSION := 4.4.0.0
 
 .PHONY: download-async-profiler
 download-async-profiler:

agent/src/main/java/io/pyroscope/javaagent/ProfilerSdk.java

@@ -46,4 +46,40 @@ public void setTracingContext(long spanId, long spanName) {
     public long registerConstant(String constant) {
         return Pyroscope.LabelsWrapper.registerConstant(constant);
     }
+
+    @Override
+    public void setTraceId(@NotNull String traceId) {
+        // W3C trace ID is 32 hex chars (128 bits). Parse directly into two longs
+        // to avoid the String#substring allocations on this hot path.
+        if (traceId.length() != 32) {
+            throw new NumberFormatException("trace_id must be 32 hex chars, got length " + traceId.length());
+        }
+        long hi = parseHex64(traceId, 0);
+        long lo = parseHex64(traceId, 16);
+        asprof.setTraceId(hi, lo);
+    }
+
+    @Override
+    public void clearTraceId() {
+        asprof.setTraceId(0L, 0L);
+    }
+
+    static long parseHex64(String s, int offset) {
+        long result = 0L;
+        for (int i = 0; i < 16; i++) {
+            int c = s.charAt(offset + i);
+            int nibble;
+            if (c >= '0' && c <= '9') {
+                nibble = c - '0';
+            } else if (c >= 'a' && c <= 'f') {
+                nibble = c - 'a' + 10;
+            } else if (c >= 'A' && c <= 'F') {
+                nibble = c - 'A' + 10;
+            } else {
+                throw new NumberFormatException("invalid hex char in trace_id at index " + (offset + i));
+            }
+            result = (result << 4) | nibble;
+        }
+        return result;
+    }
 }

agent/src/main/java/io/pyroscope/javaagent/api/ProfilerApi.java

@@ -23,4 +23,15 @@ public interface ProfilerApi {
     void setTracingContext(long spanId, long spanName);
 
     long registerConstant(String constant);
+
+    /**
+     * Attaches a {@code trace_id} label to samples produced by the current thread.
+     * Thread local; pair with {@link #clearTraceId()}.
+     */
+    // Default methods so an older ProfilerApi implementation can run against a newer ProfilerApi
+    // injected into the bootstrap classloader (no label rather than AbstractMethodError).
+    // Note: this does not help if the runtime ProfilerApi itself is older and lacks this method.
+    default void setTraceId(@NotNull String traceId) {}
+
+    default void clearTraceId() {}
 }

agent/src/test/java/io/pyroscope/javaagent/ProfilerSdkTest.java

@@ -0,0 +1,54 @@
+package io.pyroscope.javaagent;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class ProfilerSdkTest {
+
+    @Test
+    void parseHex64_lowercase() {
+        // 0x0123456789abcdef
+        assertEquals(0x0123456789abcdefL, ProfilerSdk.parseHex64("0123456789abcdef", 0));
+    }
+
+    @Test
+    void parseHex64_uppercase() {
+        // OTel emits lowercase per W3C spec, but defensive: accept uppercase too.
+        assertEquals(0x0123456789ABCDEFL, ProfilerSdk.parseHex64("0123456789ABCDEF", 0));
+    }
+
+    @Test
+    void parseHex64_allZeros() {
+        assertEquals(0L, ProfilerSdk.parseHex64("0000000000000000", 0));
+    }
+
+    @Test
+    void parseHex64_allFs() {
+        assertEquals(-1L, ProfilerSdk.parseHex64("ffffffffffffffff", 0));
+    }
+
+    @Test
+    void parseHex64_offset() {
+        // Parse the low half of a full 32 char trace id.
+        String traceId = "0123456789abcdeffedcba9876543210";
+        assertEquals(0x0123456789abcdefL, ProfilerSdk.parseHex64(traceId, 0));
+        assertEquals(0xfedcba9876543210L, ProfilerSdk.parseHex64(traceId, 16));
+    }
+
+    @Test
+    void parseHex64_invalidChar() {
+        assertThrows(NumberFormatException.class, () -> ProfilerSdk.parseHex64("0123456789abcdez", 0));
+    }
+
+    @Test
+    void setTraceId_rejectsWrongLength() {
+        ProfilerSdk sdk = new ProfilerSdk();
+        // Both shorter and longer than 32 must throw the same exception type.
+        assertThrows(NumberFormatException.class, () -> sdk.setTraceId(""));
+        assertThrows(NumberFormatException.class, () -> sdk.setTraceId("0123456789abcdef"));
+        assertThrows(NumberFormatException.class, () ->
+                sdk.setTraceId("0123456789abcdef0123456789abcdef00"));
+    }
+}

async-profiler-grafana-fork-dist/lib/async-profiler.jar

@@ -1 +1 @@
-8327292d14d9513a296a208c022adf01a1b0ca3c
+0cb3705aa2c9632dc2a0bb3344ce8e6532bc0f78

async-profiler-grafana-fork-dist/lib/libasyncProfiler-linux-arm64.so

@@ -1 +1 @@
-9c277b06ac6ffc80c9225dbef119f1e2e82f829b
+84373075b3ab8f713c0f265220da60d96d363296