Rework thread-owned tags: encapsulate locking via withLock API

Address reviewer feedback on the thread-ownership optimization:

- Add TagMap.withLock(Runnable/Supplier) API for compound operations,
  replacing all synchronized(unsafeTags) in DDSpanContext with
  unsafeTags.withLock() — locking is now fully encapsulated in TagMap
- Add missing ownership checks to OptimizedTagMap: size(), isEmpty(),
  containsValue(), freeze(), isFrozen(), toString(), immutableCopy()
- Add revokeOwnership() helper to avoid redundant volatile writes on
  the cache line after transition to shared mode
- withLock() revokes ownership before running the operation, closing
  the TOCTOU race where the owner thread could bypass the compound lock
- Add self-synchronization to LegacyTagMap (synchronized on all key
  HashMap methods + withLock override) so it is safe without outer locks
- Add TagMapConcurrencyTest: owner fast path, transition visibility,
  post-transition contention, withLock atomicity, size consistency
- Expand DDSpanContextConcurrencyTest: compound atomicity for
  setSpanSamplingPriority, getTags snapshot consistency, transition +
  concurrent read safety

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

Author: bm1549

dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java

@@ -609,11 +609,14 @@ public int getSamplingPriority() {
   }
 
   public void setSpanSamplingPriority(double rate, int limit) {
-    unsafeSetTag(SPAN_SAMPLING_MECHANISM_TAG, SamplingMechanism.SPAN_SAMPLING_RATE);
-    unsafeSetTag(SPAN_SAMPLING_RULE_RATE_TAG, rate);
-    if (limit != Integer.MAX_VALUE) {
-      unsafeSetTag(SPAN_SAMPLING_MAX_PER_SECOND_TAG, limit);
-    }
+    unsafeTags.withLock(
+        () -> {
+          unsafeSetTag(SPAN_SAMPLING_MECHANISM_TAG, SamplingMechanism.SPAN_SAMPLING_RATE);
+          unsafeSetTag(SPAN_SAMPLING_RULE_RATE_TAG, rate);
+          if (limit != Integer.MAX_VALUE) {
+            unsafeSetTag(SPAN_SAMPLING_MAX_PER_SECOND_TAG, limit);
+          }
+        });
   }
 
   /**
@@ -901,44 +904,50 @@ void setAllTags(final TagMap map, boolean needsIntercept) {
       return;
     }
 
-    if (needsIntercept) {
-      // forEach out-performs the iterator of TagMap
-      // Taking advantage of ability to pass through other context arguments
-      // to avoid using a capturing lambda
-      map.forEach(
-          this,
-          (ctx, tagEntry) -> {
-            String tag = tagEntry.tag();
-            Object value = tagEntry.objectValue();
-
-            if (!ctx.tagInterceptor.interceptTag(ctx, tag, value)) {
-              ctx.unsafeTags.set(tagEntry);
-            }
-          });
-    } else {
-      unsafeTags.putAll(map);
-    }
+    unsafeTags.withLock(
+        () -> {
+          if (needsIntercept) {
+            // forEach out-performs the iterator of TagMap
+            // Taking advantage of ability to pass through other context arguments
+            // to avoid using a capturing lambda
+            map.forEach(
+                this,
+                (ctx, tagEntry) -> {
+                  String tag = tagEntry.tag();
+                  Object value = tagEntry.objectValue();
+
+                  if (!ctx.tagInterceptor.interceptTag(ctx, tag, value)) {
+                    ctx.unsafeTags.set(tagEntry);
+                  }
+                });
+          } else {
+            unsafeTags.putAll(map);
+          }
+        });
   }
 
   void setAllTags(final TagMap.Ledger ledger) {
     if (ledger == null) {
       return;
     }
 
-    for (final TagMap.EntryChange entryChange : ledger) {
-      if (entryChange.isRemoval()) {
-        unsafeTags.remove(entryChange.tag());
-      } else {
-        TagMap.Entry entry = (TagMap.Entry) entryChange;
+    unsafeTags.withLock(
+        () -> {
+          for (final TagMap.EntryChange entryChange : ledger) {
+            if (entryChange.isRemoval()) {
+              unsafeTags.remove(entryChange.tag());
+            } else {
+              TagMap.Entry entry = (TagMap.Entry) entryChange;
 
-        String tag = entry.tag();
-        Object value = entry.objectValue();
+              String tag = entry.tag();
+              Object value = entry.objectValue();
 
-        if (!tagInterceptor.interceptTag(this, tag, value)) {
-          unsafeTags.set(entry);
-        }
-      }
-    }
+              if (!tagInterceptor.interceptTag(this, tag, value)) {
+                unsafeTags.set(entry);
+              }
+            }
+          }
+        });
   }
 
   void setAllTags(final Map<String, ?> map) {
@@ -947,11 +956,14 @@ void setAllTags(final Map<String, ?> map) {
     } else if (map instanceof TagMap) {
       setAllTags((TagMap) map);
     } else if (!map.isEmpty()) {
-      for (final Map.Entry<String, ?> tag : map.entrySet()) {
-        if (!tagInterceptor.interceptTag(this, tag.getKey(), tag.getValue())) {
-          unsafeSetTag(tag.getKey(), tag.getValue());
-        }
-      }
+      unsafeTags.withLock(
+          () -> {
+            for (final Map.Entry<String, ?> tag : map.entrySet()) {
+              if (!tagInterceptor.interceptTag(this, tag.getKey(), tag.getValue())) {
+                unsafeSetTag(tag.getKey(), tag.getValue());
+              }
+            }
+          });
     }
   }
 
@@ -1000,23 +1012,26 @@ public Object unsafeGetTag(final String tag) {
 
   @Deprecated
   public TagMap getTags() {
-    TagMap tags = unsafeTags.copy();
-
-    tags.put(DDTags.THREAD_ID, threadId);
-    // maintain previously observable type of the thread name :|
-    tags.put(DDTags.THREAD_NAME, threadName.toString());
-    if (samplingPriority != PrioritySampling.UNSET) {
-      tags.put(SAMPLE_RATE_KEY, samplingPriority);
-    }
-    if (httpStatusCode != 0) {
-      tags.put(Tags.HTTP_STATUS, (int) httpStatusCode);
-    }
-    // maintain previously observable type of http url :|
-    Object value = tags.get(Tags.HTTP_URL);
-    if (value != null) {
-      tags.put(Tags.HTTP_URL, value.toString());
-    }
-    return tags.freeze();
+    return unsafeTags.withLock(
+        () -> {
+          TagMap tags = unsafeTags.copy();
+
+          tags.put(DDTags.THREAD_ID, threadId);
+          // maintain previously observable type of the thread name :|
+          tags.put(DDTags.THREAD_NAME, threadName.toString());
+          if (samplingPriority != PrioritySampling.UNSET) {
+            tags.put(SAMPLE_RATE_KEY, samplingPriority);
+          }
+          if (httpStatusCode != 0) {
+            tags.put(Tags.HTTP_STATUS, (int) httpStatusCode);
+          }
+          // maintain previously observable type of http url :|
+          Object value = tags.get(Tags.HTTP_URL);
+          if (value != null) {
+            tags.put(Tags.HTTP_URL, value.toString());
+          }
+          return tags.freeze();
+        });
   }
 
   /**
@@ -1048,7 +1063,8 @@ public <T> void setMetaStruct(final String field, final T value) {
   }
 
   void earlyProcessTags(AppendableSpanLinks links) {
-    TagsPostProcessorFactory.eagerProcessor().processTags(unsafeTags, this, links);
+    unsafeTags.withLock(
+        () -> TagsPostProcessorFactory.eagerProcessor().processTags(unsafeTags, this, links));
   }
 
   void processTagsAndBaggage(
@@ -1058,39 +1074,44 @@ void processTagsAndBaggage(
     // This is a compromise to avoid...
     // - creating an extra wrapper object that would create significant allocation
     // - implementing an interface to read the spans that require making the read method public
-    // Tags
-    TagsPostProcessorFactory.lazyProcessor().processTags(unsafeTags, this, restrictedSpan);
-
-    String linksTag = DDSpanLink.toTag(restrictedSpan.getLinks());
-    if (linksTag != null) {
-      unsafeTags.put(SPAN_LINKS, linksTag);
-    }
-    // Baggage
-    Map<String, String> baggageItemsWithPropagationTags;
-    if (injectBaggageAsTags) {
-      baggageItemsWithPropagationTags = new HashMap<>(baggageItems);
-      if (w3cBaggage != null) {
-        injectW3CBaggageTags(baggageItemsWithPropagationTags);
-      }
-      propagationTags.fillTagMap(baggageItemsWithPropagationTags);
-    } else {
-      baggageItemsWithPropagationTags = propagationTags.createTagMap();
-    }
-
-    consumer.accept(
-        new Metadata(
-            threadId,
-            threadName,
-            unsafeTags,
-            baggageItemsWithPropagationTags,
-            samplingPriority != PrioritySampling.UNSET ? samplingPriority : getSamplingPriority(),
-            measured,
-            topLevel,
-            httpStatusCode == 0 ? null : HTTP_STATUSES.get(httpStatusCode),
-            // Get origin from rootSpan.context
-            getOrigin(),
-            longRunningVersion,
-            ProcessTags.getTagsForSerialization()));
+    unsafeTags.withLock(
+        () -> {
+          // Tags
+          TagsPostProcessorFactory.lazyProcessor().processTags(unsafeTags, this, restrictedSpan);
+
+          String linksTag = DDSpanLink.toTag(restrictedSpan.getLinks());
+          if (linksTag != null) {
+            unsafeTags.put(SPAN_LINKS, linksTag);
+          }
+          // Baggage
+          Map<String, String> baggageItemsWithPropagationTags;
+          if (injectBaggageAsTags) {
+            baggageItemsWithPropagationTags = new HashMap<>(baggageItems);
+            if (w3cBaggage != null) {
+              injectW3CBaggageTags(baggageItemsWithPropagationTags);
+            }
+            propagationTags.fillTagMap(baggageItemsWithPropagationTags);
+          } else {
+            baggageItemsWithPropagationTags = propagationTags.createTagMap();
+          }
+
+          consumer.accept(
+              new Metadata(
+                  threadId,
+                  threadName,
+                  unsafeTags,
+                  baggageItemsWithPropagationTags,
+                  samplingPriority != PrioritySampling.UNSET
+                      ? samplingPriority
+                      : getSamplingPriority(),
+                  measured,
+                  topLevel,
+                  httpStatusCode == 0 ? null : HTTP_STATUSES.get(httpStatusCode),
+                  // Get origin from rootSpan.context
+                  getOrigin(),
+                  longRunningVersion,
+                  ProcessTags.getTagsForSerialization()));
+        });
   }
 
   void injectW3CBaggageTags(Map<String, String> baggageItemsWithPropagationTags) {

dd-trace-core/src/test/java/datadog/trace/core/DDSpanContextConcurrencyTest.java

@@ -613,4 +613,145 @@ void concurrentSetMetric() throws Exception {
           "Thread " + t + " last metric missing");
     }
   }
+
+  /**
+   * Verifies compound atomicity of setSpanSamplingPriority: the three sampling tags must always
+   * appear together (all or none) when observed via getTags().
+   */
+  @RepeatedTest(10)
+  @DisplayName("setSpanSamplingPriority atomicity: 3 tags always consistent in getTags()")
+  void spanSamplingPriorityAtomicity() throws Exception {
+    AgentSpan span = TRACER.startSpan("test", "samplingAtomicity");
+    DDSpanContext ctx = ((DDSpan) span).context();
+    // Transition to shared so the compound lock is in effect
+    ctx.transitionToShared();
+
+    int iterations = 2000;
+    AtomicBoolean running = new AtomicBoolean(true);
+    CopyOnWriteArrayList<Throwable> errors = new CopyOnWriteArrayList<>();
+
+    // Writer: repeatedly sets span sampling priority
+    Thread writer =
+        new Thread(
+            () -> {
+              try {
+                for (int i = 0; i < iterations; i++) {
+                  ctx.setSpanSamplingPriority(0.5, 100);
+                }
+              } catch (Throwable e) {
+                errors.add(e);
+              } finally {
+                running.set(false);
+              }
+            });
+
+    // Reader: checks that the 3 sampling tags are consistent
+    Thread reader =
+        new Thread(
+            () -> {
+              try {
+                while (running.get()) {
+                  datadog.trace.api.TagMap tags = ctx.getTags();
+                  Object mechanism = tags.getObject("_dd.span_sampling.mechanism");
+                  Object ruleRate = tags.getObject("_dd.span_sampling.rule_rate");
+                  Object maxPerSecond = tags.getObject("_dd.span_sampling.max_per_second");
+
+                  // Either all three are present or none
+                  if (mechanism != null || ruleRate != null || maxPerSecond != null) {
+                    if (mechanism == null || ruleRate == null || maxPerSecond == null) {
+                      errors.add(
+                          new AssertionError(
+                              "Partial sampling tags: mechanism="
+                                  + mechanism
+                                  + " ruleRate="
+                                  + ruleRate
+                                  + " maxPerSecond="
+                                  + maxPerSecond));
+                    }
+                  }
+                }
+              } catch (Throwable e) {
+                errors.add(e);
+              }
+            });
+
+    writer.start();
+    reader.start();
+    writer.join(30_000);
+    reader.join(5_000);
+    span.finish();
+
+    assertEquals(0, errors.size(), "Atomicity violated: " + errors);
+  }
+
+  /**
+   * Verifies that getTags() returns a consistent snapshot: THREAD_ID and THREAD_NAME are always
+   * present and consistent with the span's owning thread.
+   */
+  @Test
+  @DisplayName("getTags() snapshot consistency: virtual fields always present")
+  void getTagsSnapshotConsistency() throws Exception {
+    AgentSpan span = TRACER.startSpan("test", "getTagsConsistency");
+
+    span.setTag("user_tag", "hello");
+    span.finish();
+
+    datadog.trace.api.TagMap tags = ((DDSpan) span).context().getTags();
+    assertNotNull(tags.getObject("thread.id"), "THREAD_ID missing from getTags()");
+    assertNotNull(tags.getObject("thread.name"), "THREAD_NAME missing from getTags()");
+    assertEquals("hello", tags.getObject("user_tag"));
+  }
+
+  /**
+   * Exercises the transition window: owner tags, finish (triggers transitionToShared), then
+   * concurrent readers verify tags are visible post-transition.
+   */
+  @RepeatedTest(5)
+  @DisplayName("Transition + concurrent read: tags written before finish visible after")
+  void transitionAndConcurrentRead() throws Exception {
+    int numReaders = 4;
+    int tagsToWrite = 200;
+    AgentSpan span = TRACER.startSpan("test", "transitionRead");
+
+    // Owner thread writes tags
+    for (int i = 0; i < tagsToWrite; i++) {
+      span.setTag("pre_" + i, "val_" + i);
+    }
+
+    // Finish triggers transitionToShared
+    span.finish();
+
+    // Multiple reader threads verify all tags are visible
+    CyclicBarrier barrier = new CyclicBarrier(numReaders);
+    ExecutorService executor = Executors.newFixedThreadPool(numReaders);
+    CopyOnWriteArrayList<Throwable> errors = new CopyOnWriteArrayList<>();
+
+    List<Future<?>> futures = new ArrayList<>();
+    for (int r = 0; r < numReaders; r++) {
+      futures.add(
+          executor.submit(
+              () -> {
+                try {
+                  barrier.await(5, TimeUnit.SECONDS);
+                  for (int i = 0; i < tagsToWrite; i++) {
+                    Object val = span.getTag("pre_" + i);
+                    if (!"val_".concat(String.valueOf(i)).equals(val)) {
+                      errors.add(
+                          new AssertionError(
+                              "Tag pre_" + i + " expected val_" + i + " got " + val));
+                    }
+                  }
+                } catch (Throwable e) {
+                  errors.add(e);
+                }
+              }));
+    }
+
+    for (Future<?> f : futures) {
+      f.get(10, TimeUnit.SECONDS);
+    }
+    executor.shutdown();
+
+    assertEquals(0, errors.size(), "Tags not visible after transition: " + errors);
+  }
 }