Fix app context attributes wiped by ddprof setContext on span activation

dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfiler.java

@@ -44,6 +44,7 @@
 import datadog.trace.bootstrap.instrumentation.api.TaskWrapper;
 import datadog.trace.util.TempLocationManager;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.attribute.PosixFilePermissions;
@@ -108,6 +109,45 @@ public static DatadogProfiler newInstance(ConfigProvider configProvider) {
 
   private final List<String> orderedContextAttributes;
 
+  // True for each attribute slot that was configured by the application (e.g. foo, bar).
+  // ddprof wipes all custom slots on setContext; these slots are re-applied via
+  // reapplyAppContext() on span activation.
+  private final boolean[] isAppOffset;
+
+  private final boolean hasAppContext;
+
+  /**
+   * Per-thread snapshot of application attribute values. Lazily allocated; only threads that call
+   * setContextValue for an app attribute ever allocate. Holds the ddprof constant ID and
+   * pre-encoded UTF-8 bytes for each slot, ready for a zero-allocation reapply via
+   * setContextValuesByIdAndBytes.
+   */
+  private final ThreadLocal<AppContextSnapshot> appContextValues = new ThreadLocal<>();
+
+  private static final class AppContextSnapshot {
+    final int[] ids;
+    final byte[][] utf8;
+    // Cached string values for change detection — avoids re-encoding and re-snapshotting
+    // the constant ID when the same value is set again on the same thread.
+    final String[] strings;
+    // Scratch buffer for snapshotTags; reused to avoid per-call allocation.
+    final int[] scratch;
+
+    AppContextSnapshot(int size) {
+      ids = new int[size];
+      utf8 = new byte[size][];
+      strings = new String[size];
+      scratch = new int[size];
+    }
+
+    boolean isEmpty() {
+      for (int id : ids) {
+        if (id != 0) return false;
+      }
+      return true;
+    }
+  }
+
   private final long queueTimeThresholdMillis;
 
   private final Path recordingsPath;
@@ -151,6 +191,19 @@ private DatadogProfiler(ConfigProvider configProvider) {
       orderedContextAttributes.add(RESOURCE);
     }
     this.contextSetter = new ContextSetter(profiler, orderedContextAttributes);
+    // ContextSetter deduplicates and truncates to 10 internally; size arrays to its actual size.
+    int contextSize = contextSetter.snapshotTags().length;
+    boolean[] appOffsets = new boolean[contextSize];
+    boolean anyApp = false;
+    for (String attribute : contextAttributes) {
+      int idx = contextSetter.offsetOf(attribute);
+      if (idx >= 0) {
+        appOffsets[idx] = true;
+        anyApp = true;
+      }
+    }
+    this.isAppOffset = appOffsets;
+    this.hasAppContext = anyApp;
     this.queueTimeThresholdMillis =
         configProvider.getLong(
             PROFILING_QUEUEING_TIME_THRESHOLD_MILLIS,
@@ -375,9 +428,13 @@ public void clearSpanContext() {
   public boolean setContextValue(int offset, CharSequence value) {
     if (contextSetter != null && offset >= 0 && value != null) {
       try {
-        return contextSetter.setContextValue(offset, value.toString());
+        String s = value.toString();
+        if (contextSetter.setContextValue(offset, s)) {
+          recordAppContextValue(offset, s);
+          return true;
+        }
       } catch (Throwable e) {
-        log.debug("Failed to set context", e);
+        log.debug("Failed to set context value", e);
       }
     }
     return false;
@@ -400,14 +457,70 @@ public boolean clearContextValue(String attribute) {
   public boolean clearContextValue(int offset) {
     if (contextSetter != null && offset >= 0) {
       try {
+        recordAppContextValue(offset, null);
         return contextSetter.clearContextValue(offset);
       } catch (Throwable t) {
-        log.debug("Failed to clear context", t);
+        log.debug("Failed to clear context value", t);
       }
     }
     return false;
   }
 
+  /**
+   * Re-applies this thread's application-managed context attributes after a span activation.
+   * ddprof's {@code setContext} clears all custom attribute slots; this restores only the
+   * app-owned ones so they remain visible during the new span's lifetime. No-op when no
+   * application attributes are configured or none have been set on this thread.
+   *
+   * <p>Uses the pre-computed constant IDs and UTF-8 bytes captured by {@link
+   * #recordAppContextValue} to restore both the DD sidecar encoding and the OTEP attrs_data value
+   * in a single call — no String allocation, no hash lookup, no OTEP sidecar gap.
+   */
+  public void reapplyAppContext() {
+    if (!hasAppContext) {
+      return;
+    }
+    AppContextSnapshot snapshot = appContextValues.get();
+    if (snapshot == null) {
+      return;
+    }
+    try {
+      contextSetter.setContextValuesByIdAndBytes(snapshot.ids, snapshot.utf8);
+    } catch (Throwable e) {
+      log.debug("Failed to reapply context", e);
+    }
+  }
+
+  private void recordAppContextValue(int offset, String value) {
+    if (!hasAppContext || offset < 0 || offset >= isAppOffset.length || !isAppOffset[offset]) {
+      return;
+    }
+    AppContextSnapshot snapshot = appContextValues.get();
+    if (value == null) {
+      if (snapshot == null) {
+        return;
+      }
+      snapshot.ids[offset] = 0;
+      snapshot.utf8[offset] = null;
+      snapshot.strings[offset] = null;
+      if (snapshot.isEmpty()) {
+        appContextValues.remove();
+      }
+      return;
+    }
+    if (snapshot == null) {
+      snapshot = new AppContextSnapshot(isAppOffset.length);
+      appContextValues.set(snapshot);
+    }
+    if (!value.equals(snapshot.strings[offset])) {
+      snapshot.strings[offset] = value;
+      snapshot.utf8[offset] = value.getBytes(StandardCharsets.UTF_8);
+      // Read back the constant ID registered by setContextValue for this slot only.
+      contextSetter.snapshotTags(snapshot.scratch);
+      snapshot.ids[offset] = snapshot.scratch[offset];
+    }
+  }
+
   private void debugLogging(long localRootSpanId) {
     if (detailedDebugLogging && log.isDebugEnabled()) {
       log.debug("localRootSpanId={}", localRootSpanId, new Throwable());

dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java

@@ -48,6 +48,7 @@ public void activate(Object context) {
                 profilerContext.getTraceIdLow());
             DDPROF.setContextValue(SPAN_NAME_INDEX, profilerContext.getOperationName());
             DDPROF.setContextValue(RESOURCE_NAME_INDEX, profilerContext.getResourceName());
+            DDPROF.reapplyAppContext();
           }
         }
       };

dd-java-agent/agent-profiling/profiling-ddprof/src/test/java/com/datadog/profiling/ddprof/DatadogProfilerTest.java

@@ -1,7 +1,9 @@
 package com.datadog.profiling.ddprof;
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -193,6 +195,91 @@ public void testContextRegistration() {
         assertFalse(Arrays.equals(snapshot2, profiler.snapshot()));
       }
     }
+
+    // setSpanContext wipes all custom slots; reapplyAppContext must restore them.
+    int fooOffset = profiler.offsetOf("foo");
+    fooSetter.set("reapply-me");
+    assertNotEquals(0, profiler.snapshot()[fooOffset]);
+
+    profiler.setSpanContext(1L, 1L, 0L, 1L);
+    assertEquals(0, profiler.snapshot()[fooOffset]);
+
+    profiler.reapplyAppContext();
+    assertNotEquals(0, profiler.snapshot()[fooOffset]);
+
+  }
+
+  @Test
+  void testReapplyAppContextIsIdempotent() {
+    DatadogProfiler profiler =
+        new DatadogProfiler(
+            configProvider(true, true, true, true), new HashSet<>(Arrays.asList("foo")));
+    DatadogProfilerContextSetter fooSetter = new DatadogProfilerContextSetter("foo", profiler);
+    int fooOffset = profiler.offsetOf("foo");
+
+    fooSetter.set("idempotent-value");
+    profiler.setSpanContext(1L, 1L, 0L, 1L);
+    assertEquals(0, profiler.snapshot()[fooOffset], "setSpanContext must wipe foo");
+
+    profiler.reapplyAppContext();
+    int afterFirst = profiler.snapshot()[fooOffset];
+    assertNotEquals(0, afterFirst, "reapplyAppContext must restore foo");
+
+    profiler.reapplyAppContext();
+    assertEquals(afterFirst, profiler.snapshot()[fooOffset],
+        "calling reapplyAppContext twice must produce the same result");
+  }
+
+  @Test
+  void testReactivationAfterChildReappliesAppContext() {
+    DatadogProfiler profiler =
+        new DatadogProfiler(
+            configProvider(true, true, true, true), new HashSet<>(Arrays.asList("foo")));
+    DatadogProfilerContextSetter fooSetter = new DatadogProfilerContextSetter("foo", profiler);
+    int fooOffset = profiler.offsetOf("foo");
+
+    fooSetter.set("parent-foo");
+    int parentEncoding = profiler.snapshot()[fooOffset];
+    assertNotEquals(0, parentEncoding, "foo must be set before child activation");
+
+    // Simulate child activation — wipes attrs.
+    profiler.setSpanContext(2L, 2L, 0L, 2L);
+    assertEquals(0, profiler.snapshot()[fooOffset], "setSpanContext must wipe foo on child activation");
+
+    // Simulate parent re-activation (ScopeStack.cleanup path).
+    profiler.setSpanContext(1L, 1L, 0L, 1L);
+    profiler.reapplyAppContext();
+    assertEquals(parentEncoding, profiler.snapshot()[fooOffset],
+        "re-activation + reapply must restore parent app attr");
+  }
+
+  @Test
+  void testAppAttrSetInChildSurvivesToNextActivation() {
+    DatadogProfiler profiler =
+        new DatadogProfiler(
+            configProvider(true, true, true, true), new HashSet<>(Arrays.asList("foo")));
+    DatadogProfilerContextSetter fooSetter = new DatadogProfilerContextSetter("foo", profiler);
+    int fooOffset = profiler.offsetOf("foo");
+
+    fooSetter.set("parent-val");
+
+    // Child activation wipes foo.
+    profiler.setSpanContext(2L, 2L, 0L, 2L);
+    assertEquals(0, profiler.snapshot()[fooOffset], "setSpanContext must wipe foo");
+
+    // Set foo in child context via the ThreadLocal setter.
+    fooSetter.set("child-val");
+    int childEncoding = profiler.snapshot()[fooOffset];
+    assertNotEquals(0, childEncoding, "foo must be set in child context");
+
+    // Parent re-activation wipes foo again.
+    profiler.setSpanContext(1L, 1L, 0L, 1L);
+    assertEquals(0, profiler.snapshot()[fooOffset], "setSpanContext must wipe foo on re-activation");
+
+    // reapplyAppContext must restore the ThreadLocal ambient value ("child-val").
+    profiler.reapplyAppContext();
+    assertEquals(childEncoding, profiler.snapshot()[fooOffset],
+        "ThreadLocal ambient value must survive into the next activation");
   }
 
   private static ConfigProvider configProvider(

gradle/libs.versions.toml

@@ -8,7 +8,7 @@ gradle-tooling-api = "8.14.5"
 spotbugs_annotations = "4.9.8"
 
 # DataDog libs and forks
-ddprof = "1.44.0"
+ddprof = "1.45.0-SNAPSHOT"
 dogstatsd = "4.4.5"
 okhttp = "3.12.15" # Datadog fork to support Java 7