feat(PROF-15098): add batch reapply-by-id-and-bytes context API

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: jbachorik

ddprof-lib/src/main/java/com/datadoghq/profiler/ContextSetter.java

@@ -7,6 +7,7 @@
 
 public class ContextSetter {
 
+    // Must equal ThreadContext.MAX_CUSTOM_SLOTS — both cap the same physical slot range.
     private static final int TAGS_STORAGE_LIMIT = 10;
     private final List<String> attributes;
     private final JavaProfiler profiler;
@@ -30,7 +31,7 @@ public int[] snapshotTags() {
     }
 
     public void snapshotTags(int[] snapshot) {
-        if (snapshot.length <= attributes.size()) {
+        if (snapshot.length >= attributes.size()) {
             profiler.copyTags(snapshot);
         }
     }
@@ -50,6 +51,17 @@ public boolean setContextValue(int offset, String value) {
         return false;
     }
 
+    /**
+     * Re-applies attribute values from precomputed constant IDs and UTF-8 bytes, indexed by slot
+     * (as produced by {@link #snapshotTags(int[])}). Restores both the DD sidecar encoding and the
+     * OTEP attrs_data value for every slot whose constantId is {@code > 0}, in a single atomic
+     * publish — no String allocation, hashing, or cache lookup. Intended for re-applying
+     * application-managed context after a {@code setContext} span activation wipes the slots.
+     */
+    public boolean setContextValuesByIdAndBytes(int[] constantIds, byte[][] utf8) {
+        return profiler.setContextAttributesByIdAndBytes(constantIds, utf8);
+    }
+
     public boolean clearContextValue(String attribute) {
         return clearContextValue(offsetOf(attribute));
     }

ddprof-lib/src/main/java/com/datadoghq/profiler/JavaProfiler.java

@@ -221,6 +221,10 @@ public void clearContextAttribute(int offset) {
         tlsContextStorage.get().clearContextAttribute(offset);
     }
 
+    public boolean setContextAttributesByIdAndBytes(int[] constantIds, byte[][] utf8) {
+        return tlsContextStorage.get().setContextAttributesByIdAndBytes(constantIds, utf8);
+    }
+
     void copyTags(int[] snapshot) {
         tlsContextStorage.get().copyCustoms(snapshot);
     }

ddprof-lib/src/main/java/com/datadoghq/profiler/ThreadContext.java

@@ -28,6 +28,7 @@
  * for minimal overhead. Only little-endian platforms are supported.
  */
 public final class ThreadContext {
+    // Must equal ContextSetter.TAGS_STORAGE_LIMIT — both cap the same physical slot range.
     private static final int MAX_CUSTOM_SLOTS = 10;
     // Max UTF-8 byte length for a custom attribute value. Matches the 1-byte length
     // field in the OTEP attrs_data entry header. Enforced up front in setContextAttribute
@@ -313,17 +314,89 @@ private boolean setContextAttributeDirect(int keyIndex, String value) {
 
         // Write both sidecar and OTEP attrs_data inside the detach/attach window
         // so a signal handler never sees a new sidecar encoding alongside old attrs_data.
-        int otepKeyIndex = keyIndex + 1;
         detach();
+        boolean written = writeSlot(keyIndex, encoding, utf8);
+        attach();
+        return written;
+    }
+
+    /**
+     * Writes one slot's sidecar encoding and OTEP attrs_data value. The caller must already hold
+     * the detach/attach window. On attrs_data overflow the old entry was compacted out and the new
+     * one couldn't fit, so the sidecar is zeroed to keep both views agreeing there is no value.
+     *
+     * @return true if the value was written; false on attrs_data overflow (sidecar left zeroed)
+     */
+    private boolean writeSlot(int keyIndex, int encoding, byte[] utf8) {
         ctxBuffer.putInt(TAG_ENCODINGS_OFFSET + keyIndex * Integer.BYTES, encoding);
-        boolean written = replaceOtepAttribute(otepKeyIndex, utf8);
-        if (!written) {
-            // attrs_data overflow: the old entry was compacted out and the new one
-            // couldn't fit. Zero the sidecar so both views agree there is no value.
+        if (!replaceOtepAttribute(keyIndex + 1, utf8)) {
             ctxBuffer.putInt(TAG_ENCODINGS_OFFSET + keyIndex * Integer.BYTES, 0);
+            return false;
+        }
+        return true;
+    }
+
+    /**
+     * Re-applies multiple custom attributes from precomputed constant IDs and UTF-8 bytes in a
+     * single detach/attach window, without Dictionary lookups or per-thread cache access.
+     *
+     * <p>Both arrays are indexed by key slot, matching the layout produced by
+     * {@link #copyCustoms(int[])}. For each slot {@code i} with {@code constantIds[i] > 0}, the
+     * sidecar encoding (DD signal handler) and the OTEP attrs_data value (external profilers) are
+     * written together; slots with {@code constantIds[i] <= 0} are left untouched. Performing all
+     * writes inside one detach/attach window means a signal handler never observes a partially
+     * re-applied set.
+     *
+     * <p>Intended for the reapply-app-context hot path: the caller already holds the constant IDs
+     * (from {@link #copyCustoms(int[])}) and the UTF-8 bytes (from the original Strings), so this
+     * does no String allocation, hashing, or cache lookup. The record is re-published only if it
+     * was valid before the call, so a cleared (span-less) record is not resurrected.
+     *
+     * <p><b>Caller contract.</b> {@code constantIds[i]} must be an ID previously returned for the
+     * same value via {@link #copyCustoms(int[])} within the current profiler session, and
+     * {@code utf8[i]} must be that value's UTF-8 bytes. The (id, bytes) pairing is not verifiable
+     * here; a mismatch silently diverges the sidecar and attrs_data views.
+     *
+     * @param constantIds per-slot Dictionary constant IDs; entries {@code <= 0} are skipped
+     * @param utf8        per-slot UTF-8 value bytes; must be non-null and at most
+     *                    {@value #MAX_VALUE_BYTES} bytes for every slot whose constantId {@code > 0}
+     * @return true if every slot with {@code constantId > 0} was written; false on invalid
+     *         arguments (null arrays, length mismatch, or a missing/oversized {@code utf8} entry),
+     *         if the record was not valid before the call (nothing published), or if any slot
+     *         overflowed attrs_data (that slot's sidecar is zeroed)
+     */
+    public boolean setContextAttributesByIdAndBytes(int[] constantIds, byte[][] utf8) {
+        if (constantIds == null || utf8 == null || constantIds.length != utf8.length) {
+            return false;
+        }
+        int len = Math.min(constantIds.length, MAX_CUSTOM_SLOTS);
+        // Validate before mutating so malformed input never leaves a half-written record.
+        for (int i = 0; i < len; i++) {
+            if (constantIds[i] > 0) {
+                byte[] bytes = utf8[i];
+                if (bytes == null || bytes.length > MAX_VALUE_BYTES) {
+                    return false;
+                }
+            }
+        }
+        // Never resurrect a cleared (span-less) record: valid=0 means no reader can observe
+        // what we write, and re-publishing would expose a record with no trace/span context.
+        if (ctxBuffer.get(validOffset) == 0) {
+            return false;
+        }
+        detach();
+        boolean allWritten = true;
+        for (int i = 0; i < len; i++) {
+            int constantId = constantIds[i];
+            if (constantId <= 0) {
+                continue;
+            }
+            if (!writeSlot(i, constantId, utf8[i])) {
+                allWritten = false;
+            }
         }
         attach();
-        return written;
+        return allWritten;
     }
 
     /**

ddprof-test/src/test/java/com/datadoghq/profiler/context/TagContextTest.java

@@ -1,5 +1,6 @@
 package com.datadoghq.profiler.context;
 
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -239,6 +240,172 @@ public void testCrossSlotIsolation() throws Exception {
         assertNull(readTag(contextSetter, "tag2"));
     }
 
+    @Test
+    public void testReapplyByIdAndBytes() throws Exception {
+        Assumptions.assumeTrue(!Platform.isJ9());
+        registerCurrentThreadForWallClockProfiling();
+        ContextSetter contextSetter = new ContextSetter(profiler, Arrays.asList("tag1", "tag2"));
+        int slot = contextSetter.offsetOf("tag1");
+        String value = "app-managed";
+
+        // Set the attribute the normal way, then capture both the constant ID (sidecar) and the
+        // UTF-8 bytes — exactly what dd-trace-java retains for the reapply hot path.
+        assertTrue(contextSetter.setContextValue("tag1", value));
+        int[] ids = contextSetter.snapshotTags();
+        int savedId = ids[slot];
+        assertNotEquals(0, savedId);
+        byte[][] bytes = new byte[ids.length][];
+        bytes[slot] = value.getBytes(StandardCharsets.UTF_8);
+
+        // setContext (span activation) wipes both views.
+        profiler.setContext(1L, 42L, 0L, 42L);
+        assertNull(readTag(contextSetter, "tag1"), "attrs_data must be wiped by setContext");
+        assertEquals(0, contextSetter.snapshotTags()[slot], "sidecar must be wiped by setContext");
+
+        // Reapply by ID + bytes restores BOTH views.
+        assertTrue(contextSetter.setContextValuesByIdAndBytes(ids, bytes));
+        assertEquals(value, readTag(contextSetter, "tag1"), "attrs_data must be restored");
+        assertEquals(savedId, contextSetter.snapshotTags()[slot], "sidecar must be restored");
+    }
+
+    @Test
+    public void testReapplyByIdAndBytesRejectsBadArgs() throws Exception {
+        Assumptions.assumeTrue(!Platform.isJ9());
+        registerCurrentThreadForWallClockProfiling();
+        ContextSetter contextSetter = new ContextSetter(profiler, Arrays.asList("tag1", "tag2"));
+        int slot = contextSetter.offsetOf("tag1");
+
+        assertTrue(contextSetter.setContextValue("tag1", "v"));
+        int id = contextSetter.snapshotTags()[slot];
+        assertNotEquals(0, id);
+
+        // Null arrays and length mismatch.
+        assertFalse(contextSetter.setContextValuesByIdAndBytes(null, new byte[1][]));
+        assertFalse(contextSetter.setContextValuesByIdAndBytes(new int[1], null));
+        assertFalse(contextSetter.setContextValuesByIdAndBytes(new int[2], new byte[3][]));
+
+        // A slot with constantId > 0 requires non-null bytes within the size limit.
+        assertFalse(contextSetter.setContextValuesByIdAndBytes(
+                new int[] {id, 0}, new byte[][] {null, null}));
+        assertFalse(contextSetter.setContextValuesByIdAndBytes(
+                new int[] {id, 0}, new byte[][] {new byte[256], null}));
+
+        // 255 bytes is the boundary and must be accepted.
+        byte[] ok255 = new byte[255];
+        Arrays.fill(ok255, (byte) 'x');
+        assertTrue(contextSetter.setContextValuesByIdAndBytes(
+                new int[] {id, 0}, new byte[][] {ok255, null}));
+    }
+
+    @Test
+    public void testReapplyByIdAndBytesReplacesExistingValue() throws Exception {
+        Assumptions.assumeTrue(!Platform.isJ9());
+        registerCurrentThreadForWallClockProfiling();
+        ContextSetter contextSetter = new ContextSetter(profiler, Arrays.asList("tag1", "tag2"));
+        int slot = contextSetter.offsetOf("tag1");
+
+        // Capture the ID + bytes for "first".
+        assertTrue(contextSetter.setContextValue("tag1", "first"));
+        int[] idsFirst = contextSetter.snapshotTags();
+        byte[][] bytesFirst = new byte[idsFirst.length][];
+        bytesFirst[slot] = "first".getBytes(StandardCharsets.UTF_8);
+
+        // Overwrite the live slot with a different value.
+        assertTrue(contextSetter.setContextValue("tag1", "second"));
+        assertEquals("second", readTag(contextSetter, "tag1"));
+
+        // Reapply "first" by ID + bytes over the live "second" — exercises the
+        // compact-then-insert path in replaceOtepAttribute.
+        assertTrue(contextSetter.setContextValuesByIdAndBytes(idsFirst, bytesFirst));
+        assertEquals("first", readTag(contextSetter, "tag1"));
+        assertEquals(idsFirst[slot], contextSetter.snapshotTags()[slot]);
+    }
+
+    @Test
+    public void testReapplyByIdAndBytesAfterClear() throws Exception {
+        Assumptions.assumeTrue(!Platform.isJ9());
+        registerCurrentThreadForWallClockProfiling();
+        ContextSetter contextSetter = new ContextSetter(profiler, Arrays.asList("tag1", "tag2"));
+        int slot = contextSetter.offsetOf("tag1");
+
+        assertTrue(contextSetter.setContextValue("tag1", "live"));
+        int[] ids = contextSetter.snapshotTags();
+        byte[][] bytes = new byte[ids.length][];
+        bytes[slot] = "live".getBytes(StandardCharsets.UTF_8);
+
+        assertTrue(contextSetter.clearContextValue("tag1"));
+        assertNull(readTag(contextSetter, "tag1"));
+        assertEquals(0, contextSetter.snapshotTags()[slot]);
+
+        assertTrue(contextSetter.setContextValuesByIdAndBytes(ids, bytes));
+        assertEquals("live", readTag(contextSetter, "tag1"));
+        assertEquals(ids[slot], contextSetter.snapshotTags()[slot]);
+    }
+
+    @Test
+    public void testReapplyByIdAndBytesClearedRecord() throws Exception {
+        // Verifies that setContextValuesByIdAndBytes never resurrects a cleared (span-less) record.
+        // A cleared record has valid=0 and no trace/span context; re-publishing it would expose
+        // attribute values with no associated trace, which is meaningless to the signal handler.
+        Assumptions.assumeTrue(!Platform.isJ9());
+        registerCurrentThreadForWallClockProfiling();
+        ContextSetter contextSetter = new ContextSetter(profiler, Arrays.asList("tag1", "tag2"));
+        int slot = contextSetter.offsetOf("tag1");
+
+        // Establish a live record with tag1 set.
+        profiler.setContext(1L, 42L, 0L, 42L);
+        assertTrue(contextSetter.setContextValue("tag1", "will-be-cleared"));
+        int[] ids = contextSetter.snapshotTags();
+        byte[][] bytes = new byte[ids.length][];
+        bytes[slot] = "will-be-cleared".getBytes(StandardCharsets.UTF_8);
+
+        // Drive valid=0 via the all-zero clear path (clearContext → put(0,0,0,0) → no attach()).
+        profiler.clearContext();
+        // readContextAttribute respects valid=0 and returns null, confirming the record is dark.
+        assertNull(readTag(contextSetter, "tag1"));
+
+        // Reapply must return false and must not resurrect the cleared record.
+        assertFalse(contextSetter.setContextValuesByIdAndBytes(ids, bytes),
+                "setContextValuesByIdAndBytes must return false when the record is cleared (valid=0)");
+        assertNull(readTag(contextSetter, "tag1"),
+                "cleared record must not be resurrected by setContextValuesByIdAndBytes");
+    }
+
+    @Test
+    public void testReapplyByIdAndBytesOverflowRollback() throws Exception {
+        Assumptions.assumeTrue(!Platform.isJ9());
+        registerCurrentThreadForWallClockProfiling();
+        List<String> attrs = new ArrayList<>();
+        for (int i = 1; i <= 10; i++) {
+            attrs.add("tag" + i);
+        }
+        ContextSetter contextSetter = new ContextSetter(profiler, attrs);
+
+        // Register one 255-byte value to obtain a valid constant ID.
+        char[] chars = new char[255];
+        Arrays.fill(chars, 'x');
+        String bigValue = new String(chars);
+        assertTrue(contextSetter.setContextValue("tag1", bigValue));
+        int bigId = contextSetter.snapshotTags()[contextSetter.offsetOf("tag1")];
+        assertNotEquals(0, bigId);
+        byte[] bigBytes = bigValue.getBytes(StandardCharsets.UTF_8);
+
+        // Reapply the same 255-byte value to all 10 slots — attrs_data cannot hold them all.
+        int[] ids = new int[10];
+        byte[][] bytes = new byte[10][];
+        Arrays.fill(ids, bigId);
+        Arrays.fill(bytes, bigBytes);
+        assertFalse(contextSetter.setContextValuesByIdAndBytes(ids, bytes),
+                "10 x 255-byte values must overflow attrs_data");
+
+        // The last slot certainly overflowed: its sidecar must be zeroed and attrs_data empty.
+        int lastSlot = contextSetter.offsetOf("tag10");
+        assertEquals(0, contextSetter.snapshotTags()[lastSlot],
+                "overflowed slot's sidecar must be zeroed");
+        assertNull(readTag(contextSetter, "tag10"),
+                "overflowed slot must read null — the entry never landed in attrs_data");
+    }
+
     private void work(ContextSetter contextSetter, String contextAttribute, String contextValue)
             throws InterruptedException {
         assertTrue(contextSetter.setContextValue(contextAttribute, contextValue));