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

ddprof-lib/src/main/cpp/sframe.cpp

@@ -290,7 +290,7 @@ bool SFrameParser::parse() {
         const SFrameFDE* fde = &fde_array[i];
         if (SFRAME_FUNC_FDE_TYPE(fde->info) != 0) continue;  // skip PCMASK
         if (fde->fre_num == 0)                    continue;  // empty FDE
-        int saved_count             = _count;
+        int saved_count             = _count;  // safe: addRecord capacity guard keeps _count <= INT_MAX/2
         int saved_linked_frame_size = _linked_frame_size;
         if (!parseFDE(hdr, fde, fre_section, fre_end)) {
             if (_oom) return false;  // OOM: partial table is not safe to use

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

@@ -1,21 +1,36 @@
+/*
+ * Copyright 2026, Datadog, Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 package com.datadoghq.profiler;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 
 public class ContextSetter {
 
-    private static final int TAGS_STORAGE_LIMIT = 10;
     private final List<String> attributes;
     private final JavaProfiler profiler;
 
     public ContextSetter(JavaProfiler profiler, List<String> attributes) {
         this.profiler = profiler;
         Set<String> unique = new HashSet<>(attributes);
         this.attributes = new ArrayList<>(unique.size());
-        for (int i = 0; i < Math.min(attributes.size(), TAGS_STORAGE_LIMIT); i++) {
+        for (int i = 0; i < Math.min(attributes.size(), ThreadContext.MAX_CUSTOM_SLOTS); i++) {
             String attribute = attributes.get(i);
             if (unique.remove(attribute)) {
                 this.attributes.add(attribute);
@@ -25,13 +40,21 @@ public ContextSetter(JavaProfiler profiler, List<String> attributes) {
 
     public int[] snapshotTags() {
         int[] snapshot = new int[attributes.size()];
-        snapshotTags(snapshot);
+        profiler.copyTags(snapshot);
         return snapshot;
     }
 
+    /**
+     * Copies current sidecar encodings into {@code snapshot}. The array must have at least
+     * {@code attributes.size()} elements; arrays shorter than {@code attributes.size()} are
+     * silently ignored. Indices {@code [attributes.size(), snapshot.length)} are zeroed after
+     * copying to prevent stale data from leaking to the caller.
+     * Use the no-arg {@link #snapshotTags()} overload to obtain a correctly sized array.
+     */
     public void snapshotTags(int[] snapshot) {
-        if (snapshot.length <= attributes.size()) {
+        if (snapshot.length >= attributes.size()) {
             profiler.copyTags(snapshot);
+            Arrays.fill(snapshot, attributes.size(), snapshot.length, 0);
         }
     }
 
@@ -50,6 +73,22 @@ 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.
+     *
+     * <p><b>Partial-write on overflow.</b> A {@code false} return does not mean the record is
+     * unchanged: slots that were written before an attrs_data overflow remain published. Overflowed
+     * slots are zeroed in both the sidecar and attrs_data views. Callers must not assume the record
+     * is unmodified when {@code false} is returned.
+     */
+    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

@@ -1,5 +1,6 @@
 /*
  * Copyright 2018 Andrei Pangin
+ * Copyright 2026, Datadog, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -213,14 +214,58 @@ public void clearContext() {
         tlsContextStorage.get().put(0, 0, 0, 0);
     }
 
+    /**
+     * Sets a custom context attribute at the given slot offset for the current thread.
+     *
+     * @param offset slot index (0-based, in [0, 9]); out-of-range values return {@code false}
+     * @param value  the string value to record; {@code null} returns {@code false} without
+     *               writing; an empty string is written as a zero-length entry (not a clear —
+     *               use {@link #clearContextAttribute(int)} to remove a value)
+     * @return true if the value was recorded; false if {@code offset} is out of range,
+     *         {@code value} is null, the Dictionary is full, or {@code attrs_data} overflows
+     *         for this slot
+     */
     public boolean setContextAttribute(int offset, String value) {
         return tlsContextStorage.get().setContextAttribute(offset, value);
     }
 
+    /**
+     * Clears the custom context attribute at the given slot offset for the current thread.
+     * Zeros the sidecar encoding and removes it from OTEP {@code attrs_data}.
+     *
+     * @param offset slot index (0-based, in [0, 9]); out-of-range values are silently ignored
+     */
     public void clearContextAttribute(int offset) {
         tlsContextStorage.get().clearContextAttribute(offset);
     }
 
+    /**
+     * Re-applies multiple custom attributes from precomputed constant IDs and UTF-8 bytes for
+     * the current thread in a single detach/attach window.
+     *
+     * <ul>
+     *   <li>Slots with {@code constantIds[i] <= 0} are skipped.</li>
+     *   <li>Returns {@code false} without writing if the thread's record is not currently valid
+     *       (span-less), to avoid resurrecting a cleared record.</li>
+     *   <li>On {@code attrs_data} overflow, the overflowed slot's sidecar is zeroed and
+     *       {@code false} is returned; slots written before the overflow are retained.</li>
+     * </ul>
+     *
+     * @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 255 bytes
+     *                    (the OTEP attrs_data entry length field is one byte) for every slot
+     *                    whose {@code constantId > 0}
+     * @return true if every slot with {@code constantId > 0} was written; false on a cleared
+     *         (span-less) record, or {@code attrs_data} overflow for any slot
+     * @throws NullPointerException     if {@code constantIds}, {@code utf8}, or any active
+     *                                  {@code utf8[i]} is null
+     * @throws IllegalArgumentException if the arrays have different lengths, exceed the slot limit,
+     *                                  or any active {@code utf8[i]} exceeds 255 bytes
+     */
+    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

@@ -19,6 +19,7 @@
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.charset.StandardCharsets;
+import java.util.Objects;
 
 /**
  * Thread-local context for trace/span identification.
@@ -28,7 +29,7 @@
  * for minimal overhead. Only little-endian platforms are supported.
  */
 public final class ThreadContext {
-    private static final int MAX_CUSTOM_SLOTS = 10;
+    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
     // so replaceOtepAttribute can assume the input always fits.
@@ -313,17 +314,101 @@ 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 successfully; false if
+     *         the record was not valid before the call (nothing is published), or if any slot
+     *         overflowed {@code attrs_data} (that slot's sidecar is zeroed). Note: a {@code false}
+     *         return due to {@code attrs_data} overflow does <em>not</em> mean the record is
+     *         unmodified — slots processed before the overflowed one are durably written.
+     * @throws NullPointerException     if {@code constantIds}, {@code utf8}, or any active
+     *                                  {@code utf8[i]} (where {@code constantIds[i] > 0}) is null
+     * @throws IllegalArgumentException if the arrays have different lengths,
+     *                                  {@code constantIds.length > MAX_CUSTOM_SLOTS}, or any
+     *                                  active {@code utf8[i].length > MAX_VALUE_BYTES}
+     */
+    public boolean setContextAttributesByIdAndBytes(int[] constantIds, byte[][] utf8) {
+        Objects.requireNonNull(constantIds, "constantIds");
+        Objects.requireNonNull(utf8, "utf8");
+        if (constantIds.length != utf8.length) {
+            throw new IllegalArgumentException("constantIds and utf8 must have the same length");
+        }
+        if (constantIds.length > MAX_CUSTOM_SLOTS) {
+            throw new IllegalArgumentException("constantIds.length exceeds MAX_CUSTOM_SLOTS");
+        }
+        int len = constantIds.length;
+        // Validate active slots before touching the buffer so a bad input never leaves
+        // the record detached (valid=0) after an exception unwinds past attach().
+        for (int i = 0; i < len; i++) {
+            if (constantIds[i] > 0) {
+                byte[] bytes = Objects.requireNonNull(utf8[i], "utf8[" + i + "]");
+                if (bytes.length > MAX_VALUE_BYTES) {
+                    throw new IllegalArgumentException("utf8[" + i + "].length exceeds MAX_VALUE_BYTES");
+                }
+            }
+        }
+        // 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-stresstest/build.gradle.kts

@@ -1,3 +1,18 @@
+/*
+ * Copyright 2026, Datadog, Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 import com.datadoghq.native.util.PlatformUtils
 
 plugins {
@@ -67,6 +82,9 @@ dependencies {
   // dd-trace-api: annotations only at compile time. The patched dd-java-agent
   // provides the (relocated) runtime classes and intercepts @Trace.
   "chaosCompileOnly"(libs.dd.trace.api)
+  // ddprof-lib public API: compile-only; the patched dd-java-agent provides the
+  // classes at runtime for antagonists that call JavaProfiler/ThreadContext directly.
+  "chaosCompileOnly"(project(mapOf("path" to ":ddprof-lib", "configuration" to "debug")))
 }
 
 tasks.register<Jar>("chaosJar") {

ddprof-stresstest/src/chaos/java/com/datadoghq/profiler/chaos/Main.java

@@ -90,6 +90,8 @@ private static Antagonist create(String name) {
                 return new WeakRefWaveAntagonist();
             case "dump-storm":
                 return new DumpStormAntagonist();
+            case "reapply-context":
+                return new ReapplyContextAntagonist();
             // Deferred: dlopen-churn (needs per-arch dummy .so built in CI prep).
             default:
                 throw new IllegalArgumentException("unknown antagonist: " + name);