From b38bab3261d91de7f726ea6f0d36c78158025fb5 Mon Sep 17 00:00:00 2001 From: Douglas Q Hawkins Date: Fri, 11 Sep 2026 08:17:42 -0400 Subject: [PATCH 01/13] Add find-or-insert Reservation API to ConcurrentHashtable Ports the ConcurrentHashtable Reservation API (auto-cancelling, deferred-construction find-or-insert handle; self-bound Entry replacing separate D1/D2 entry hierarchies) onto bric3's merged LogCollector change, and updates LogCollector to use it. Taking the reservation inside the existing table write lock (rather than lock-free ahead of it) keeps LogCollector's find-or-insert decision serialized with drain() and other writers, so no duplicate reservation can ever be mistaken for a full table. Co-Authored-By: Claude Sonnet 5 --- .../util/ThreadSafeMapCounterBenchmark.java | 29 +- .../trace/util/ThreadSafeMapD1Benchmark.java | 27 +- .../trace/util/ThreadSafeMapD2Benchmark.java | 7 +- .../trace/api/telemetry/LogCollector.java | 41 +- .../trace/util/ConcurrentHashtable.java | 751 ++++++++++++++---- .../trace/util/ConcurrentHashtableD1Test.java | 28 +- .../trace/util/ConcurrentHashtableD2Test.java | 19 +- .../ConcurrentHashtableReservationTest.java | 184 +++++ .../ConcurrentHashtableSizeManagerTest.java | 35 +- .../util/ConcurrentHashtableStaticsTest.java | 7 +- 10 files changed, 892 insertions(+), 236 deletions(-) create mode 100644 internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java diff --git a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapCounterBenchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapCounterBenchmark.java index bc8bc07b9f5..239dd3890f4 100644 --- a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapCounterBenchmark.java +++ b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapCounterBenchmark.java @@ -23,11 +23,10 @@ * Measures lookup followed by an atomic counter increment in a shared, pre-populated table. Models * per-class or per-method hit counters in the tracer. * - *

The {@link ConcurrentHashtable.D1} case embeds a {@code volatile long} in each entry. {@link - * AtomicLongFieldUpdater} updates that field atomically without allocating an {@link AtomicLong} - * per key. The map baselines store a separate {@link AtomicLong} or {@link LongAdder}; {@code - * LongAdder} spreads contention across internal cells at the cost of more memory and a more - * expensive read. + *

The {@link ConcurrentHashtable.D1} case embeds a {@code volatile long} counter directly in the + * entry and increments it via {@link AtomicLongFieldUpdater}, avoiding a second heap object. The + * map baselines store a separate {@link AtomicLong} or {@link LongAdder}; {@code LongAdder} spreads + * contention across internal cells at the cost of more memory and a more expensive read. * *

Lookups reuse the key instances installed during setup. {@code Objects.equals} therefore * returns on its identity check without dispatching to {@code equals}, so this measures the @@ -50,8 +49,8 @@ *

  • {@code LongAdder} is marginally faster (79 vs 71 ops/us) because it shards the counter * across cells to reduce CAS contention; the advantage grows with thread count. *
  • {@code ConcurrentHashtable} matches {@code AtomicLong} throughput (69 vs 71 ops/us) while - * embedding the counter directly in the entry — one object instead of two, with no throughput - * penalty. + * embedding the counter directly in the entry via {@code AtomicLongFieldUpdater} — one object + * instead of two, with no throughput penalty. * */ @Fork(2) @@ -73,8 +72,12 @@ public class ThreadSafeMapCounterBenchmark { } } + /** + * Shared state ({@link Scope#Benchmark}): one instance of each map across all threads, modelling + * a shared instrumentation counter table. + */ static final class CounterEntry extends ConcurrentHashtable.D1.Entry { - private static final AtomicLongFieldUpdater COUNT = + static final AtomicLongFieldUpdater COUNT = AtomicLongFieldUpdater.newUpdater(CounterEntry.class, "count"); volatile long count; @@ -82,16 +85,8 @@ static final class CounterEntry extends ConcurrentHashtable.D1.Entry { CounterEntry(String key) { super(key); } - - long increment() { - return COUNT.incrementAndGet(this); - } } - /** - * Shared state ({@link Scope#Benchmark}): one instance of each map across all threads, modelling - * a shared instrumentation counter table. - */ @State(Scope.Benchmark) public static class SharedState { ConcurrentHashtable.D1 table; @@ -125,7 +120,7 @@ int next() { @Benchmark public long increment_concurrentHashtable(SharedState s, ThreadState t) { - return s.table.get(KEYS[t.next()]).increment(); + return CounterEntry.COUNT.incrementAndGet(s.table.get(KEYS[t.next()])); } @Benchmark diff --git a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD1Benchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD1Benchmark.java index 8fd07544264..6f5a8dcd769 100644 --- a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD1Benchmark.java +++ b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD1Benchmark.java @@ -83,34 +83,35 @@ public class ThreadSafeMapD1Benchmark { } } - static final class D1Entry extends ConcurrentHashtable.D1.Entry { + /** + * Shared state ({@link Scope#Benchmark}): one instance of each map across all threads, modelling + * a shared instrumentation cache. + */ + static final class LongEntry extends ConcurrentHashtable.D1.Entry { final long value; - D1Entry(String key) { + LongEntry(String key, long value) { super(key); - this.value = 1L; + this.value = value; } } - /** - * Shared state ({@link Scope#Benchmark}): one instance of each map across all threads, modelling - * a shared instrumentation cache. - */ @State(Scope.Benchmark) public static class SharedState { - ConcurrentHashtable.D1 table; + ConcurrentHashtable.D1 table; ConcurrentHashMap concurrentHashMap; ConcurrentSkipListMap skipListMap; Map synchronizedHashMap; @Setup(Level.Iteration) public void setUp() { - table = ConcurrentHashtable.D1.createBounded(D1Entry.class, CAPACITY); + table = ConcurrentHashtable.D1.createBounded(LongEntry.class, CAPACITY); concurrentHashMap = new ConcurrentHashMap<>(CAPACITY); skipListMap = new ConcurrentSkipListMap<>(); synchronizedHashMap = Collections.synchronizedMap(new HashMap<>(CAPACITY)); for (int i = 0; i < N_KEYS; ++i) { - table.tryGetOrCreateOrNull(KEYS[i], D1Entry::new); + long value = i; + table.tryGetOrCreateOrNull(KEYS[i], k -> new LongEntry(k, value)); concurrentHashMap.put(KEYS[i], (long) i); skipListMap.put(KEYS[i], (long) i); synchronizedHashMap.put(KEYS[i], (long) i); @@ -131,7 +132,7 @@ int next() { } @Benchmark - public D1Entry get_concurrentHashtable(SharedState s, ThreadState t) { + public LongEntry get_concurrentHashtable(SharedState s, ThreadState t) { return s.table.get(KEYS[t.next()]); } @@ -151,8 +152,8 @@ public Long get_synchronizedHashMap(SharedState s, ThreadState t) { } @Benchmark - public D1Entry getOrCreate_concurrentHashtable(SharedState s, ThreadState t) { - return s.table.tryGetOrCreateOrNull(KEYS[t.next()], D1Entry::new); + public LongEntry getOrCreate_concurrentHashtable(SharedState s, ThreadState t) { + return s.table.tryGetOrCreateOrNull(KEYS[t.next()], k -> new LongEntry(k, 0L)); } /** diff --git a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD2Benchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD2Benchmark.java index aff30dd0a33..57506a12230 100644 --- a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD2Benchmark.java +++ b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD2Benchmark.java @@ -108,7 +108,7 @@ static final class D2Entry extends ConcurrentHashtable.D2.Entry * Entry used with the static helpers. Its primitive second key keeps storage and lookup unboxed, * independently of {@link Integer} caching or JVM escape analysis. */ - static final class SupportEntry extends ConcurrentHashtable.Entry { + static final class SupportEntry extends ConcurrentHashtable.Entry { final String k1; final int k2; final long value; @@ -127,6 +127,11 @@ static long hash(String k1, int k2) { boolean matches(String k1, int k2) { return this.k2 == k2 && this.k1.equals(k1); } + + @Override + public boolean matches(SupportEntry other) { + return matches(other.k1, other.k2); + } } /** Composite key for map-based baselines. */ diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java index a6d23b6e73e..160cdd31c12 100644 --- a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java +++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java @@ -4,12 +4,12 @@ import static datadog.trace.util.ConcurrentHashtable.bucketIndex; import static datadog.trace.util.ConcurrentHashtable.estimateSize; import static datadog.trace.util.ConcurrentHashtable.getTableWriteLock; -import static datadog.trace.util.ConcurrentHashtable.insertReserved; import static datadog.trace.util.ConcurrentHashtable.isFull; import static datadog.trace.util.LongHashingUtils.hash; import datadog.trace.api.internal.VisibleForTesting; import datadog.trace.util.ConcurrentHashtable; +import datadog.trace.util.ConcurrentHashtable.Reservation; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.ArrayList; import java.util.Collection; @@ -40,7 +40,7 @@ private LogCollector() { value = "SING_SINGLETON_HAS_NONPRIVATE_CONSTRUCTOR", justification = "Usage in tests") LogCollector(int maxCapacity) { - this.rawLogMessages = ConcurrentHashtable.State.createBounded(RawLogMessage.class, maxCapacity); + this.rawLogMessages = ConcurrentHashtable.createBounded(RawLogMessage.class, maxCapacity); } public void addLogMessage(String logLevel, String message, @Nullable Throwable throwable) { @@ -81,26 +81,28 @@ public void addLogMessage( return; } - // Slow path after a miss: repeat the lookup and capacity checks under the table write lock - // because another writer or drain() may have changed the table. + // Slow path after a miss: repeat the lookup under the table write lock because another writer + // or drain() may have changed the table. Taking the reservation under this same lock, rather + // than lock-free ahead of it, keeps the whole find-or-insert decision serialized with other + // writers and with drain() -- a writer here genuinely waits for an in-progress drain instead + // of being told, incorrectly, that the table is full because of a duplicate reservation that + // hasn't cancelled yet. Serialized this way, the reservation can never lose a race to insert, + // so there's no existing-vs-new distinction to make: it always inserts fresh. synchronized (getTableWriteLock(rawLogMessages)) { rawLogMessage = find(bucketIndex, keyHash, logLevel, message, throwable); if (rawLogMessage != null) { rawLogMessage.increment(); return; } - // Capacity may have been released by drain() or consumed by another writer while waiting. - if (isFull(rawLogMessages)) { - return; - } - - // Allocate before reserving because a reservation cannot - // be rolled back if construction fails. - rawLogMessage = - new RawLogMessage(logLevel, message, throwable, tags, System.currentTimeMillis() / 1000); - // Reserve before linking so every published entry is included in the capacity count. - if (rawLogMessages.sizeManager.tryReserve()) { - insertReserved(rawLogMessages, keyHash, rawLogMessage); + try (Reservation reservation = + ConcurrentHashtable.tryReserve(rawLogMessages)) { + if (reservation.isReserved()) { + // Allocate only after the reservation succeeds. + reservation.tryGetOrInsertOrNull( + new RawLogMessage( + logLevel, message, throwable, tags, System.currentTimeMillis() / 1000)); + } + // TODO: We could emit a metric for dropped logs. } } } @@ -190,7 +192,7 @@ private RawLogMessage find( * match. The first message supplies the tags and timestamp; later messages only increment the * occurrence count. */ - public static final class RawLogMessage extends ConcurrentHashtable.Entry { + public static final class RawLogMessage extends ConcurrentHashtable.Entry { private static final AtomicIntegerFieldUpdater LIVE_OCCURRENCE_COUNT_UPDATER = AtomicIntegerFieldUpdater.newUpdater(RawLogMessage.class, "liveOccurrenceCount"); @@ -239,6 +241,11 @@ private void snapshotCount() { count = LIVE_OCCURRENCE_COUNT_UPDATER.get(this); } + @Override + public boolean matches(RawLogMessage that) { + return equals(that); + } + @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java index cc0b53ee10e..73276ccf594 100644 --- a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java +++ b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java @@ -1,6 +1,10 @@ package datadog.trace.util; +import datadog.trace.api.function.Strategy; +import datadog.trace.api.function.StrategyConsumer; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.util.Iterator; +import java.util.NoSuchElementException; import java.util.Objects; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReferenceArray; @@ -69,10 +73,16 @@ private ConcurrentHashtable() {} * arity. See {@link D1.Entry} and {@link D2.Entry}; for higher arities, or for primitive key * components, subclass this directly and drive the table with the static building blocks on * {@link ConcurrentHashtable}. + * + *

    The self-bound type parameter ({@code TEntry extends Entry}, mirroring {@code Enum>}) exists so {@link #matches(Entry)} can compare two already-built entries + * without an {@code Object}/{@code instanceof} boundary. It requires every concrete subclass to + * declare itself as its own type argument (see {@link D1.Entry}/{@link D2.Entry}); Java has no + * true {@code Self} type, so this is enforced by convention, not the compiler. */ - public abstract static class Entry { + public abstract static class Entry> { public final long keyHash; - private volatile Entry next = null; + private volatile TEntry next = null; protected Entry(long keyHash) { this.keyHash = keyHash; @@ -81,15 +91,26 @@ protected Entry(long keyHash) { // Package-private: the only writers are the static insert/remove building blocks // (insertHeadEntry, unlink) on the enclosing class, which reach it via the Entry bound. Custom // tables mutate chains through those helpers, never by touching next directly. - final void setNext(TEntry next) { + final void setNext(TEntry next) { this.next = next; } - @SuppressWarnings("unchecked") @Nullable - public final TEntry next() { - return (TEntry) this.next; + public final TEntry next() { + return this.next; } + + /** + * Returns {@code true} if {@code other} is logically the same entry as this one (same key(s)), + * used to compare two already-built entries under the write lock (e.g. in {@link + * Reservation#tryGetOrInsertOrNull}). Deliberately narrower than {@code equals}/{@code + * hashCode}: this class doesn't need reflexive/symmetric-with-null-and-unrelated-types contract + * baggage, and a bespoke method avoids entries accidentally working as {@code HashSet}/{@code + * HashMap} keys via an unrelated identity notion. {@link D1.Entry}/{@link D2.Entry} implement + * this in terms of their existing key-based {@code matches(...)}, so most callers never write + * it directly. + */ + public abstract boolean matches(@Nonnull TEntry other); } /** @@ -105,10 +126,17 @@ public static final class D1> { * Abstract base for {@link D1} entries. Subclass to add value fields you wish to mutate in * place after retrieving the entry via {@link D1#get}. * + *

    Deliberately parameterized on {@code K} alone, not self-bound on the concrete subclass: + * {@link D1} stores and links entries internally as {@code Entry} and casts back to {@code + * TEntry} at its API boundary, trading one unchecked cast (always sound -- the table only ever + * holds instances the caller's own {@code creator} produced) for a simpler subclass signature, + * e.g. {@code class MyEntry extends D1.Entry} rather than {@code D1.Entry}. + * * @param the key type */ - public abstract static class Entry extends ConcurrentHashtable.Entry { - final K key; + public abstract static class Entry extends ConcurrentHashtable.Entry> { + @Nullable final K key; protected Entry(@Nullable K key) { super(hash(key)); @@ -127,6 +155,12 @@ public boolean matches(@Nullable Object key) { return Objects.equals(key, this.key); } + /** {@link ConcurrentHashtable.Entry#matches(Entry)} in terms of the key-based overload. */ + @Override + public final boolean matches(@Nonnull Entry other) { + return matches(other.key); + } + /** * Returns the 64-bit lookup hash for {@code key}. Null keys map to {@link Long#MIN_VALUE} so * they don't collide with a real key that hashes to 0; real-key collisions in chains are @@ -137,21 +171,51 @@ public static long hash(@Nullable Object key) { } } - private final State state; + private final State> state; - private D1(State state) { + private D1(State> state) { this.state = state; } /** * Creates a fixed-size table holding at most {@code maxCapacity} entries. {@code entryClass} is - * used only to infer the concrete entry type; entries are created by the functions passed to - * the insertion methods. The table does not resize. + * used only to allocate the backing array with the right component type; entries themselves are + * created by the {@code creator}/{@code evictable} functions passed to the insertion methods. + * The table does not resize. */ @Nonnull + @SuppressWarnings("unchecked") public static > D1 createBounded( @Nonnull Class entryClass, int maxCapacity) { - return new D1<>(State.createBounded(entryClass, maxCapacity)); + // entryClass is erased away (see ConcurrentHashtable#createFixedBuckets), so treating it as + // Class> instead of the caller's concrete Class is safe. + Class> baseEntryClass = (Class>) (Class) entryClass; + return new D1<>(ConcurrentHashtable.createBounded(baseEntryClass, maxCapacity)); + } + + /** + * Sound because every entry ever inserted into {@link #state} was produced as a {@code TEntry}. + */ + @SuppressWarnings("unchecked") + private TEntry cast(@Nullable Entry entry) { + return (TEntry) entry; + } + + /** See {@link #cast(Entry)}; casts the functional-interface reference, not each element. */ + @SuppressWarnings("unchecked") + private Predicate> castPredicate(Predicate predicate) { + return (Predicate>) predicate; + } + + @SuppressWarnings("unchecked") + private Consumer> castConsumer(Consumer consumer) { + return (Consumer>) consumer; + } + + @SuppressWarnings("unchecked") + private BiConsumer> castConsumer( + BiConsumer consumer) { + return (BiConsumer>) consumer; } public int size() { @@ -165,11 +229,11 @@ public boolean isFull() { @Nullable public TEntry get(@Nullable K key) { long keyHash = D1.Entry.hash(key); - for (TEntry curEntry = bucketFor(state, keyHash); + for (Entry curEntry = bucketFor(state, keyHash); curEntry != null; curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key)) { - return curEntry; + return cast(curEntry); } } return null; @@ -182,7 +246,7 @@ public TEntry get(@Nullable K key) { */ @Nonnull public Maybe tryGetOrCreate( - @Nullable K key, @Nonnull Function creator) { + @Nullable K key, @Strategy @Nonnull Function creator) { return Maybe.of(tryGetOrCreateOrNull(key, creator)); } @@ -192,22 +256,25 @@ public Maybe tryGetOrCreate( * {@code key} was not already present. Re-checks under the lock to avoid duplicate entries * under concurrent misses. */ + @StrategyConsumer @Nullable public TEntry tryGetOrCreateOrNull( - @Nullable K key, @Nonnull Function creator) { + @Nullable K key, @Strategy @Nonnull Function creator) { long keyHash = D1.Entry.hash(key); int index = bucketIndex(state.buckets, keyHash); - for (TEntry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { + for (Entry curEntry = bucketAt(state, index); + curEntry != null; + curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key)) { - return curEntry; + return cast(curEntry); } } synchronized (getTableWriteLock(state)) { - for (TEntry curEntry = bucketAt(state, index); + for (Entry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key)) { - return curEntry; + return cast(curEntry); } } // isFull() is checked before creating the entry, not before reserving a slot for it: @@ -233,8 +300,8 @@ public TEntry tryGetOrCreateOrNull( @Nonnull public Maybe tryGetOrCreateOrEvict( @Nullable K key, - @Nonnull Function creator, - @Nonnull Predicate evictable) { + @Strategy @Nonnull Function creator, + @Strategy @Nonnull Predicate evictable) { return Maybe.of(tryGetOrCreateOrEvictOrNull(key, creator, evictable)); } @@ -245,28 +312,31 @@ public Maybe tryGetOrCreateOrEvict( * ever leaving a slot double-booked. A creator that throws after a successful eviction simply * leaves the table one entry smaller — no corruption, just a wasted eviction. */ + @StrategyConsumer @Nullable public TEntry tryGetOrCreateOrEvictOrNull( @Nullable K key, - @Nonnull Function creator, - @Nonnull Predicate evictable) { + @Strategy @Nonnull Function creator, + @Strategy @Nonnull Predicate evictable) { long keyHash = D1.Entry.hash(key); int index = bucketIndex(state.buckets, keyHash); - for (TEntry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { + for (Entry curEntry = bucketAt(state, index); + curEntry != null; + curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key)) { - return curEntry; + return cast(curEntry); } } synchronized (getTableWriteLock(state)) { - for (TEntry curEntry = bucketAt(state, index); + for (Entry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key)) { - return curEntry; + return cast(curEntry); } } if (state.sizeManager.isFull() - && state.sizeManager.evictOne(state.buckets, evictable) == null) { + && state.sizeManager.evictOne(state.buckets, castPredicate(evictable)) == null) { return null; } TEntry newEntry = creator.apply(key); @@ -286,14 +356,14 @@ public TEntry remove(@Nullable K key) { long keyHash = D1.Entry.hash(key); int index = bucketIndex(state.buckets, keyHash); synchronized (getTableWriteLock(state)) { - TEntry prev = null; - for (TEntry curEntry = bucketAt(state, index); + Entry prev = null; + for (Entry curEntry = bucketAt(state, index); curEntry != null; prev = curEntry, curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key)) { unlink(state, index, prev, curEntry); state.sizeManager.decrement(); - return curEntry; + return cast(curEntry); } } return null; @@ -305,8 +375,8 @@ public TEntry remove(@Nullable K key) { * Holds the table-level lock for the whole sweep, so the predicate sees a stable table and * concurrent writers are excluded; lock-free readers continue throughout. */ - public boolean removeIf(@Nonnull Predicate predicate) { - return ConcurrentHashtable.removeIf(state, predicate); + public boolean removeIf(@Strategy @Nonnull Predicate predicate) { + return ConcurrentHashtable.removeIf(state, castPredicate(predicate)); } /** @@ -321,8 +391,8 @@ public boolean removeIf(@Nonnull Predicate predicate) { * * @param drainedEntryConsumer action invoked for each removed entry */ - public void drain(@Nonnull Consumer drainedEntryConsumer) { - ConcurrentHashtable.drain(state, drainedEntryConsumer); + public void drain(@Strategy @Nonnull Consumer drainedEntryConsumer) { + ConcurrentHashtable.drain(state, castConsumer(drainedEntryConsumer)); } /** @@ -340,8 +410,8 @@ public void drain(@Nonnull Consumer drainedEntryConsumer) { * @param drainedEntryConsumer action invoked with the context and each removed entry */ public void drain( - C context, @Nonnull BiConsumer drainedEntryConsumer) { - ConcurrentHashtable.drain(state, context, drainedEntryConsumer); + C context, @Strategy @Nonnull BiConsumer drainedEntryConsumer) { + ConcurrentHashtable.drain(state, context, castConsumer(drainedEntryConsumer)); } /** Removes all entries. Lock-free readers mid-walk complete against the entries they hold. */ @@ -349,16 +419,17 @@ public void clear() { ConcurrentHashtable.clear(state); } - public void forEach(@Nonnull Consumer consumer) { - ConcurrentHashtable.forEach(state, consumer); + public void forEach(@Strategy @Nonnull Consumer consumer) { + ConcurrentHashtable.forEach(state, castConsumer(consumer)); } /** * Context-passing forEach. Avoids a capturing-lambda allocation — pass a non-capturing {@link * BiConsumer} (typically a {@code static final}) plus whatever side-band state it needs. */ - public void forEach(C context, @Nonnull BiConsumer consumer) { - ConcurrentHashtable.forEach(state, context, consumer); + public void forEach( + C context, @Strategy @Nonnull BiConsumer consumer) { + ConcurrentHashtable.forEach(state, context, castConsumer(consumer)); } } @@ -376,14 +447,17 @@ public static final class D2> { /** * Abstract base for {@link D2} entries. Subclass to add value fields you wish to mutate in - * place. + * place after retrieving the entry via {@link D2#get}. + * + *

    Deliberately parameterized on {@code K1}/{@code K2} alone, not self-bound on the concrete + * subclass -- see {@link D1.Entry} for why. * * @param first key type * @param second key type */ - public abstract static class Entry extends ConcurrentHashtable.Entry { - final K1 key1; - final K2 key2; + public abstract static class Entry extends ConcurrentHashtable.Entry> { + @Nullable final K1 key1; + @Nullable final K2 key2; protected Entry(@Nullable K1 key1, @Nullable K2 key2) { super(hash(key1, key2)); @@ -409,27 +483,63 @@ public boolean matches(@Nullable K1 key1, @Nullable K2 key2) { return Objects.equals(key1, this.key1) && Objects.equals(key2, this.key2); } + /** {@link ConcurrentHashtable.Entry#matches(Entry)} in terms of the key-based overload. */ + @Override + public final boolean matches(@Nonnull Entry other) { + return matches(other.key1, other.key2); + } + /** Returns the 64-bit lookup hash combining both key parts via {@link LongHashingUtils}. */ public static long hash(@Nullable Object key1, @Nullable Object key2) { return LongHashingUtils.hash(key1, key2); } } - private final State state; + private final State> state; - private D2(State state) { + private D2(State> state) { this.state = state; } /** * Creates a fixed-size table holding at most {@code maxCapacity} entries. {@code entryClass} is - * used only to infer the concrete entry type; entries are created by the functions passed to - * the insertion methods. The table does not resize. + * used only to allocate the backing array with the right component type; entries themselves are + * created by the {@code creator}/{@code evictable} functions passed to the insertion methods. + * The table does not resize. */ @Nonnull + @SuppressWarnings("unchecked") public static > D2 createBounded( @Nonnull Class entryClass, int maxCapacity) { - return new D2<>(State.createBounded(entryClass, maxCapacity)); + // entryClass is erased away (see ConcurrentHashtable#createFixedBuckets), so treating it as + // Class> instead of the caller's concrete Class is safe. + Class> baseEntryClass = (Class>) (Class) entryClass; + return new D2<>(ConcurrentHashtable.createBounded(baseEntryClass, maxCapacity)); + } + + /** + * Sound because every entry ever inserted into {@link #state} was produced as a {@code TEntry}. + */ + @SuppressWarnings("unchecked") + private TEntry cast(@Nullable Entry entry) { + return (TEntry) entry; + } + + /** See {@link #cast(Entry)}; casts the functional-interface reference, not each element. */ + @SuppressWarnings("unchecked") + private Predicate> castPredicate(Predicate predicate) { + return (Predicate>) predicate; + } + + @SuppressWarnings("unchecked") + private Consumer> castConsumer(Consumer consumer) { + return (Consumer>) consumer; + } + + @SuppressWarnings("unchecked") + private BiConsumer> castConsumer( + BiConsumer consumer) { + return (BiConsumer>) consumer; } public int size() { @@ -443,11 +553,11 @@ public boolean isFull() { @Nullable public TEntry get(@Nullable K1 key1, @Nullable K2 key2) { long keyHash = D2.Entry.hash(key1, key2); - for (TEntry curEntry = bucketFor(state, keyHash); + for (Entry curEntry = bucketFor(state, keyHash); curEntry != null; curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key1, key2)) { - return curEntry; + return cast(curEntry); } } return null; @@ -457,15 +567,12 @@ public TEntry get(@Nullable K1 key1, @Nullable K2 key2) { * Returns the entry for {@code (key1, key2)}, creating one via {@code creator} if absent and * the table is under capacity. Lock-free on hit; acquires a table-level lock on miss. Wraps * {@link #tryGetOrCreateOrNull} — see that method for the refusal and ordering details. - * - *

    The {@code creator} should build an entry whose {@code keyHash} equals {@link - * D2.Entry#hash(Object, Object) D2.Entry.hash(key1, key2)}. */ @Nonnull public Maybe tryGetOrCreate( @Nullable K1 key1, @Nullable K2 key2, - @Nonnull BiFunction creator) { + @Strategy @Nonnull BiFunction creator) { return Maybe.of(tryGetOrCreateOrNull(key1, key2, creator)); } @@ -475,24 +582,27 @@ public Maybe tryGetOrCreate( * {@code (key1, key2)} was not already present. Re-checks under the lock to avoid duplicate * entries under concurrent misses. */ + @StrategyConsumer @Nullable public TEntry tryGetOrCreateOrNull( @Nullable K1 key1, @Nullable K2 key2, - @Nonnull BiFunction creator) { + @Strategy @Nonnull BiFunction creator) { long keyHash = D2.Entry.hash(key1, key2); int index = bucketIndex(state.buckets, keyHash); - for (TEntry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { + for (Entry curEntry = bucketAt(state, index); + curEntry != null; + curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key1, key2)) { - return curEntry; + return cast(curEntry); } } synchronized (getTableWriteLock(state)) { - for (TEntry curEntry = bucketAt(state, index); + for (Entry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key1, key2)) { - return curEntry; + return cast(curEntry); } } // isFull() is checked before creating the entry, not before reserving a slot for it: @@ -519,8 +629,8 @@ public TEntry tryGetOrCreateOrNull( public Maybe tryGetOrCreateOrEvict( @Nullable K1 key1, @Nullable K2 key2, - @Nonnull BiFunction creator, - @Nonnull Predicate evictable) { + @Strategy @Nonnull BiFunction creator, + @Strategy @Nonnull Predicate evictable) { return Maybe.of(tryGetOrCreateOrEvictOrNull(key1, key2, creator, evictable)); } @@ -531,29 +641,32 @@ public Maybe tryGetOrCreateOrEvict( * ever leaving a slot double-booked. A creator that throws after a successful eviction simply * leaves the table one entry smaller — no corruption, just a wasted eviction. */ + @StrategyConsumer @Nullable public TEntry tryGetOrCreateOrEvictOrNull( @Nullable K1 key1, @Nullable K2 key2, - @Nonnull BiFunction creator, - @Nonnull Predicate evictable) { + @Strategy @Nonnull BiFunction creator, + @Strategy @Nonnull Predicate evictable) { long keyHash = D2.Entry.hash(key1, key2); int index = bucketIndex(state.buckets, keyHash); - for (TEntry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { + for (Entry curEntry = bucketAt(state, index); + curEntry != null; + curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key1, key2)) { - return curEntry; + return cast(curEntry); } } synchronized (getTableWriteLock(state)) { - for (TEntry curEntry = bucketAt(state, index); + for (Entry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key1, key2)) { - return curEntry; + return cast(curEntry); } } if (state.sizeManager.isFull() - && state.sizeManager.evictOne(state.buckets, evictable) == null) { + && state.sizeManager.evictOne(state.buckets, castPredicate(evictable)) == null) { return null; } TEntry newEntry = creator.apply(key1, key2); @@ -573,14 +686,14 @@ public TEntry remove(@Nullable K1 key1, @Nullable K2 key2) { long keyHash = D2.Entry.hash(key1, key2); int index = bucketIndex(state.buckets, keyHash); synchronized (getTableWriteLock(state)) { - TEntry prev = null; - for (TEntry curEntry = bucketAt(state, index); + Entry prev = null; + for (Entry curEntry = bucketAt(state, index); curEntry != null; prev = curEntry, curEntry = curEntry.next()) { if (curEntry.keyHash == keyHash && curEntry.matches(key1, key2)) { unlink(state, index, prev, curEntry); state.sizeManager.decrement(); - return curEntry; + return cast(curEntry); } } return null; @@ -592,8 +705,8 @@ public TEntry remove(@Nullable K1 key1, @Nullable K2 key2) { * Holds the table-level lock for the whole sweep, so the predicate sees a stable table and * concurrent writers are excluded; lock-free readers continue throughout. */ - public boolean removeIf(@Nonnull Predicate predicate) { - return ConcurrentHashtable.removeIf(state, predicate); + public boolean removeIf(@Strategy @Nonnull Predicate predicate) { + return ConcurrentHashtable.removeIf(state, castPredicate(predicate)); } /** @@ -608,8 +721,8 @@ public boolean removeIf(@Nonnull Predicate predicate) { * * @param drainedEntryConsumer action invoked for each removed entry */ - public void drain(@Nonnull Consumer drainedEntryConsumer) { - ConcurrentHashtable.drain(state, drainedEntryConsumer); + public void drain(@Strategy @Nonnull Consumer drainedEntryConsumer) { + ConcurrentHashtable.drain(state, castConsumer(drainedEntryConsumer)); } /** @@ -627,8 +740,8 @@ public void drain(@Nonnull Consumer drainedEntryConsumer) { * @param drainedEntryConsumer action invoked with the context and each removed entry */ public void drain( - C context, @Nonnull BiConsumer drainedEntryConsumer) { - ConcurrentHashtable.drain(state, context, drainedEntryConsumer); + C context, @Strategy @Nonnull BiConsumer drainedEntryConsumer) { + ConcurrentHashtable.drain(state, context, castConsumer(drainedEntryConsumer)); } /** Removes all entries. Lock-free readers mid-walk complete against the entries they hold. */ @@ -636,16 +749,17 @@ public void clear() { ConcurrentHashtable.clear(state); } - public void forEach(@Nonnull Consumer consumer) { - ConcurrentHashtable.forEach(state, consumer); + public void forEach(@Strategy @Nonnull Consumer consumer) { + ConcurrentHashtable.forEach(state, castConsumer(consumer)); } /** * Context-passing forEach. Avoids a capturing-lambda allocation — pass a non-capturing {@link * BiConsumer} (typically a {@code static final}) plus whatever side-band state it needs. */ - public void forEach(C context, @Nonnull BiConsumer consumer) { - ConcurrentHashtable.forEach(state, context, consumer); + public void forEach( + C context, @Strategy @Nonnull BiConsumer consumer) { + ConcurrentHashtable.forEach(state, context, castConsumer(consumer)); } } @@ -693,8 +807,9 @@ public boolean isFull() { * cannot both acquire the last slot. Returns {@code false} with the count unchanged when the * table is full. * - *

    Build the entry before reserving: there is no cancellation operation, so abandoning a - * successful reservation permanently consumes capacity. + *

    Prefer {@link #cancelReservation()} plus deferred entry construction (see {@link + * ConcurrentHashtable#reserve}) over abandoning a reservation outright: this method by itself + * still has no way to give back a slot once claimed. */ public boolean tryReserve() { if (size.incrementAndGet() > capacity) { @@ -704,6 +819,27 @@ public boolean tryReserve() { return true; } + /** + * Gives back a slot claimed by {@link #tryReserve()} that was never filled — e.g. a concurrent + * match was found under the write lock instead of inserting. Lock-free, symmetric with {@link + * #decrement()}. + * + *

    Caller must call this at most once per successful {@link #tryReserve()}; double-cancelling + * corrupts the count the same way double-incrementing would. {@link Reservation#close()} + * handles this bookkeeping automatically and should be preferred over calling this directly. + * + *

    Under heavy contention on the same logical duplicate, multiple threads can each reserve a + * slot for what turns out to be the same entry before any of them cancels, transiently + * inflating {@code size} above the table's true occupancy. This can cause an unrelated, + * genuinely distinct concurrent insert to see the table as full when it isn't, until the losing + * reservations cancel. The effect is self-correcting (bounded by in-flight reservations, not + * sustained) and considered an acceptable tradeoff for tables expecting bursts of identical + * inserts (e.g. deduplication). + */ + public void cancelReservation() { + size.decrementAndGet(); + } + /** * Reserves one slot, evicting an entry matching {@code evictable} when the table is full. * Returns {@code false} without changing the table when no entry can be evicted. @@ -712,9 +848,9 @@ public boolean tryReserve() { * an abandoned reservation permanently consumes capacity. */ @GuardedBy("getTableWriteLock(buckets)") - public boolean tryReserveOrEvict( + public > boolean tryReserveOrEvict( @Nonnull AtomicReferenceArray buckets, - @Nonnull Predicate evictable) { + @Strategy @Nonnull Predicate evictable) { if (tryReserve()) { return true; } @@ -763,9 +899,9 @@ public void release(int removed) { */ @GuardedBy("getTableWriteLock(buckets)") @Nullable - public TEntry evictOne( + public > TEntry evictOne( @Nonnull AtomicReferenceArray buckets, - @Nonnull Predicate evictable) { + @Strategy @Nonnull Predicate evictable) { TEntry evicted = evictOneInRange(buckets, evictable, evictionCursor, buckets.length()); if (evicted == null && evictionCursor != 0) { evicted = evictOneInRange(buckets, evictable, 0, evictionCursor); @@ -787,10 +923,11 @@ public TEntry evictOne( justification = "evictionCursor is read and written only under synchronized (getTableWriteLock(buckets)); SpotBugs" + " cannot model that dynamic guard") + @StrategyConsumer @Nullable - private TEntry evictOneInRange( + private > TEntry evictOneInRange( @Nonnull AtomicReferenceArray buckets, - @Nonnull Predicate evictable, + @Strategy @Nonnull Predicate evictable, int startBucket, int endBucket) { for (int i = startBucket; i < endBucket; i++) { @@ -818,9 +955,10 @@ private TEntry evictOneInRange( justification = "evictionCursor is read and written only under synchronized (getTableWriteLock(buckets)); SpotBugs" + " cannot model that dynamic guard") - public int evictAll( + @StrategyConsumer + public > int evictAll( @Nonnull AtomicReferenceArray buckets, - @Nonnull Predicate evictable) { + @Strategy @Nonnull Predicate evictable) { int count = 0; for (int i = 0; i < buckets.length(); i++) { TEntry prev = null; @@ -842,25 +980,30 @@ public int evictAll( /** * Bucket array and occupancy manager for a caller-defined capped table. Keep them paired and * prefer the {@code State}-accepting helpers so structural changes update the count consistently. + * + *

    {@code sizeManager} is intentionally package-private: callers outside this class must go + * through the {@code State}-accepting static helpers ({@link #estimateSize}, {@link #isFull}, + * {@link #tryReserve}, {@link #tryReserveOrEvict}, {@link #evictOne}, {@link #evictAll}) rather + * than reach into the manager directly. */ - public static final class State { + public static final class State> { public final AtomicReferenceArray buckets; - public final SizeManager sizeManager; + final SizeManager sizeManager; private State(AtomicReferenceArray buckets, int maxCapacity) { this.buckets = buckets; this.sizeManager = new SizeManager(maxCapacity); } + } - /** - * Creates a bucket array for {@code maxCapacity} entries and pairs it with a manager enforcing - * that cap. {@code entryClass} is used only to infer {@code TEntry}. - */ - @Nonnull - public static State createBounded( - @Nonnull Class entryClass, int maxCapacity) { - return new State<>(createFixedBuckets(entryClass, maxCapacity), maxCapacity); - } + /** + * Creates a bucket array for {@code maxCapacity} entries and pairs it with a manager enforcing + * that cap. {@code entryClass} is used only to infer {@code TEntry}. + */ + @Nonnull + public static > State createBounded( + @Nonnull Class entryClass, int maxCapacity) { + return new State<>(createFixedBuckets(entryClass, maxCapacity), maxCapacity); } /** Live entries in {@code state}; see {@link SizeManager#estimateSize()}. Lock-free. */ @@ -875,6 +1018,237 @@ public static boolean isFull(@Nonnull State state) { return state.sizeManager.isFull(); } + /** + * Reserves one slot in {@code state} without evicting; see {@link SizeManager#tryReserve()}. + * Lock-free — does not acquire the table write lock. Returns {@code false} with the table + * unchanged when it is full. + * + *

    Complete it with {@link #insertReserved}, or prefer {@link #tryReserve} for a higher-level, + * auto-cancelling handle that also defers entry construction until the reservation succeeds. + */ + public static > boolean tryReserveSlot( + @Nonnull State state) { + return state.sizeManager.tryReserve(); + } + + /** + * Claims one slot in {@code state} and returns a handle for completing the find-or-insert + * protocol, or an empty handle if the table is full. Lock-free — does not acquire the table write + * lock. Never returns {@code null}, so this always composes with try-with-resources: + * + * 

    {@code
+   * try (Reservation r = ConcurrentHashtable.tryReserve(state)) {
+   *   return r.tryGetOrInsertOrNull(component1, component2, component3, TEntry::new);
+   * }
+   * }
    + * + *

    Run a lock-free scan first (see {@link #bucketFor}/{@link #bucketAt}) and only call this + * once that scan has missed — a successful reservation isn't required for correctness (the + * reservation itself, and the locked comparison inside {@link Reservation#tryGetOrInsertOrNull}, + * are the source of truth), it just avoids paying for a lock and a factory call when a hit was + * already visible lock-free. + * + *

    Always returns a non-null handle — even when the table is full — so the caller must check + * {@link Reservation#isReserved()} (or simply call {@link Reservation#tryGetOrInsertOrNull}, + * which returns {@code null} on an absent reservation) rather than assume every reservation is + * real. + */ + @Nonnull + public static > Reservation tryReserve( + @Nonnull State state) { + return new Reservation<>(state.sizeManager.tryReserve() ? state : null); + } + + /** + * Handle returned by {@link #tryReserve}, gating {@link #tryGetOrInsertOrNull} behind a claimed + * slot and auto-cancelling it on {@link #close} if it's never consumed. A single {@code + * Reservation} must be used for at most one {@code tryGetOrInsertOrNull} call. + * + *

    Overloaded up to 4 key components ({@link #tryGetOrInsertOrNull(Object, Function)} through + * {@link #tryGetOrInsertOrNull(Object, Object, Object, Object, Function4)}) so a non-capturing + * method reference can build the entry directly from its natural constructor arguments, without + * an intermediate holder object or a capturing lambda. + * + * @param the table's entry type, itself self-bound (see {@link + * ConcurrentHashtable.Entry}) + */ + public static final class Reservation> implements AutoCloseable { + @Nullable private final State state; + private boolean consumed; + + private Reservation(@Nullable State state) { + this.state = state; + } + + /** {@code true} if this is a real, claimed reservation rather than an empty one. */ + public boolean isReserved() { + return state != null; + } + + /** + * Escape hatch for a caller that already built {@code newEntry} itself -- e.g. more than 4 key + * components, or components the caller wants to keep as primitives rather than boxing them into + * a {@code Function}'s type argument: + * + * 

    {@code
+     * try (Reservation r = ConcurrentHashtable.tryReserve(state)) {
+     *   if (!r.isReserved()) {
+     *     return null;
+     *   }
+     *   return r.tryGetOrInsertOrNull(new TEntry(longComponent1, longComponent2));
+     * }
+     * }
    + * + * See {@link #tryGetOrInsertOrNull(Object, Object, BiFunction)} for the general contract. + */ + @Nullable + public TEntry tryGetOrInsertOrNull(@Nonnull TEntry newEntry) { + return state == null ? null : finish(newEntry); + } + + /** + * One key component; see {@link #tryGetOrInsertOrNull(Object, Object, BiFunction)} for the + * general contract. + */ + @StrategyConsumer + @Nullable + public TEntry tryGetOrInsertOrNull( + A a, @Strategy @Nonnull Function factory) { + return state == null ? null : finish(factory.apply(a)); + } + + /** + * Two key components. Builds {@code factory.apply(...)} (skipped entirely if this reservation + * is empty — the table was full) and either links the result as a new entry or discards it in + * favor of an existing match found under the write lock. Returns {@code null} only when this + * reservation is empty; otherwise always returns a real entry (the newly built one, or the + * concurrent match). + * + *

    Building the entry here, after the reservation already succeeded, keeps the write lock's + * critical section limited to the comparison/link/discard decision rather than whatever + * construction cost {@code factory} pays. See {@link ConcurrentHashtable.Entry#matches} — the + * under-lock comparison is entry-to-entry, so it needs {@code newEntry} already built. + */ + @StrategyConsumer + @Nullable + public TEntry tryGetOrInsertOrNull( + A a, B b, @Strategy @Nonnull BiFunction factory) { + return state == null ? null : finish(factory.apply(a, b)); + } + + /** + * Three key components; see {@link #tryGetOrInsertOrNull(Object, Object, BiFunction)} for the + * general contract. + */ + @StrategyConsumer + @Nullable + public TEntry tryGetOrInsertOrNull( + A a, + B b, + C c, + @Strategy @Nonnull Function3 factory) { + return state == null ? null : finish(factory.apply(a, b, c)); + } + + /** Four key components; see {@link #tryGetOrInsertOrNull(Object, Object, BiFunction)}. */ + @StrategyConsumer + @Nullable + public TEntry tryGetOrInsertOrNull( + A a, + B b, + C c, + D d, + @Strategy @Nonnull + Function4 factory) { + return state == null ? null : finish(factory.apply(a, b, c, d)); + } + + /** {@link Maybe}-wrapping counterpart of {@link #tryGetOrInsertOrNull(Entry)}. */ + @Nonnull + public Maybe tryGetOrInsert(@Nonnull TEntry newEntry) { + return Maybe.of(tryGetOrInsertOrNull(newEntry)); + } + + /** + * {@link Maybe}-wrapping counterpart of {@link #tryGetOrInsertOrNull(Object, Function)}, for + * callers who'd rather make the "this can fail unlike unbounded collections" outcome visible in + * the return type than rely on a {@code null} check — mirrors {@link D1#tryGetOrCreate} / + * {@link D2#tryGetOrCreate} wrapping their own {@code ...OrNull} methods. + */ + @Nonnull + public Maybe tryGetOrInsert( + A a, @Strategy @Nonnull Function factory) { + return Maybe.of(tryGetOrInsertOrNull(a, factory)); + } + + /** Two key components; see {@link #tryGetOrInsert(Object, Function)}. */ + @Nonnull + public Maybe tryGetOrInsert( + A a, B b, @Strategy @Nonnull BiFunction factory) { + return Maybe.of(tryGetOrInsertOrNull(a, b, factory)); + } + + /** Three key components; see {@link #tryGetOrInsert(Object, Function)}. */ + @Nonnull + public Maybe tryGetOrInsert( + A a, + B b, + C c, + @Strategy @Nonnull Function3 factory) { + return Maybe.of(tryGetOrInsertOrNull(a, b, c, factory)); + } + + /** Four key components; see {@link #tryGetOrInsert(Object, Function)}. */ + @Nonnull + public Maybe tryGetOrInsert( + A a, + B b, + C c, + D d, + @Strategy @Nonnull + Function4 factory) { + return Maybe.of(tryGetOrInsertOrNull(a, b, c, d, factory)); + } + + private TEntry finish(@Nonnull TEntry newEntry) { + synchronized (getTableWriteLock(state)) { + int index = bucketIndex(state.buckets, newEntry.keyHash); + for (TEntry curEntry = bucketAt(state, index); + curEntry != null; + curEntry = curEntry.next()) { + if (curEntry.keyHash == newEntry.keyHash && curEntry.matches(newEntry)) { + return curEntry; + } + } + insertHeadEntryAt(state, index, newEntry); + consumed = true; + return newEntry; + } + } + + /** Gives back an unconsumed reservation's slot; a no-op on an empty reservation. */ + @Override + public void close() { + if (state != null && !consumed) { + state.sizeManager.cancelReservation(); + } + } + } + + /** Three-argument analogue of {@link java.util.function.BiFunction}. */ + @Strategy + @FunctionalInterface + public interface Function3 { + R apply(A a, B b, C c); + } + + /** Four-argument analogue of {@link java.util.function.BiFunction}. */ + @Strategy + @FunctionalInterface + public interface Function4 { + R apply(A a, B b, C c, D d); + } + /** * Reserves one slot in {@code state}, evicting an entry matching {@code evictable} when * necessary. Returns {@code false} if the table is full and nothing can be evicted. This method @@ -883,8 +1257,8 @@ public static boolean isFull(@Nonnull State state) { *

    The reservation survives drain and clear operations. Complete it with {@link * #insertReserved}; abandoning it permanently consumes capacity. */ - public static boolean tryReserveOrEvict( - @Nonnull State state, @Nonnull Predicate evictable) { + public static > boolean tryReserveOrEvict( + @Nonnull State state, @Strategy @Nonnull Predicate evictable) { synchronized (getTableWriteLock(state)) { return state.sizeManager.tryReserveOrEvict(state.buckets, evictable); } @@ -896,8 +1270,8 @@ public static boolean tryReserveOrEvict( * Self-locking. */ @Nullable - public static TEntry evictOne( - @Nonnull State state, @Nonnull Predicate evictable) { + public static > TEntry evictOne( + @Nonnull State state, @Strategy @Nonnull Predicate evictable) { synchronized (getTableWriteLock(state)) { return state.sizeManager.evictOne(state.buckets, evictable); } @@ -907,8 +1281,8 @@ public static TEntry evictOne( * Unlinks every entry in {@code state} matching {@code evictable}, decrementing per removal, and * returns how many went. Self-locking. */ - public static int evictAll( - @Nonnull State state, @Nonnull Predicate evictable) { + public static > int evictAll( + @Nonnull State state, @Strategy @Nonnull Predicate evictable) { synchronized (getTableWriteLock(state)) { return state.sizeManager.evictAll(state.buckets, evictable); } @@ -931,7 +1305,7 @@ public static int evictAll( * reflective allocation or runtime type checks; it only lets the compiler infer {@code TEntry}. */ @Nonnull - public static AtomicReferenceArray createFixedBuckets( + public static > AtomicReferenceArray createFixedBuckets( @Nonnull Class entryClass, int capacity) { return new AtomicReferenceArray<>(sizeFor(capacity)); } @@ -1026,14 +1400,14 @@ public static int bucketIndex(@Nonnull AtomicReferenceArray buckets, long key * bucket. */ @Nullable - public static TEntry bucketFor( + public static > TEntry bucketFor( @Nonnull AtomicReferenceArray buckets, long keyHash) { return buckets.get(bucketIndex(buckets, keyHash)); } /** {@link #bucketFor(AtomicReferenceArray, long)} over a {@link State}. */ @Nullable - public static TEntry bucketFor( + public static > TEntry bucketFor( @Nonnull State state, long keyHash) { return bucketFor(state.buckets, keyHash); } @@ -1045,17 +1419,86 @@ public static TEntry bucketFor( * overload of it. */ @Nullable - public static TEntry bucketAt( + public static > TEntry bucketAt( @Nonnull AtomicReferenceArray buckets, int index) { return buckets.get(index); } /** {@link #bucketAt(AtomicReferenceArray, int)} over a {@link State}. */ @Nullable - public static TEntry bucketAt(@Nonnull State state, int index) { + public static > TEntry bucketAt( + @Nonnull State state, int index) { return bucketAt(state.buckets, index); } + /** + * Returns a lock-free iterator over the candidates for {@code keyHash}: entries in the bucket + * chain that {@code keyHash} maps to (starting from {@link #bucketFor(AtomicReferenceArray, + * long)}) whose own {@link Entry#keyHash} equals it, skipping any other entry sharing the same + * bucket via hash collision on {@link #bucketIndex}. Callers only need a {@code matches} check + * against the entries this yields, not a {@code keyHash} check of their own. + * + *

    Each step follows {@link Entry#next()}, so the iterator reflects entries linked at the time + * each step runs rather than a point-in-time snapshot -- entries inserted ahead of the iterator's + * current position after iteration starts may or may not be observed, and a concurrently removed + * entry remains reachable because {@code unlink()} deliberately retains its {@code next} link for + * in-flight readers. + */ + @Nonnull + public static > Iterator hashIterator( + @Nonnull AtomicReferenceArray buckets, long keyHash) { + return new Iterator() { + private TEntry next = advance(bucketFor(buckets, keyHash)); + + private TEntry advance(TEntry candidate) { + while (candidate != null && candidate.keyHash != keyHash) { + candidate = candidate.next(); + } + return candidate; + } + + @Override + public boolean hasNext() { + return next != null; + } + + @Override + public TEntry next() { + TEntry current = next; + if (current == null) { + throw new NoSuchElementException(); + } + next = advance(current.next()); + return current; + } + }; + } + + /** {@link #hashIterator(AtomicReferenceArray, long)} over a {@link State}. */ + @Nonnull + public static > Iterator hashIterator( + @Nonnull State state, long keyHash) { + return hashIterator(state.buckets, keyHash); + } + + /** + * {@link Iterable} wrapper around {@link #hashIterator(AtomicReferenceArray, long)}, for callers + * that want a plain for-each loop over the candidates for {@code keyHash} rather than driving the + * {@link Iterator} by hand. + */ + @Nonnull + public static > Iterable hashIterable( + @Nonnull AtomicReferenceArray buckets, long keyHash) { + return () -> hashIterator(buckets, keyHash); + } + + /** {@link #hashIterable(AtomicReferenceArray, long)} over a {@link State}. */ + @Nonnull + public static > Iterable hashIterable( + @Nonnull State state, long keyHash) { + return hashIterable(state.buckets, keyHash); + } + /** * Publishes {@code entry} as the head of bucket {@code index}. The helper writes the entry's * {@code next} link before the volatile {@link AtomicReferenceArray#set}; a volatile bucket read @@ -1066,7 +1509,7 @@ public static TEntry bucketAt(@Nonnull State stat * retains its {@code next} link for readers already traversing that chain. */ @GuardedBy("getWriteLockAt(buckets, index)") - public static void insertHeadEntryAt( + public static > void insertHeadEntryAt( @Nonnull AtomicReferenceArray buckets, int index, @Nonnull TEntry entry) { assert Thread.holdsLock(getWriteLockAt(buckets, index)) : "insertHeadEntryAt called without holding getWriteLockAt(buckets, index)"; @@ -1080,7 +1523,7 @@ public static void insertHeadEntryAt( /** {@link #insertHeadEntryAt(AtomicReferenceArray, int, Entry)} over a {@link State}. */ @GuardedBy("getWriteLockAt(state, index)") - public static void insertHeadEntryAt( + public static > void insertHeadEntryAt( @Nonnull State state, int index, @Nonnull TEntry entry) { insertHeadEntryAt(state.buckets, index, entry); } @@ -1091,7 +1534,7 @@ public static void insertHeadEntryAt( * getOrCreate} that reuses it across the lock-free pre-check). */ @GuardedBy("getWriteLock(buckets, keyHash)") - public static void insertHeadEntryFor( + public static > void insertHeadEntryFor( @Nonnull AtomicReferenceArray buckets, long keyHash, @Nonnull TEntry entry) { insertHeadEntryAt(buckets, bucketIndex(buckets, keyHash), entry); } @@ -1105,7 +1548,7 @@ public static void insertHeadEntryFor( * reservations. */ @GuardedBy("getTableWriteLock(state)") - public static void insertReserved( + public static > void insertReserved( @Nonnull State state, long keyHash, @Nonnull TEntry entry) { insertHeadEntryFor(state.buckets, keyHash, entry); } @@ -1120,7 +1563,7 @@ public static void insertReserved( * Does not touch size accounting. */ @GuardedBy("getWriteLockAt(buckets, index)") - public static void unlink( + public static > void unlink( @Nonnull AtomicReferenceArray buckets, int index, @Nullable TEntry prev, @@ -1137,7 +1580,7 @@ public static void unlink( /** {@link #unlink(AtomicReferenceArray, int, Entry, Entry)} over a {@link State}. */ @GuardedBy("getWriteLockAt(state, index)") - public static void unlink( + public static > void unlink( @Nonnull State state, int index, @Nullable TEntry prev, @Nonnull TEntry entry) { unlink(state.buckets, index, prev, entry); } @@ -1148,10 +1591,11 @@ public static void unlink( * predicate sees a stable table and concurrent writers are excluded; lock-free readers continue * throughout. */ - public static boolean removeIf( + @StrategyConsumer + public static > boolean removeIf( @Nonnull AtomicReferenceArray buckets, @Nonnull AtomicInteger size, - @Nonnull Predicate predicate) { + @Strategy @Nonnull Predicate predicate) { synchronized (getTableWriteLock(buckets)) { boolean removed = false; for (int i = 0; i < buckets.length(); i++) { @@ -1176,8 +1620,9 @@ public static boolean removeIf( * occupancy with a {@link State} instead of a bare counter — used by {@link D1#removeIf} and * {@link D2#removeIf}. */ - public static boolean removeIf( - @Nonnull State state, @Nonnull Predicate predicate) { + @StrategyConsumer + public static > boolean removeIf( + @Nonnull State state, @Strategy @Nonnull Predicate predicate) { AtomicReferenceArray buckets = state.buckets; synchronized (getTableWriteLock(state)) { boolean removed = false; @@ -1211,9 +1656,9 @@ public static boolean removeIf( * @param buckets bucket array to drain * @param drainedEntryConsumer action invoked for each entry after its bucket is detached */ - public static void drain( + public static > void drain( @Nonnull AtomicReferenceArray buckets, - @Nonnull Consumer drainedEntryConsumer) { + @Strategy @Nonnull Consumer drainedEntryConsumer) { drainCounting(buckets, drainedEntryConsumer); } @@ -1221,9 +1666,10 @@ public static void drain( * {@link #drain(AtomicReferenceArray, Consumer)} returning the number of entries passed to {@code * drainedEntryConsumer} for size accounting. */ - private static int drainCounting( + @StrategyConsumer + private static > int drainCounting( @Nonnull AtomicReferenceArray buckets, - @Nonnull Consumer drainedEntryConsumer) { + @Strategy @Nonnull Consumer drainedEntryConsumer) { int removed = 0; synchronized (getTableWriteLock(buckets)) { for (int i = 0; i < buckets.length(); i++) { @@ -1256,18 +1702,19 @@ private static int drainCounting( * @param context context passed to each invocation of {@code drainedEntryConsumer} * @param drainedEntryConsumer action invoked with the context and each removed entry */ - public static void drain( + public static > void drain( @Nonnull AtomicReferenceArray buckets, C context, - @Nonnull BiConsumer drainedEntryConsumer) { + @Strategy @Nonnull BiConsumer drainedEntryConsumer) { drainCounting(buckets, context, drainedEntryConsumer); } /** {@link #drainCounting(AtomicReferenceArray, Consumer)}, context-passing form. */ - private static int drainCounting( + @StrategyConsumer + private static > int drainCounting( @Nonnull AtomicReferenceArray buckets, C context, - @Nonnull BiConsumer drainedEntryConsumer) { + @Strategy @Nonnull BiConsumer drainedEntryConsumer) { int removed = 0; synchronized (getTableWriteLock(buckets)) { for (int i = 0; i < buckets.length(); i++) { @@ -1301,8 +1748,9 @@ private static int drainCounting( * @param state table state to drain * @param drainedEntryConsumer action invoked for each entry after its bucket is detached */ - public static void drain( - @Nonnull State state, @Nonnull Consumer drainedEntryConsumer) { + public static > void drain( + @Nonnull State state, + @Strategy @Nonnull Consumer drainedEntryConsumer) { synchronized (getTableWriteLock(state)) { state.sizeManager.release(drainCounting(state.buckets, drainedEntryConsumer)); } @@ -1326,10 +1774,10 @@ public static void drain( * @param context context passed to each invocation of {@code drainedEntryConsumer} * @param drainedEntryConsumer action invoked with the context and each removed entry */ - public static void drain( + public static > void drain( @Nonnull State state, C context, - @Nonnull BiConsumer drainedEntryConsumer) { + @Strategy @Nonnull BiConsumer drainedEntryConsumer) { synchronized (getTableWriteLock(state)) { state.sizeManager.release(drainCounting(state.buckets, context, drainedEntryConsumer)); } @@ -1351,16 +1799,16 @@ public static void clear(@Nonnull AtomicReferenceArray buckets) { * rather than O(buckets); clear is a rare, whole-table operation, so the walk is affordable and * keeping the count honest is worth more than the constant. */ - private static int clearCounting(@Nonnull AtomicReferenceArray buckets) { + private static int clearCounting(@Nonnull AtomicReferenceArray> buckets) { int removed = 0; synchronized (getTableWriteLock(buckets)) { for (int i = 0; i < buckets.length(); i++) { - Entry head = buckets.get(i); + Entry head = buckets.get(i); if (head == null) { continue; } buckets.set(i, null); - for (Entry e = head; e != null; e = e.next()) { + for (Entry e = head; e != null; e = e.next()) { removed++; } } @@ -1377,8 +1825,10 @@ public static void clear(@Nonnull State state) { } } - public static void forEach( - @Nonnull AtomicReferenceArray buckets, @Nonnull Consumer consumer) { + @StrategyConsumer + public static > void forEach( + @Nonnull AtomicReferenceArray buckets, + @Strategy @Nonnull Consumer consumer) { for (int i = 0; i < buckets.length(); i++) { for (TEntry curEntry = buckets.get(i); curEntry != null; curEntry = curEntry.next()) { consumer.accept(curEntry); @@ -1386,10 +1836,11 @@ public static void forEach( } } - public static void forEach( + @StrategyConsumer + public static > void forEach( @Nonnull AtomicReferenceArray buckets, C context, - @Nonnull BiConsumer consumer) { + @Strategy @Nonnull BiConsumer consumer) { for (int i = 0; i < buckets.length(); i++) { for (TEntry curEntry = buckets.get(i); curEntry != null; curEntry = curEntry.next()) { consumer.accept(context, curEntry); @@ -1398,16 +1849,16 @@ public static void forEach( } /** {@link #forEach(AtomicReferenceArray, Consumer)} over a {@link State}. */ - public static void forEach( - @Nonnull State state, @Nonnull Consumer consumer) { + public static > void forEach( + @Nonnull State state, @Strategy @Nonnull Consumer consumer) { forEach(state.buckets, consumer); } /** {@link #forEach(AtomicReferenceArray, Object, BiConsumer)} over a {@link State}. */ - public static void forEach( + public static > void forEach( @Nonnull State state, C context, - @Nonnull BiConsumer consumer) { + @Strategy @Nonnull BiConsumer consumer) { forEach(state.buckets, context, consumer); } } diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableD1Test.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableD1Test.java index 540eb036c2f..46963236366 100644 --- a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableD1Test.java +++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableD1Test.java @@ -80,7 +80,7 @@ void forEachVisitsAllEntries() { table.tryGetOrCreateOrNull("b", k -> new StringEntry(k, 2)); table.tryGetOrCreateOrNull("c", k -> new StringEntry(k, 3)); Set seen = new HashSet<>(); - table.forEach(e -> seen.add(e.key)); + table.forEach(e -> seen.add(e.key())); assertEquals(3, seen.size()); assertTrue(seen.contains("a")); assertTrue(seen.contains("b")); @@ -94,7 +94,7 @@ void forEachWithContextPassesContext() { table.tryGetOrCreateOrNull("x", k -> new StringEntry(k, 10)); table.tryGetOrCreateOrNull("y", k -> new StringEntry(k, 20)); Set seen = new HashSet<>(); - table.forEach(seen, (ctx, e) -> ctx.add(e.key)); + table.forEach(seen, (ctx, e) -> ctx.add(e.key())); assertEquals(2, seen.size()); assertTrue(seen.contains("x")); assertTrue(seen.contains("y")); @@ -258,7 +258,7 @@ void removeIfRemovesMatchingEntries() { assertTrue(removed); assertEquals(5, table.size()); Set seen = new HashSet<>(); - table.forEach(e -> seen.add(e.key)); + table.forEach(e -> seen.add(e.key())); assertEquals(5, seen.size()); for (String key : seen) { assertNotNull(table.get(key)); @@ -301,7 +301,7 @@ void drainRemovesEveryEntryAndFeedsSink() { int[] sum = {0}; table.drain( e -> { - drained.add(e.key); + drained.add(e.key()); sum[0] += e.value; }); @@ -323,7 +323,7 @@ void drainWithContextFeedsSink() { table.tryGetOrCreateOrNull("b", k -> new StringEntry(k, 2)); Set drained = new HashSet<>(); - table.drain(drained, (ctx, e) -> ctx.add(e.key)); + table.drain(drained, (ctx, e) -> ctx.add(e.key())); assertEquals(new HashSet<>(Arrays.asList("a", "b")), drained); assertEquals(0, table.size()); @@ -421,7 +421,7 @@ void tryGetOrCreateOrEvictEvictsWhenFullAndInsertsNewEntry() { Maybe created = table.tryGetOrCreateOrEvict("new", k -> new StringEntry(k, 2), e -> true); assertTrue(created.isPresent()); - assertEquals("new", created.getOrNull().key); + assertEquals("new", created.getOrNull().key()); assertEquals(1, table.size()); assertNull(table.get("old")); assertSame(created.getOrNull(), table.get("new")); @@ -464,8 +464,9 @@ void tryGetOrCreateOrEvictOrNullEvictionRunsBeforeThrowingCreator() { assertNull(table.get("new")); } + /** Entry holding a key plus one mutable {@code int} payload. */ private static final class StringEntry extends ConcurrentHashtable.D1.Entry { - final int value; + volatile int value; StringEntry(String key, int value) { super(key); @@ -473,6 +474,13 @@ private static final class StringEntry extends ConcurrentHashtable.D1.Entry { + CollidingEntry(CollidingKey key) { + super(key); + } + } + /** Key with a fixed hashCode to force deterministic bucket placement. */ private static final class CollidingKey { final String label; @@ -497,10 +505,4 @@ public boolean equals(Object o) { return fixedHash == that.fixedHash && label.equals(that.label); } } - - private static final class CollidingEntry extends ConcurrentHashtable.D1.Entry { - CollidingEntry(CollidingKey key) { - super(key); - } - } } diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableD2Test.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableD2Test.java index 182cb4c4f25..e2b3f8dd353 100644 --- a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableD2Test.java +++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableD2Test.java @@ -45,8 +45,8 @@ void getOrCreateOnMissBuildsEntryViaCreator() { return new PairEntry(k1, k2); }); assertNotNull(created); - assertEquals("a", created.key1); - assertEquals(Integer.valueOf(1), created.key2); + assertEquals("a", created.key1()); + assertEquals(Integer.valueOf(1), created.key2()); assertEquals(1, table.size()); assertEquals(1, createCount[0]); assertSame(created, table.get("a", 1)); @@ -78,7 +78,7 @@ void forEachVisitsBothPairs() { table.tryGetOrCreateOrNull("a", 1, PairEntry::new); table.tryGetOrCreateOrNull("b", 2, PairEntry::new); Set seen = new HashSet<>(); - table.forEach(e -> seen.add(e.key1 + ":" + e.key2)); + table.forEach(e -> seen.add(e.key1() + ":" + e.key2())); assertEquals(2, seen.size()); assertTrue(seen.contains("a:1")); assertTrue(seen.contains("b:2")); @@ -91,7 +91,7 @@ void forEachWithContextPassesContextToConsumer() { table.tryGetOrCreateOrNull("a", 1, PairEntry::new); table.tryGetOrCreateOrNull("b", 2, PairEntry::new); Set seen = new HashSet<>(); - table.forEach(seen, (ctx, e) -> ctx.add(e.key1 + ":" + e.key2)); + table.forEach(seen, (ctx, e) -> ctx.add(e.key1() + ":" + e.key2())); assertEquals(2, seen.size()); assertTrue(seen.contains("a:1")); assertTrue(seen.contains("b:2")); @@ -246,11 +246,11 @@ void removeIfRemovesMatchingEntries() { for (int i = 0; i < 10; i++) { table.tryGetOrCreateOrNull("k", i, PairEntry::new); } - boolean removed = table.removeIf(e -> e.key2 % 2 == 0); // removes key2 0,2,4,6,8 + boolean removed = table.removeIf(e -> e.key2() % 2 == 0); // removes key2 0,2,4,6,8 assertTrue(removed); assertEquals(5, table.size()); Set seen = new HashSet<>(); - table.forEach(e -> seen.add(e.key1 + ":" + e.key2)); + table.forEach(e -> seen.add(e.key1() + ":" + e.key2())); assertEquals(5, seen.size()); } @@ -286,7 +286,7 @@ void drainRemovesEveryEntryAndFeedsSink() { table.tryGetOrCreateOrNull("b", 1, PairEntry::new); Set drained = new HashSet<>(); - table.drain(e -> drained.add(e.key1 + ":" + e.key2)); + table.drain(e -> drained.add(e.key1() + ":" + e.key2())); assertEquals(new HashSet<>(Arrays.asList("a:1", "a:2", "b:1")), drained); assertEquals(0, table.size()); @@ -304,7 +304,7 @@ void drainWithContextFeedsSink() { table.tryGetOrCreateOrNull("b", 2, PairEntry::new); Set drained = new HashSet<>(); - table.drain(drained, (ctx, e) -> ctx.add(e.key1 + ":" + e.key2)); + table.drain(drained, (ctx, e) -> ctx.add(e.key1() + ":" + e.key2())); assertEquals(new HashSet<>(Arrays.asList("a:1", "b:2")), drained); assertEquals(0, table.size()); @@ -348,7 +348,7 @@ void tryGetOrCreateOrEvictEvictsWhenFullAndInsertsNewEntry() { Maybe created = table.tryGetOrCreateOrEvict("new", 2, PairEntry::new, e -> true); assertTrue(created.isPresent()); - assertEquals("new", created.getOrNull().key1); + assertEquals("new", created.getOrNull().key1()); assertEquals(1, table.size()); assertNull(table.get("old", 1)); assertSame(created.getOrNull(), table.get("new", 2)); @@ -391,6 +391,7 @@ void tryGetOrCreateOrEvictOrNullEvictionRunsBeforeThrowingCreator() { assertNull(table.get("new", 2)); } + /** Entry with no payload beyond its two key parts, used to exercise the D2 identity/API. */ private static final class PairEntry extends ConcurrentHashtable.D2.Entry { PairEntry(String key1, Integer key2) { super(key1, key2); diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java new file mode 100644 index 00000000000..f61876b2e66 --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java @@ -0,0 +1,184 @@ +package datadog.trace.util; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nonnull; +import org.junit.jupiter.api.Test; + +/** Exercises {@link ConcurrentHashtable#tryReserve} and {@link ConcurrentHashtable.Reservation}. */ +class ConcurrentHashtableReservationTest { + + private static final class TestEntry extends ConcurrentHashtable.Entry { + final int value; + + TestEntry(int value) { + super(value); + this.value = value; + } + + @Override + public boolean matches(@Nonnull TestEntry other) { + return value == other.value; + } + } + + @Test + void tryGetOrInsertOrNullInsertsOnMissAndFindsOnHit() { + ConcurrentHashtable.State state = + ConcurrentHashtable.createBounded(TestEntry.class, 4); + + TestEntry first; + try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) { + assertTrue(r.isReserved()); + first = r.tryGetOrInsertOrNull(1, TestEntry::new); + } + assertEquals(1, first.value); + assertEquals(1, ConcurrentHashtable.estimateSize(state)); + + // Reserving again for a key that already exists should discard the reservation and return the + // existing entry, not double-insert or leak the claimed slot. + TestEntry second; + try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) { + second = r.tryGetOrInsertOrNull(1, TestEntry::new); + } + assertSame(first, second); + assertEquals(1, ConcurrentHashtable.estimateSize(state)); + } + + @Test + void reserveOnFullTableIsAbsentAndSkipsTheFactory() { + ConcurrentHashtable.State state = + ConcurrentHashtable.createBounded(TestEntry.class, 1); + try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) { + r.tryGetOrInsertOrNull(1, TestEntry::new); + } + assertTrue(ConcurrentHashtable.isFull(state)); + + AtomicInteger factoryCalls = new AtomicInteger(); + TestEntry result; + try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) { + assertFalse(r.isReserved()); + result = + r.tryGetOrInsertOrNull( + 2, + v -> { + factoryCalls.incrementAndGet(); + return new TestEntry(v); + }); + } + assertNull(result); + assertEquals(0, factoryCalls.get()); + assertEquals(1, ConcurrentHashtable.estimateSize(state)); + } + + @Test + void closeCancelsAnUnconsumedReservation() { + ConcurrentHashtable.State state = + ConcurrentHashtable.createBounded(TestEntry.class, 1); + try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) { + assertTrue(r.isReserved()); + // Deliberately not consuming the reservation. + } + assertEquals(0, ConcurrentHashtable.estimateSize(state)); + assertFalse(ConcurrentHashtable.isFull(state)); + } + + @Test + void tryGetOrInsertWrapsResultInMaybe() { + ConcurrentHashtable.State state = + ConcurrentHashtable.createBounded(TestEntry.class, 1); + + Maybe present; + try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) { + present = r.tryGetOrInsert(1, TestEntry::new); + } + assertTrue(present.isPresent()); + assertEquals(1, present.getOrNull().value); + + Maybe absent; + try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) { + absent = r.tryGetOrInsert(2, TestEntry::new); + } + assertFalse(absent.isPresent()); + assertNull(absent.getOrNull()); + } + + @Test + void closeOnAnAbsentReservationIsANoOp() { + ConcurrentHashtable.State state = + ConcurrentHashtable.createBounded(TestEntry.class, 0); + try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) { + assertFalse(r.isReserved()); + } + assertEquals(0, ConcurrentHashtable.estimateSize(state)); + } + + private static final class ThreePartEntry extends ConcurrentHashtable.Entry { + final String a; + final String b; + final String c; + + ThreePartEntry(String a, String b, String c) { + super(HashingUtils.hash(a, b, c)); + this.a = a; + this.b = b; + this.c = c; + } + + @Override + public boolean matches(@Nonnull ThreePartEntry other) { + return a.equals(other.a) && b.equals(other.b) && c.equals(other.c); + } + } + + private static final class FourPartEntry extends ConcurrentHashtable.Entry { + final String a; + final String b; + final String c; + final String d; + + FourPartEntry(String a, String b, String c, String d) { + super(HashingUtils.hash(a, b, c, d)); + this.a = a; + this.b = b; + this.c = c; + this.d = d; + } + + @Override + public boolean matches(@Nonnull FourPartEntry other) { + return a.equals(other.a) && b.equals(other.b) && c.equals(other.c) && d.equals(other.d); + } + } + + @Test + void tryGetOrInsertOrNullSupportsUpToFourComponents() { + ConcurrentHashtable.State state3 = + ConcurrentHashtable.createBounded(ThreePartEntry.class, 2); + ThreePartEntry three; + try (ConcurrentHashtable.Reservation r = + ConcurrentHashtable.tryReserve(state3)) { + three = r.tryGetOrInsertOrNull("x", "y", "z", ThreePartEntry::new); + } + assertEquals("x", three.a); + assertEquals("y", three.b); + assertEquals("z", three.c); + + ConcurrentHashtable.State state4 = + ConcurrentHashtable.createBounded(FourPartEntry.class, 2); + FourPartEntry four; + try (ConcurrentHashtable.Reservation r = + ConcurrentHashtable.tryReserve(state4)) { + four = r.tryGetOrInsertOrNull("w", "x", "y", "z", FourPartEntry::new); + } + assertEquals("w", four.a); + assertEquals("x", four.b); + assertEquals("y", four.c); + assertEquals("z", four.d); + } +} diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableSizeManagerTest.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableSizeManagerTest.java index 24c581b2380..2fceb4ce3ea 100644 --- a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableSizeManagerTest.java +++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableSizeManagerTest.java @@ -40,7 +40,7 @@ void tryReserveSucceedsUnderCapacityAndFailsWhenFull() { @Test void tryReserveOrEvictReservesDirectlyWhenUnderCapacity() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 2); + ConcurrentHashtable.createBounded(TestEntry.class, 2); boolean reserved = tryReserveOrEvict(state, e -> true); assertTrue(reserved); @@ -51,7 +51,7 @@ void tryReserveOrEvictReservesDirectlyWhenUnderCapacity() { @Test void tryReserveOrEvictEvictsWhenFullAndSomethingMatches() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 1); + ConcurrentHashtable.createBounded(TestEntry.class, 1); TestEntry existing = insertAt(state, 0, "existing"); assertTrue(state.sizeManager.tryReserve()); assertTrue(state.sizeManager.isFull()); @@ -65,7 +65,7 @@ void tryReserveOrEvictEvictsWhenFullAndSomethingMatches() { @Test void tryReserveOrEvictFailsAndLeavesTableUntouchedWhenNothingEvictable() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 1); + ConcurrentHashtable.createBounded(TestEntry.class, 1); TestEntry existing = insertAt(state, 0, "existing"); assertTrue(state.sizeManager.tryReserve()); @@ -78,7 +78,7 @@ void tryReserveOrEvictFailsAndLeavesTableUntouchedWhenNothingEvictable() { @Test void insertReservedSplicesWithoutTouchingTheCountAfterATryReserve() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 2); + ConcurrentHashtable.createBounded(TestEntry.class, 2); assertTrue(state.sizeManager.tryReserve()); assertEquals(1, state.sizeManager.estimateSize()); @@ -95,7 +95,7 @@ void insertReservedSplicesWithoutTouchingTheCountAfterATryReserve() { @Test void evictOneReturnsNullAndLeavesCountUnchangedWhenNothingMatches() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 4); + ConcurrentHashtable.createBounded(TestEntry.class, 4); insertAt(state, 0, "a"); state.sizeManager.increment(); @@ -107,7 +107,7 @@ void evictOneReturnsNullAndLeavesCountUnchangedWhenNothingMatches() { @Test void evictOneUnlinksMatchAndDecrementsCount() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 4); + ConcurrentHashtable.createBounded(TestEntry.class, 4); TestEntry a = insertAt(state, 0, "a"); TestEntry b = insertAt(state, 1, "b"); state.sizeManager.increment(); @@ -130,7 +130,7 @@ void evictOneUnlinksMatchAndDecrementsCount() { void evictOneResumesFromLastEvictedBucketAndWrapsAround() { // Bucket-array length 4: keyHash i lands in bucket i. ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 4); + ConcurrentHashtable.createBounded(TestEntry.class, 4); TestEntry e0 = insertAt(state, 0, "e0"); insertAt(state, 2, "e2"); TestEntry e3 = insertAt(state, 3, "e3"); @@ -159,7 +159,7 @@ void evictOneResumesFromLastEvictedBucketAndWrapsAround() { @Test void evictAllRemovesEveryMatchAndReturnsCount() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 8); + ConcurrentHashtable.createBounded(TestEntry.class, 8); for (int i = 0; i < 6; i++) { insertAt(state, i, "e" + i); state.sizeManager.increment(); @@ -179,7 +179,7 @@ void evictAllRemovesEveryMatchAndReturnsCount() { @Test void evictAllResetsCursorSoSubsequentEvictOneScansFromBucketZero() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 4); + ConcurrentHashtable.createBounded(TestEntry.class, 4); insertAt(state, 2, "a"); state.sizeManager.increment(); // Advance the cursor away from 0 via a successful eviction at bucket 2. @@ -203,7 +203,7 @@ void evictAllResetsCursorSoSubsequentEvictOneScansFromBucketZero() { @Test void releaseGivesBackRemovedSlotsAndRestartsScan() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 4); + ConcurrentHashtable.createBounded(TestEntry.class, 4); insertAt(state, 2, "a"); state.sizeManager.increment(); evictOne(state, e -> true); // advances the cursor to 2, count back to 0 @@ -226,7 +226,7 @@ void releaseGivesBackRemovedSlotsAndRestartsScan() { @Test void stateCreateCappedBundlesBucketsAndSizeManager() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 3); + ConcurrentHashtable.createBounded(TestEntry.class, 3); assertEquals(0, state.sizeManager.estimateSize()); assertEquals(3, state.sizeManager.capacity()); assertTrue(state.buckets.length() >= 3); @@ -235,7 +235,7 @@ void stateCreateCappedBundlesBucketsAndSizeManager() { @Test void stateLevelTryReserveOrEvictAndEvictOneAndEvictAllDelegateToSizeManager() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 1); + ConcurrentHashtable.createBounded(TestEntry.class, 1); synchronized (ConcurrentHashtable.getWriteLockAt(state, 0)) { ConcurrentHashtable.insertHeadEntryAt(state, 0, new TestEntry(0, "a")); } @@ -280,7 +280,7 @@ void stateLevelTryReserveOrEvictAndEvictOneAndEvictAllDelegateToSizeManager() { @Test void clearCannotInterleaveBetweenReservationAndInsert() throws InterruptedException { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 1); + ConcurrentHashtable.createBounded(TestEntry.class, 1); insertAt(state, 0, "a"); state.sizeManager.increment(); assertTrue(ConcurrentHashtable.isFull(state)); @@ -318,7 +318,7 @@ void clearCannotInterleaveBetweenReservationAndInsert() throws InterruptedExcept @Test void reservationSurvivesAClearLandingBetweenReserveAndInsert() { ConcurrentHashtable.State state = - ConcurrentHashtable.State.createBounded(TestEntry.class, 2); + ConcurrentHashtable.createBounded(TestEntry.class, 2); insertAt(state, 0, "a"); state.sizeManager.increment(); @@ -420,12 +420,17 @@ private static int evictAll( } /** Entry with a caller-controlled {@code keyHash} so tests can place it in an exact bucket. */ - private static final class TestEntry extends ConcurrentHashtable.Entry { + private static final class TestEntry extends ConcurrentHashtable.Entry { final String label; TestEntry(long keyHash, String label) { super(keyHash); this.label = label; } + + @Override + public boolean matches(TestEntry other) { + return label.equals(other.label); + } } } diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableStaticsTest.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableStaticsTest.java index 8191e396e14..8be6e791b9e 100644 --- a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableStaticsTest.java +++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableStaticsTest.java @@ -329,7 +329,7 @@ private static boolean assertionsEnabled() { } /** Primitive-{@code int}-key entry: no boxing, keyHash is the key itself. */ - private static final class IntEntry extends ConcurrentHashtable.Entry { + private static final class IntEntry extends ConcurrentHashtable.Entry { final int key; final int value; @@ -342,6 +342,11 @@ private static final class IntEntry extends ConcurrentHashtable.Entry { boolean matches(int key) { return this.key == key; } + + @Override + public boolean matches(IntEntry other) { + return matches(other.key); + } } /** From b8b503706ebaf5a9e0924c792998cef4afc27838 Mon Sep 17 00:00:00 2001 From: Douglas Q Hawkins Date: Fri, 11 Sep 2026 08:25:24 -0400 Subject: [PATCH 02/13] Move RawLogMessage equality logic into matches(), have equals() delegate equals() now only performs the type check and defers the substantive comparison to matches(), avoiding the duplicated logic. Co-Authored-By: Claude Sonnet 5 --- .../trace/api/telemetry/LogCollector.java | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java index 160cdd31c12..d5c1113a737 100644 --- a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java +++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java @@ -243,15 +243,6 @@ private void snapshotCount() { @Override public boolean matches(RawLogMessage that) { - return equals(that); - } - - @Override - public boolean equals(Object o) { - if (this == o) return true; - if (o == null || getClass() != o.getClass()) return false; - RawLogMessage that = (RawLogMessage) o; - if (!Objects.equals(logLevel, that.logLevel)) return false; if (!Objects.equals(message, that.message)) return false; @@ -271,6 +262,13 @@ public boolean equals(Object o) { } } + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + return matches((RawLogMessage) o); + } + @Override public int hashCode() { return (int) keyHash; From fe622b3a8e503907f40a86d795d86d3f60f907b5 Mon Sep 17 00:00:00 2001 From: Douglas Q Hawkins Date: Fri, 11 Sep 2026 08:34:20 -0400 Subject: [PATCH 03/13] Simplify LogCollector's find-or-insert and fix a missed increment Use ConcurrentHashtable.hashIterable() in find() instead of a manual bucket walk with a keyHash pre-check, and drop the redundant isFull() fast-reject and the second find() done just before taking a reservation -- Reservation#tryGetOrInsertOrNull already performs that same locked comparison via matches(), so repeating it added nothing. That simplification exposed a real bug: tryGetOrInsertOrNull can return an existing match instead of the newly built entry, and that occurrence was never counted. RawLogMessage's live occurrence count now starts at zero so every find-or-insert can unconditionally increment() the returned entry, whether it's the new instance or an existing match. Co-Authored-By: Claude Sonnet 5 --- .../trace/api/telemetry/LogCollector.java | 83 +++++++------------ 1 file changed, 29 insertions(+), 54 deletions(-) diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java index d5c1113a737..ca812ceb54b 100644 --- a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java +++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java @@ -1,10 +1,8 @@ package datadog.trace.api.telemetry; -import static datadog.trace.util.ConcurrentHashtable.bucketAt; -import static datadog.trace.util.ConcurrentHashtable.bucketIndex; import static datadog.trace.util.ConcurrentHashtable.estimateSize; import static datadog.trace.util.ConcurrentHashtable.getTableWriteLock; -import static datadog.trace.util.ConcurrentHashtable.isFull; +import static datadog.trace.util.ConcurrentHashtable.hashIterable; import static datadog.trace.util.LongHashingUtils.hash; import datadog.trace.api.internal.VisibleForTesting; @@ -59,50 +57,33 @@ public void addLogMessage(String logLevel, String message, @Nullable Throwable t public void addLogMessage( String logLevel, String message, @Nullable Throwable throwable, @Nullable String tags) { long keyHash = RawLogMessage.computeHash(logLevel, message, throwable); - int bucketIndex = bucketIndex(rawLogMessages.buckets, keyHash); - // Fast path for duplicates: search the target bucket without locking. - RawLogMessage rawLogMessage = find(bucketIndex, keyHash, logLevel, message, throwable); + // Fast path for duplicates: search lock-free before ever taking the table lock. + RawLogMessage rawLogMessage = find(keyHash, logLevel, message, throwable); if (rawLogMessage != null) { rawLogMessage.increment(); return; } - // Fast path after a miss for a full table: reject without locking when the target bucket is - // populated. If the bucket is empty, drain() may have detached it before releasing capacity, - // so continue to the locked capacity check. - if (isFull(rawLogMessages) && bucketAt(rawLogMessages, bucketIndex) != null) { - // Mitigate a race where another writer could claim the bucket before the previous find - rawLogMessage = find(bucketIndex, keyHash, logLevel, message, throwable); - if (rawLogMessage != null) { - rawLogMessage.increment(); - return; - } - // TODO: We could emit a metric for dropped logs. - return; - } - - // Slow path after a miss: repeat the lookup under the table write lock because another writer - // or drain() may have changed the table. Taking the reservation under this same lock, rather - // than lock-free ahead of it, keeps the whole find-or-insert decision serialized with other - // writers and with drain() -- a writer here genuinely waits for an in-progress drain instead - // of being told, incorrectly, that the table is full because of a duplicate reservation that - // hasn't cancelled yet. Serialized this way, the reservation can never lose a race to insert, - // so there's no existing-vs-new distinction to make: it always inserts fresh. + // Slow path after a miss: take the reservation under the table write lock, rather than + // lock-free ahead of it, so concurrent reservations for the same logical duplicate are + // serialized with each other, with drain(), and with the locked find-or-insert inside + // Reservation#finish() -- a losing reservation cancels immediately instead of transiently + // inflating size and starving a genuinely distinct concurrent insert. finish() does its own + // locked comparison, so there's no need to repeat find() here first. synchronized (getTableWriteLock(rawLogMessages)) { - rawLogMessage = find(bucketIndex, keyHash, logLevel, message, throwable); - if (rawLogMessage != null) { - rawLogMessage.increment(); - return; - } try (Reservation reservation = ConcurrentHashtable.tryReserve(rawLogMessages)) { - if (reservation.isReserved()) { - // Allocate only after the reservation succeeds. - reservation.tryGetOrInsertOrNull( - new RawLogMessage( - logLevel, message, throwable, tags, System.currentTimeMillis() / 1000)); + if (!reservation.isReserved()) { + // TODO: We could emit a metric for dropped logs. + return; } - // TODO: We could emit a metric for dropped logs. + // Built zeroed, so this occurrence can be counted uniformly below whether or not + // tryGetOrInsertOrNull ends up returning this instance or an existing match. + rawLogMessage = + reservation.tryGetOrInsertOrNull( + new RawLogMessage( + logLevel, message, throwable, tags, System.currentTimeMillis() / 1000)); + rawLogMessage.increment(); } } } @@ -136,14 +117,13 @@ public Collection drain() { /** * Finds a log group with the same level, message, and - * throwable in the selected bucket. + * throwable as {@code keyHash}'s candidates. * *

    Note, throwables are matched by identity or by class and stack trace. * *

    The bucket chain supports lock-free reads. A caller that inserts after a miss must repeat * the search under the table write lock. * - * @param bucketIndex bucket selected for {@code keyHash} * @param keyHash precomputed hash of the level, message, and throwable class * @param logLevel log level to match * @param message message to match @@ -152,19 +132,10 @@ public Collection drain() { */ @Nullable private RawLogMessage find( - int bucketIndex, - long keyHash, - String logLevel, - String message, - @Nullable Throwable throwable) { - // Start searching from given bucket, and follow the entry next links + long keyHash, String logLevel, String message, @Nullable Throwable throwable) { StackTraceElement[] stackTrace = null; - for (RawLogMessage entry = bucketAt(rawLogMessages, bucketIndex); - entry != null; - entry = entry.next()) { - if (entry.keyHash != keyHash - || !Objects.equals(logLevel, entry.logLevel) - || !Objects.equals(message, entry.message)) { + for (RawLogMessage entry : hashIterable(rawLogMessages, keyHash)) { + if (!Objects.equals(logLevel, entry.logLevel) || !Objects.equals(message, entry.message)) { continue; } // throwables are more costly to compare, check first the identity @@ -205,8 +176,12 @@ public static final class RawLogMessage extends ConcurrentHashtable.Entry Date: Fri, 11 Sep 2026 09:53:01 -0400 Subject: [PATCH 04/13] Finish ReentrantLock migration for ConcurrentHashtable and shorten drain's lock hold Releases capacity per-entry during drain instead of once at the end, and acquires/releases the table lock per bucket rather than for the whole sweep, so writers can interleave with an in-progress drain instead of blocking for its entire duration. Also cleans up leftover synchronized(...) usage in ConcurrentHashtableSizeManagerTest and simplifies LogCollector.addLogMessage now that Reservation holds the lock itself. Co-Authored-By: Claude Sonnet 5 --- .../util/ThreadSafeMapCounterBenchmark.java | 29 +- .../trace/util/ThreadSafeMapD1Benchmark.java | 28 +- .../trace/api/telemetry/LogCollector.java | 38 +- .../trace/util/ConcurrentHashtable.java | 336 +++++++++++++----- .../trace/api/telemetry/LogCollectorTest.java | 5 +- .../ConcurrentHashtableSizeManagerTest.java | 64 +++- 6 files changed, 356 insertions(+), 144 deletions(-) diff --git a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapCounterBenchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapCounterBenchmark.java index 239dd3890f4..bc8bc07b9f5 100644 --- a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapCounterBenchmark.java +++ b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapCounterBenchmark.java @@ -23,10 +23,11 @@ * Measures lookup followed by an atomic counter increment in a shared, pre-populated table. Models * per-class or per-method hit counters in the tracer. * - *

    The {@link ConcurrentHashtable.D1} case embeds a {@code volatile long} counter directly in the - * entry and increments it via {@link AtomicLongFieldUpdater}, avoiding a second heap object. The - * map baselines store a separate {@link AtomicLong} or {@link LongAdder}; {@code LongAdder} spreads - * contention across internal cells at the cost of more memory and a more expensive read. + *

    The {@link ConcurrentHashtable.D1} case embeds a {@code volatile long} in each entry. {@link + * AtomicLongFieldUpdater} updates that field atomically without allocating an {@link AtomicLong} + * per key. The map baselines store a separate {@link AtomicLong} or {@link LongAdder}; {@code + * LongAdder} spreads contention across internal cells at the cost of more memory and a more + * expensive read. * *

    Lookups reuse the key instances installed during setup. {@code Objects.equals} therefore * returns on its identity check without dispatching to {@code equals}, so this measures the @@ -49,8 +50,8 @@ *

  • {@code LongAdder} is marginally faster (79 vs 71 ops/us) because it shards the counter * across cells to reduce CAS contention; the advantage grows with thread count. *
  • {@code ConcurrentHashtable} matches {@code AtomicLong} throughput (69 vs 71 ops/us) while - * embedding the counter directly in the entry via {@code AtomicLongFieldUpdater} — one object - * instead of two, with no throughput penalty. + * embedding the counter directly in the entry — one object instead of two, with no throughput + * penalty. * */ @Fork(2) @@ -72,12 +73,8 @@ public class ThreadSafeMapCounterBenchmark { } } - /** - * Shared state ({@link Scope#Benchmark}): one instance of each map across all threads, modelling - * a shared instrumentation counter table. - */ static final class CounterEntry extends ConcurrentHashtable.D1.Entry { - static final AtomicLongFieldUpdater COUNT = + private static final AtomicLongFieldUpdater COUNT = AtomicLongFieldUpdater.newUpdater(CounterEntry.class, "count"); volatile long count; @@ -85,8 +82,16 @@ static final class CounterEntry extends ConcurrentHashtable.D1.Entry { CounterEntry(String key) { super(key); } + + long increment() { + return COUNT.incrementAndGet(this); + } } + /** + * Shared state ({@link Scope#Benchmark}): one instance of each map across all threads, modelling + * a shared instrumentation counter table. + */ @State(Scope.Benchmark) public static class SharedState { ConcurrentHashtable.D1 table; @@ -120,7 +125,7 @@ int next() { @Benchmark public long increment_concurrentHashtable(SharedState s, ThreadState t) { - return CounterEntry.COUNT.incrementAndGet(s.table.get(KEYS[t.next()])); + return s.table.get(KEYS[t.next()]).increment(); } @Benchmark diff --git a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD1Benchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD1Benchmark.java index 6f5a8dcd769..43d9b849f48 100644 --- a/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD1Benchmark.java +++ b/internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD1Benchmark.java @@ -83,35 +83,33 @@ public class ThreadSafeMapD1Benchmark { } } - /** - * Shared state ({@link Scope#Benchmark}): one instance of each map across all threads, modelling - * a shared instrumentation cache. - */ - static final class LongEntry extends ConcurrentHashtable.D1.Entry { - final long value; + static final class D1Entry extends ConcurrentHashtable.D1.Entry { + volatile long value; - LongEntry(String key, long value) { + D1Entry(String key) { super(key); - this.value = value; } } + /** + * Shared state ({@link Scope#Benchmark}): one instance of each map across all threads, modelling + * a shared instrumentation cache. + */ @State(Scope.Benchmark) public static class SharedState { - ConcurrentHashtable.D1 table; + ConcurrentHashtable.D1 table; ConcurrentHashMap concurrentHashMap; ConcurrentSkipListMap skipListMap; Map synchronizedHashMap; @Setup(Level.Iteration) public void setUp() { - table = ConcurrentHashtable.D1.createBounded(LongEntry.class, CAPACITY); + table = ConcurrentHashtable.D1.createBounded(D1Entry.class, CAPACITY); concurrentHashMap = new ConcurrentHashMap<>(CAPACITY); skipListMap = new ConcurrentSkipListMap<>(); synchronizedHashMap = Collections.synchronizedMap(new HashMap<>(CAPACITY)); for (int i = 0; i < N_KEYS; ++i) { - long value = i; - table.tryGetOrCreateOrNull(KEYS[i], k -> new LongEntry(k, value)); + table.tryGetOrCreateOrNull(KEYS[i], D1Entry::new).value = i; concurrentHashMap.put(KEYS[i], (long) i); skipListMap.put(KEYS[i], (long) i); synchronizedHashMap.put(KEYS[i], (long) i); @@ -132,7 +130,7 @@ int next() { } @Benchmark - public LongEntry get_concurrentHashtable(SharedState s, ThreadState t) { + public D1Entry get_concurrentHashtable(SharedState s, ThreadState t) { return s.table.get(KEYS[t.next()]); } @@ -152,8 +150,8 @@ public Long get_synchronizedHashMap(SharedState s, ThreadState t) { } @Benchmark - public LongEntry getOrCreate_concurrentHashtable(SharedState s, ThreadState t) { - return s.table.tryGetOrCreateOrNull(KEYS[t.next()], k -> new LongEntry(k, 0L)); + public D1Entry getOrCreate_concurrentHashtable(SharedState s, ThreadState t) { + return s.table.tryGetOrCreateOrNull(KEYS[t.next()], D1Entry::new); } /** diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java index ca812ceb54b..875ee129f29 100644 --- a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java +++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java @@ -1,7 +1,6 @@ package datadog.trace.api.telemetry; import static datadog.trace.util.ConcurrentHashtable.estimateSize; -import static datadog.trace.util.ConcurrentHashtable.getTableWriteLock; import static datadog.trace.util.ConcurrentHashtable.hashIterable; import static datadog.trace.util.LongHashingUtils.hash; @@ -64,27 +63,24 @@ public void addLogMessage( return; } - // Slow path after a miss: take the reservation under the table write lock, rather than - // lock-free ahead of it, so concurrent reservations for the same logical duplicate are - // serialized with each other, with drain(), and with the locked find-or-insert inside - // Reservation#finish() -- a losing reservation cancels immediately instead of transiently - // inflating size and starving a genuinely distinct concurrent insert. finish() does its own - // locked comparison, so there's no need to repeat find() here first. - synchronized (getTableWriteLock(rawLogMessages)) { - try (Reservation reservation = - ConcurrentHashtable.tryReserve(rawLogMessages)) { - if (!reservation.isReserved()) { - // TODO: We could emit a metric for dropped logs. - return; - } - // Built zeroed, so this occurrence can be counted uniformly below whether or not - // tryGetOrInsertOrNull ends up returning this instance or an existing match. - rawLogMessage = - reservation.tryGetOrInsertOrNull( - new RawLogMessage( - logLevel, message, throwable, tags, System.currentTimeMillis() / 1000)); - rawLogMessage.increment(); + // Slow path after a miss: tryReserve holds the table write lock for the reservation's whole + // lifetime, so concurrent reservations for the same logical duplicate are serialized with each + // other, with drain(), and with the locked find-or-insert inside Reservation#finish() -- a + // losing reservation cancels immediately instead of transiently inflating size and starving a + // genuinely distinct concurrent insert. finish() does its own locked comparison, so there's no + // need to repeat find() here first. + try (Reservation reservation = ConcurrentHashtable.tryReserve(rawLogMessages)) { + if (!reservation.isReserved()) { + // TODO: We could emit a metric for dropped logs. + return; } + // Built zeroed, so this occurrence can be counted uniformly below whether or not + // tryGetOrInsertOrNull ends up returning this instance or an existing match. + rawLogMessage = + reservation.tryGetOrInsertOrNull( + new RawLogMessage( + logLevel, message, throwable, tags, System.currentTimeMillis() / 1000)); + rawLogMessage.increment(); } } diff --git a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java index 73276ccf594..ea1c58888e7 100644 --- a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java +++ b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java @@ -8,6 +8,7 @@ import java.util.Objects; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReferenceArray; +import java.util.concurrent.locks.ReentrantLock; import java.util.function.BiConsumer; import java.util.function.BiFunction; import java.util.function.Consumer; @@ -269,7 +270,9 @@ public TEntry tryGetOrCreateOrNull( return cast(curEntry); } } - synchronized (getTableWriteLock(state)) { + ReentrantLock lock = getTableWriteLock(state); + lock.lock(); + try { for (Entry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { @@ -288,6 +291,8 @@ public TEntry tryGetOrCreateOrNull( insertHeadEntryAt(state, index, newEntry); state.sizeManager.increment(); return newEntry; + } finally { + lock.unlock(); } } @@ -327,7 +332,9 @@ public TEntry tryGetOrCreateOrEvictOrNull( return cast(curEntry); } } - synchronized (getTableWriteLock(state)) { + ReentrantLock lock = getTableWriteLock(state); + lock.lock(); + try { for (Entry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { @@ -343,6 +350,8 @@ public TEntry tryGetOrCreateOrEvictOrNull( insertHeadEntryAt(state, index, newEntry); state.sizeManager.increment(); return newEntry; + } finally { + lock.unlock(); } } @@ -355,7 +364,9 @@ public TEntry tryGetOrCreateOrEvictOrNull( public TEntry remove(@Nullable K key) { long keyHash = D1.Entry.hash(key); int index = bucketIndex(state.buckets, keyHash); - synchronized (getTableWriteLock(state)) { + ReentrantLock lock = getTableWriteLock(state); + lock.lock(); + try { Entry prev = null; for (Entry curEntry = bucketAt(state, index); curEntry != null; @@ -367,6 +378,8 @@ public TEntry remove(@Nullable K key) { } } return null; + } finally { + lock.unlock(); } } @@ -597,7 +610,9 @@ public TEntry tryGetOrCreateOrNull( return cast(curEntry); } } - synchronized (getTableWriteLock(state)) { + ReentrantLock lock = getTableWriteLock(state); + lock.lock(); + try { for (Entry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { @@ -616,6 +631,8 @@ public TEntry tryGetOrCreateOrNull( insertHeadEntryAt(state, index, newEntry); state.sizeManager.increment(); return newEntry; + } finally { + lock.unlock(); } } @@ -657,7 +674,9 @@ public TEntry tryGetOrCreateOrEvictOrNull( return cast(curEntry); } } - synchronized (getTableWriteLock(state)) { + ReentrantLock lock = getTableWriteLock(state); + lock.lock(); + try { for (Entry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) { @@ -673,6 +692,8 @@ public TEntry tryGetOrCreateOrEvictOrNull( insertHeadEntryAt(state, index, newEntry); state.sizeManager.increment(); return newEntry; + } finally { + lock.unlock(); } } @@ -685,7 +706,9 @@ public TEntry tryGetOrCreateOrEvictOrNull( public TEntry remove(@Nullable K1 key1, @Nullable K2 key2) { long keyHash = D2.Entry.hash(key1, key2); int index = bucketIndex(state.buckets, keyHash); - synchronized (getTableWriteLock(state)) { + ReentrantLock lock = getTableWriteLock(state); + lock.lock(); + try { Entry prev = null; for (Entry curEntry = bucketAt(state, index); curEntry != null; @@ -697,6 +720,8 @@ public TEntry remove(@Nullable K1 key1, @Nullable K2 key2) { } } return null; + } finally { + lock.unlock(); } } @@ -780,7 +805,7 @@ public static final class SizeManager { * Bucket index the last eviction removed from. The next scan resumes here, so a sustained * eviction stream doesn't repeatedly re-walk the same hot entries clustered near bucket 0. */ - @GuardedBy("getTableWriteLock(buckets)") + @GuardedBy("the corresponding State's getTableWriteLock(state)") private int evictionCursor; public SizeManager(int capacity) { @@ -847,7 +872,7 @@ public void cancelReservation() { *

    The reservation survives concurrent drain and clear operations. The caller must fill it; * an abandoned reservation permanently consumes capacity. */ - @GuardedBy("getTableWriteLock(buckets)") + @GuardedBy("the corresponding State's getTableWriteLock(state)") public > boolean tryReserveOrEvict( @Nonnull AtomicReferenceArray buckets, @Strategy @Nonnull Predicate evictable) { @@ -876,12 +901,12 @@ public void decrement() { * Releases {@code removed} slots after a sweep and resets the {@code evictionCursor}. * Outstanding reservations remain counted. */ - @GuardedBy("getTableWriteLock(buckets)") + @GuardedBy("the corresponding State's getTableWriteLock(state)") @SuppressFBWarnings( value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", justification = - "evictionCursor is read and written only under synchronized (getTableWriteLock(buckets)); SpotBugs" - + " cannot model that dynamic guard") + "evictionCursor is read and written only while the corresponding State's write lock is" + + " held; SpotBugs cannot model that dynamic guard") public void release(int removed) { if (removed != 0) { size.addAndGet(-removed); @@ -897,7 +922,7 @@ public void release(int removed) { *

    This operation may inspect every live entry while holding the table write lock, so the * predicate should be quick. */ - @GuardedBy("getTableWriteLock(buckets)") + @GuardedBy("the corresponding State's getTableWriteLock(state)") @Nullable public > TEntry evictOne( @Nonnull AtomicReferenceArray buckets, @@ -917,12 +942,12 @@ public > TEntry evictOne( return null; } - @GuardedBy("getTableWriteLock(buckets)") + @GuardedBy("the corresponding State's getTableWriteLock(state)") @SuppressFBWarnings( value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", justification = - "evictionCursor is read and written only under synchronized (getTableWriteLock(buckets)); SpotBugs" - + " cannot model that dynamic guard") + "evictionCursor is read and written only while the corresponding State's write lock is" + + " held; SpotBugs cannot model that dynamic guard") @StrategyConsumer @Nullable private > TEntry evictOneInRange( @@ -934,7 +959,7 @@ private > TEntry evictOneInRange( TEntry prev = null; for (TEntry e = buckets.get(i); e != null; e = e.next()) { if (evictable.test(e)) { - unlink(buckets, i, prev, e); + unlinkUnchecked(buckets, i, prev, e); evictionCursor = i; return e; } @@ -949,12 +974,12 @@ private > TEntry evictOneInRange( * each, and returns how many were removed. Resets the scan position, since a full pass leaves * nothing later to resume from. */ - @GuardedBy("getTableWriteLock(buckets)") + @GuardedBy("the corresponding State's getTableWriteLock(state)") @SuppressFBWarnings( value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", justification = - "evictionCursor is read and written only under synchronized (getTableWriteLock(buckets)); SpotBugs" - + " cannot model that dynamic guard") + "evictionCursor is read and written only while the corresponding State's write lock is" + + " held; SpotBugs cannot model that dynamic guard") @StrategyConsumer public > int evictAll( @Nonnull AtomicReferenceArray buckets, @@ -964,7 +989,7 @@ public > int evictAll( TEntry prev = null; for (TEntry e = buckets.get(i); e != null; e = e.next()) { if (evictable.test(e)) { - unlink(buckets, i, prev, e); + unlinkUnchecked(buckets, i, prev, e); size.decrementAndGet(); count++; } else { @@ -990,6 +1015,14 @@ public static final class State> { public final AtomicReferenceArray buckets; final SizeManager sizeManager; + /** + * The table's write lock. Unlike the bare-array helpers' opaque {@code synchronized} monitor, + * this is a real {@link ReentrantLock} so a {@link Reservation} can acquire it in {@link + * #tryReserve} and release it later in {@link Reservation#finish}/{@link Reservation#close}, + * across separate calls -- something a lexically-scoped {@code synchronized} block cannot do. + */ + final ReentrantLock writeLock = new ReentrantLock(); + private State(AtomicReferenceArray buckets, int maxCapacity) { this.buckets = buckets; this.sizeManager = new SizeManager(maxCapacity); @@ -1033,8 +1066,8 @@ public static > boolean tryReserveSlot( /** * Claims one slot in {@code state} and returns a handle for completing the find-or-insert - * protocol, or an empty handle if the table is full. Lock-free — does not acquire the table write - * lock. Never returns {@code null}, so this always composes with try-with-resources: + * protocol, or an empty handle if the table is full. Never returns {@code null}, so this always + * composes with try-with-resources: * * 

    {@code
    * try (Reservation r = ConcurrentHashtable.tryReserve(state)) {
@@ -1048,6 +1081,17 @@ public static > boolean tryReserveSlot(
    * are the source of truth), it just avoids paying for a lock and a factory call when a hit was
    * already visible lock-free.
    *
+   * 
    Checks {@link SizeManager#isFull()} lock-free first and returns an empty reservation
+   * immediately if the table already looks full, without touching the write lock. Otherwise it
+   * acquires the write lock, re-checks (the lock-free peek may be stale), and either reserves the
+   * slot -- returning a real reservation that holds the lock open across this call, to be
+   * released later by {@link Reservation#finish}/{@link Reservation#close} -- or releases the lock
+   * and returns empty. Holding the lock for the reservation's whole lifetime (rather than just
+   * across this call, the way {@link SizeManager#tryReserve()} does) is what lets {@link
+   * Reservation#tryGetOrInsertOrNull} skip a second, separately-locked comparison pass: the
+   * reserve-then-insert-or-discard sequence is one uninterrupted critical section, so no concurrent
+   * reservation for the same logical duplicate can slip in between.
+   *
    * 
    Always returns a non-null handle — even when the table is full — so the caller must check
    * {@link Reservation#isReserved()} (or simply call {@link Reservation#tryGetOrInsertOrNull},
    * which returns {@code null} on an absent reservation) rather than assume every reservation is
@@ -1056,7 +1100,17 @@ public static > boolean tryReserveSlot(
   @Nonnull
   public static > Reservation tryReserve(
       @Nonnull State state) {
-    return new Reservation<>(state.sizeManager.tryReserve() ? state : null);
+    if (state.sizeManager.isFull()) {
+      return new Reservation<>(null);
+    }
+    ReentrantLock lock = state.writeLock;
+    lock.lock();
+    if (state.sizeManager.isFull()) {
+      lock.unlock();
+      return new Reservation<>(null);
+    }
+    state.sizeManager.increment();
+    return new Reservation<>(state);
   }
 
   /**
@@ -1069,6 +1123,13 @@ public static > Reservation tryReserve(
    * method reference can build the entry directly from its natural constructor arguments, without
    * an intermediate holder object or a capturing lambda.
    *
+   * 
    A real (non-empty) reservation holds the table's write lock from the moment {@link
+   * #tryReserve} returns until {@link #finish} or {@link #close} releases it -- unlike a
+   * lexically-scoped {@code synchronized} block, whose acquisition and release can't span two
+   * separate calls. Keep a reservation's lifetime short: nothing else can write to the table while
+   * one is open, and (having deliberately dropped the CAS-based fast path that would let two
+   * threads race for the same slot) nothing else can even reserve.
+   *
    * @param  the table's entry type, itself self-bound (see {@link
    *     ConcurrentHashtable.Entry})
    */
@@ -1210,27 +1271,38 @@ public  Maybe tryGetOrInsert(
       return Maybe.of(tryGetOrInsertOrNull(a, b, c, d, factory));
     }
 
+    /**
+     * Runs the locked comparison/link decision. The write lock is already held -- acquired by
+     * {@link #tryReserve} and not yet released -- so this needs no {@code synchronized}/{@code
+     * lock()} of its own.
+     */
     private TEntry finish(@Nonnull TEntry newEntry) {
-      synchronized (getTableWriteLock(state)) {
-        int index = bucketIndex(state.buckets, newEntry.keyHash);
-        for (TEntry curEntry = bucketAt(state, index);
-            curEntry != null;
-            curEntry = curEntry.next()) {
-          if (curEntry.keyHash == newEntry.keyHash && curEntry.matches(newEntry)) {
-            return curEntry;
-          }
+      int index = bucketIndex(state.buckets, newEntry.keyHash);
+      for (TEntry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) {
+        if (curEntry.keyHash == newEntry.keyHash && curEntry.matches(newEntry)) {
+          return curEntry;
         }
-        insertHeadEntryAt(state, index, newEntry);
-        consumed = true;
-        return newEntry;
       }
+      insertHeadEntryAt(state, index, newEntry);
+      consumed = true;
+      return newEntry;
     }
 
-    /** Gives back an unconsumed reservation's slot; a no-op on an empty reservation. */
+    /**
+     * Releases the write lock {@link #tryReserve} acquired, giving back the slot first if it was
+     * never consumed. A no-op on an empty reservation, which never acquired the lock.
+     */
     @Override
     public void close() {
-      if (state != null && !consumed) {
-        state.sizeManager.cancelReservation();
+      if (state == null) {
+        return;
+      }
+      try {
+        if (!consumed) {
+          state.sizeManager.decrement();
+        }
+      } finally {
+        state.writeLock.unlock();
       }
     }
   }
@@ -1259,8 +1331,12 @@ public interface Function4 {
    */
   public static > boolean tryReserveOrEvict(
       @Nonnull State state, @Strategy @Nonnull Predicate evictable) {
-    synchronized (getTableWriteLock(state)) {
+    ReentrantLock lock = getTableWriteLock(state);
+    lock.lock();
+    try {
       return state.sizeManager.tryReserveOrEvict(state.buckets, evictable);
+    } finally {
+      lock.unlock();
     }
   }
 
@@ -1272,8 +1348,12 @@ public static > boolean tryReserveOrEvict(
   @Nullable
   public static > TEntry evictOne(
       @Nonnull State state, @Strategy @Nonnull Predicate evictable) {
-    synchronized (getTableWriteLock(state)) {
+    ReentrantLock lock = getTableWriteLock(state);
+    lock.lock();
+    try {
       return state.sizeManager.evictOne(state.buckets, evictable);
+    } finally {
+      lock.unlock();
     }
   }
 
@@ -1283,8 +1363,12 @@ public static > TEntry evictOne(
    */
   public static > int evictAll(
       @Nonnull State state, @Strategy @Nonnull Predicate evictable) {
-    synchronized (getTableWriteLock(state)) {
+    ReentrantLock lock = getTableWriteLock(state);
+    lock.lock();
+    try {
       return state.sizeManager.evictAll(state.buckets, evictable);
+    } finally {
+      lock.unlock();
     }
   }
 
@@ -1345,10 +1429,15 @@ public static Object getWriteLock(@Nonnull AtomicReferenceArray buckets, long
     return getWriteLockAt(buckets, bucketIndex(buckets, keyHash));
   }
 
-  /** {@link #getWriteLock(AtomicReferenceArray, long)} over a {@link State}. */
+  /**
+   * Returns the {@link State}'s {@link ReentrantLock}, held for the same keyed scan-and-mutate
+   * operations as {@link #getWriteLock(AtomicReferenceArray, long)}. Unlike the bare-array
+   * overload's opaque {@code synchronized} monitor, this lock can be acquired in one call and
+   * released in a later one -- see {@link Reservation}.
+   */
   @Nonnull
-  public static Object getWriteLock(@Nonnull State state, long keyHash) {
-    return getWriteLock(state.buckets, keyHash);
+  public static ReentrantLock getWriteLock(@Nonnull State state, long keyHash) {
+    return state.writeLock;
   }
 
   /**
@@ -1362,10 +1451,10 @@ public static Object getWriteLockAt(@Nonnull AtomicReferenceArray buckets, in
     return buckets;
   }
 
-  /** {@link #getWriteLockAt(AtomicReferenceArray, int)} over a {@link State}. */
+  /** {@link #getWriteLock(State, long)}, for a bucket index that has already been computed. */
   @Nonnull
-  public static Object getWriteLockAt(@Nonnull State state, int bucketIndex) {
-    return getWriteLockAt(state.buckets, bucketIndex);
+  public static ReentrantLock getWriteLockAt(@Nonnull State state, int bucketIndex) {
+    return state.writeLock;
   }
 
   /**
@@ -1377,10 +1466,14 @@ public static Object getTableWriteLock(@Nonnull AtomicReferenceArray buckets)
     return buckets;
   }
 
-  /** {@link #getTableWriteLock(AtomicReferenceArray)} over a {@link State}. */
+  /**
+   * Returns the {@link State}'s {@link ReentrantLock}, guarding operations that span all buckets.
+   * Unlike {@link #getTableWriteLock(AtomicReferenceArray)}'s opaque {@code synchronized} monitor,
+   * this lock can be acquired in one call and released in a later one -- see {@link Reservation}.
+   */
   @Nonnull
-  public static Object getTableWriteLock(@Nonnull State state) {
-    return getTableWriteLock(state.buckets);
+  public static ReentrantLock getTableWriteLock(@Nonnull State state) {
+    return state.writeLock;
   }
 
   public static int bucketIndex(@Nonnull AtomicReferenceArray buckets, long keyHash) {
@@ -1513,19 +1606,27 @@ public static > void insertHeadEntryAt(
       @Nonnull AtomicReferenceArray buckets, int index, @Nonnull TEntry entry) {
     assert Thread.holdsLock(getWriteLockAt(buckets, index))
         : "insertHeadEntryAt called without holding getWriteLockAt(buckets, index)";
-    assert entry.next() == null
-        : "Entry already linked -- inserting the same Entry instance twice corrupts the chain"
-            + " (unlink() deliberately leaves a removed entry's next intact for in-flight"
-            + " readers, so a removed entry must never be reinserted)";
-    entry.setNext(buckets.get(index));
-    buckets.set(index, entry);
+    insertHeadEntryAtUnchecked(buckets, index, entry);
   }
 
   /** {@link #insertHeadEntryAt(AtomicReferenceArray, int, Entry)} over a {@link State}. */
   @GuardedBy("getWriteLockAt(state, index)")
   public static > void insertHeadEntryAt(
       @Nonnull State state, int index, @Nonnull TEntry entry) {
-    insertHeadEntryAt(state.buckets, index, entry);
+    assert getWriteLockAt(state, index).isHeldByCurrentThread()
+        : "insertHeadEntryAt called without holding getWriteLockAt(state, index)";
+    insertHeadEntryAtUnchecked(state.buckets, index, entry);
+  }
+
+  /** Splicing logic shared by both {@link #insertHeadEntryAt} overloads, without their assert. */
+  private static > void insertHeadEntryAtUnchecked(
+      @Nonnull AtomicReferenceArray buckets, int index, @Nonnull TEntry entry) {
+    assert entry.next() == null
+        : "Entry already linked -- inserting the same Entry instance twice corrupts the chain"
+            + " (unlink() deliberately leaves a removed entry's next intact for in-flight"
+            + " readers, so a removed entry must never be reinserted)";
+    entry.setNext(buckets.get(index));
+    buckets.set(index, entry);
   }
 
   /**
@@ -1550,7 +1651,9 @@ public static > void insertHeadEntryFor(
   @GuardedBy("getTableWriteLock(state)")
   public static > void insertReserved(
       @Nonnull State state, long keyHash, @Nonnull TEntry entry) {
-    insertHeadEntryFor(state.buckets, keyHash, entry);
+    assert getTableWriteLock(state).isHeldByCurrentThread()
+        : "insertReserved called without holding getTableWriteLock(state)";
+    insertHeadEntryAtUnchecked(state.buckets, bucketIndex(state.buckets, keyHash), entry);
   }
 
   /**
@@ -1570,19 +1673,30 @@ public static > void unlink(
       @Nonnull TEntry entry) {
     assert Thread.holdsLock(getWriteLockAt(buckets, index))
         : "unlink called without holding getWriteLockAt(buckets, index)";
-    TEntry next = entry.next();
-    if (prev == null) {
-      buckets.set(index, next);
-    } else {
-      prev.setNext(next);
-    }
+    unlinkUnchecked(buckets, index, prev, entry);
   }
 
   /** {@link #unlink(AtomicReferenceArray, int, Entry, Entry)} over a {@link State}. */
   @GuardedBy("getWriteLockAt(state, index)")
   public static > void unlink(
       @Nonnull State state, int index, @Nullable TEntry prev, @Nonnull TEntry entry) {
-    unlink(state.buckets, index, prev, entry);
+    assert getWriteLockAt(state, index).isHeldByCurrentThread()
+        : "unlink called without holding getWriteLockAt(state, index)";
+    unlinkUnchecked(state.buckets, index, prev, entry);
+  }
+
+  /** Splicing logic shared by both {@link #unlink} overloads, without their assert. */
+  private static > void unlinkUnchecked(
+      @Nonnull AtomicReferenceArray buckets,
+      int index,
+      @Nullable TEntry prev,
+      @Nonnull TEntry entry) {
+    TEntry next = entry.next();
+    if (prev == null) {
+      buckets.set(index, next);
+    } else {
+      prev.setNext(next);
+    }
   }
 
   /**
@@ -1624,13 +1738,15 @@ public static > boolean removeIf(
   public static > boolean removeIf(
       @Nonnull State state, @Strategy @Nonnull Predicate predicate) {
     AtomicReferenceArray buckets = state.buckets;
-    synchronized (getTableWriteLock(state)) {
+    ReentrantLock lock = getTableWriteLock(state);
+    lock.lock();
+    try {
       boolean removed = false;
       for (int i = 0; i < buckets.length(); i++) {
         TEntry prev = null;
         for (TEntry e = buckets.get(i); e != null; e = e.next()) {
           if (predicate.test(e)) {
-            unlink(buckets, i, prev, e);
+            unlink(state, i, prev, e);
             state.sizeManager.decrement();
             removed = true;
           } else {
@@ -1639,6 +1755,8 @@ public static > boolean removeIf(
         }
       }
       return removed;
+    } finally {
+      lock.unlock();
     }
   }
 
@@ -1736,13 +1854,16 @@ private static > int drainCounting(
    * Removes all entries from {@code state}, invokes {@code drainedEntryConsumer} for each one, and
    * releases one capacity slot for each removed entry.
    *
-   * 
    The drain holds the table write lock while detaching buckets and invoking the consumer. For
-   * each removed entry, the consumer is invoked synchronously after its bucket is detached.
-   * Capacity for removed entries is released only after all invocations return. Outstanding
-   * reservations remain counted.
+   * 
    The lock is acquired and released bucket by bucket rather than held for the whole sweep, so
+   * a table-level operation (an insert, a reservation, an eviction) can interleave between buckets
+   * instead of waiting out the entire drain. This means the drain is no longer an atomic snapshot:
+   * an entry inserted into a not-yet-visited bucket while the drain is in progress is swept up too,
+   * and each bucket's capacity slot is released immediately before its consumer invocation runs,
+   * rather than once for the whole sweep at the end. Outstanding reservations remain counted.
    *
-   * 
    The consumer should be quick and must not throw; if it throws, removed entries are not
-   * restored and their capacity is not released.
+   * 
    The consumer should be quick and must not throw; if it throws, entries not yet reached are
+   * left in the table with their capacity still counted, but entries already detached keep their
+   * capacity released regardless.
    *
    * @param  entry type
    * @param state table state to drain
@@ -1751,8 +1872,29 @@ private static > int drainCounting(
   public static > void drain(
       @Nonnull State state,
       @Strategy @Nonnull Consumer drainedEntryConsumer) {
-    synchronized (getTableWriteLock(state)) {
-      state.sizeManager.release(drainCounting(state.buckets, drainedEntryConsumer));
+    ReentrantLock lock = getTableWriteLock(state);
+    AtomicReferenceArray buckets = state.buckets;
+    for (int i = 0; i < buckets.length(); i++) {
+      lock.lock();
+      try {
+        TEntry head = buckets.get(i);
+        if (head == null) {
+          continue;
+        }
+        buckets.set(i, null);
+        for (TEntry e = head; e != null; e = e.next()) {
+          state.sizeManager.decrement();
+          drainedEntryConsumer.accept(e);
+        }
+      } finally {
+        lock.unlock();
+      }
+    }
+    lock.lock();
+    try {
+      state.sizeManager.release(0); // full sweep: reset the scan position
+    } finally {
+      lock.unlock();
     }
   }
 
@@ -1760,13 +1902,16 @@ public static > void drain(
    * Removes all entries from {@code state}, invokes {@code drainedEntryConsumer} with {@code
    * context} and each removed entry, and releases one capacity slot for each removed entry.
    *
-   * 
    The drain holds the table write lock while detaching buckets and invoking the consumer. For
-   * each removed entry, the consumer is invoked synchronously after its bucket is detached.
-   * Capacity for removed entries is released only after all invocations return. Outstanding
-   * reservations remain counted.
+   * 
    The lock is acquired and released bucket by bucket rather than held for the whole sweep, so
+   * a table-level operation (an insert, a reservation, an eviction) can interleave between buckets
+   * instead of waiting out the entire drain. This means the drain is no longer an atomic snapshot:
+   * an entry inserted into a not-yet-visited bucket while the drain is in progress is swept up too,
+   * and each bucket's capacity slot is released immediately before its consumer invocation runs,
+   * rather than once for the whole sweep at the end. Outstanding reservations remain counted.
    *
-   * 
    The consumer should be quick and must not throw; if it throws, removed entries are not
-   * restored and their capacity is not released.
+   * 
    The consumer should be quick and must not throw; if it throws, entries not yet reached are
+   * left in the table with their capacity still counted, but entries already detached keep their
+   * capacity released regardless.
    *
    * @param  context type
    * @param  entry type
@@ -1778,8 +1923,29 @@ public static > void drain(
       @Nonnull State state,
       C context,
       @Strategy @Nonnull BiConsumer drainedEntryConsumer) {
-    synchronized (getTableWriteLock(state)) {
-      state.sizeManager.release(drainCounting(state.buckets, context, drainedEntryConsumer));
+    ReentrantLock lock = getTableWriteLock(state);
+    AtomicReferenceArray buckets = state.buckets;
+    for (int i = 0; i < buckets.length(); i++) {
+      lock.lock();
+      try {
+        TEntry head = buckets.get(i);
+        if (head == null) {
+          continue;
+        }
+        buckets.set(i, null);
+        for (TEntry e = head; e != null; e = e.next()) {
+          state.sizeManager.decrement();
+          drainedEntryConsumer.accept(context, e);
+        }
+      } finally {
+        lock.unlock();
+      }
+    }
+    lock.lock();
+    try {
+      state.sizeManager.release(0); // full sweep: reset the scan position
+    } finally {
+      lock.unlock();
     }
   }
 
@@ -1820,8 +1986,12 @@ private static int clearCounting(@Nonnull AtomicReferenceArray state) {
-    synchronized (getTableWriteLock(state)) {
+    ReentrantLock lock = getTableWriteLock(state);
+    lock.lock();
+    try {
       state.sizeManager.release(clearCounting(state.buckets));
+    } finally {
+      lock.unlock();
     }
   }
 
diff --git a/internal-api/src/test/java/datadog/trace/api/telemetry/LogCollectorTest.java b/internal-api/src/test/java/datadog/trace/api/telemetry/LogCollectorTest.java
index 23b65cc523b..ef051e4db0e 100644
--- a/internal-api/src/test/java/datadog/trace/api/telemetry/LogCollectorTest.java
+++ b/internal-api/src/test/java/datadog/trace/api/telemetry/LogCollectorTest.java
@@ -162,9 +162,12 @@ void acceptsNewLogGroupAfterDrainDetachesFullBucket() throws Exception {
     try {
       drainThread.start();
       await(bucketDetached);
+      // Capacity for the detached entry is released before its consumer runs, so the writer's
+      // lock-free isFull() peek already sees room and proceeds to the table lock -- which the
+      // drain still holds for the whole sweep, so the writer blocks there instead of on isFull().
       writerThread.start();
       new PollingConditions(TIMEOUT_SECONDS)
-          .eventually(() -> assertThat(writerThread.getState()).isEqualTo(Thread.State.BLOCKED));
+          .eventually(() -> assertThat(writerThread.getState()).isEqualTo(Thread.State.WAITING));
     } finally {
       releaseDrain.countDown();
     }
diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableSizeManagerTest.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableSizeManagerTest.java
index 2fceb4ce3ea..cd6a6aafc1e 100644
--- a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableSizeManagerTest.java
+++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableSizeManagerTest.java
@@ -8,6 +8,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import datadog.trace.test.util.PollingConditions;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Predicate;
 import org.junit.jupiter.api.Test;
 
@@ -83,8 +84,12 @@ void insertReservedSplicesWithoutTouchingTheCountAfterATryReserve() {
     assertEquals(1, state.sizeManager.estimateSize());
 
     TestEntry entry = new TestEntry(0, "reserved");
-    synchronized (ConcurrentHashtable.getTableWriteLock(state)) {
+    ReentrantLock lock = ConcurrentHashtable.getTableWriteLock(state);
+    lock.lock();
+    try {
       ConcurrentHashtable.insertReserved(state, entry.keyHash, entry);
+    } finally {
+      lock.unlock();
     }
 
     assertSame(entry, state.buckets.get(0));
@@ -236,8 +241,12 @@ void stateCreateCappedBundlesBucketsAndSizeManager() {
   void stateLevelTryReserveOrEvictAndEvictOneAndEvictAllDelegateToSizeManager() {
     ConcurrentHashtable.State state =
         ConcurrentHashtable.createBounded(TestEntry.class, 1);
-    synchronized (ConcurrentHashtable.getWriteLockAt(state, 0)) {
+    ReentrantLock lockAt0 = ConcurrentHashtable.getWriteLockAt(state, 0);
+    lockAt0.lock();
+    try {
       ConcurrentHashtable.insertHeadEntryAt(state, 0, new TestEntry(0, "a"));
+    } finally {
+      lockAt0.unlock();
     }
     state.sizeManager.increment();
     assertTrue(ConcurrentHashtable.isFull(state));
@@ -247,12 +256,16 @@ void stateLevelTryReserveOrEvictAndEvictOneAndEvictAllDelegateToSizeManager() {
     // the reservation already happened and a plain insertHeadEntryAt/increment would double-count.
     // Both steps go in ONE critical section: tryReserveOrEvict is self-locking, so on its own it
     // leaves a window where a drain/clear could reset the count out from under the reservation.
-    synchronized (ConcurrentHashtable.getTableWriteLock(state)) {
+    ReentrantLock tableLock = ConcurrentHashtable.getTableWriteLock(state);
+    tableLock.lock();
+    try {
       boolean reserved = ConcurrentHashtable.tryReserveOrEvict(state, e -> true);
       assertTrue(reserved);
       assertEquals(1, ConcurrentHashtable.estimateSize(state));
       assertNull(state.buckets.get(0)); // "a" was evicted; the reserved slot has no entry yet
       ConcurrentHashtable.insertReserved(state, 0, new TestEntry(0, "reserved"));
+    } finally {
+      tableLock.unlock();
     }
 
     int evicted = ConcurrentHashtable.evictAll(state, e -> true);
@@ -260,8 +273,11 @@ void stateLevelTryReserveOrEvictAndEvictOneAndEvictAllDelegateToSizeManager() {
     assertEquals(0, ConcurrentHashtable.estimateSize(state));
     assertFalse(ConcurrentHashtable.isFull(state));
 
-    synchronized (ConcurrentHashtable.getWriteLockAt(state, 0)) {
+    lockAt0.lock();
+    try {
       ConcurrentHashtable.insertHeadEntryAt(state, 0, new TestEntry(0, "b"));
+    } finally {
+      lockAt0.unlock();
     }
     state.sizeManager.increment();
     TestEntry viaEvictOne = ConcurrentHashtable.evictOne(state, e -> e.label.equals("b"));
@@ -286,16 +302,20 @@ void clearCannotInterleaveBetweenReservationAndInsert() throws InterruptedExcept
     assertTrue(ConcurrentHashtable.isFull(state));
 
     Thread clearer = new Thread(() -> ConcurrentHashtable.clear(state), "clearer");
-    synchronized (ConcurrentHashtable.getTableWriteLock(state)) {
+    ReentrantLock lock = ConcurrentHashtable.getTableWriteLock(state);
+    lock.lock();
+    try {
       clearer.start();
-      // Wait until the clear is definitely queued on the monitor we hold, so the interleaving under
+      // Wait until the clear is definitely queued on the lock we hold, so the interleaving under
       // test is the one actually attempted rather than one the scheduler happened to avoid.
       new PollingConditions()
-          .eventually(() -> assertEquals(Thread.State.BLOCKED, clearer.getState()));
+          .eventually(() -> assertEquals(Thread.State.WAITING, clearer.getState()));
 
       assertTrue(ConcurrentHashtable.tryReserveOrEvict(state, e -> true));
       ConcurrentHashtable.insertReserved(state, 0, new TestEntry(0, "reserved"));
       assertEquals(1, ConcurrentHashtable.estimateSize(state));
+    } finally {
+      lock.unlock();
     }
     clearer.join();
 
@@ -332,8 +352,12 @@ void reservationSurvivesAClearLandingBetweenReserveAndInsert() {
     assertEquals(1, ConcurrentHashtable.estimateSize(state));
 
     // The reservation is still good, and filling it leaves the count matching the entries present.
-    synchronized (ConcurrentHashtable.getWriteLockAt(state, 0)) {
+    ReentrantLock lockAt0 = ConcurrentHashtable.getWriteLockAt(state, 0);
+    lockAt0.lock();
+    try {
       ConcurrentHashtable.insertReserved(state, 0, new TestEntry(0, "reserved"));
+    } finally {
+      lockAt0.unlock();
     }
     assertEquals(1, ConcurrentHashtable.estimateSize(state));
     assertNotNullLabel(state, 0, "reserved");
@@ -389,8 +413,12 @@ private static void assertNotNullLabel(
   private static TestEntry insertAt(
       ConcurrentHashtable.State state, int index, String label) {
     TestEntry entry = new TestEntry(index, label);
-    synchronized (ConcurrentHashtable.getWriteLockAt(state, index)) {
+    ReentrantLock lock = ConcurrentHashtable.getWriteLockAt(state, index);
+    lock.lock();
+    try {
       ConcurrentHashtable.insertHeadEntryAt(state, index, entry);
+    } finally {
+      lock.unlock();
     }
     return entry;
   }
@@ -398,24 +426,36 @@ private static TestEntry insertAt(
   /** {@code sizeManager.tryReserveOrEvict}, taking the write lock {@code @GuardedBy} requires. */
   private static boolean tryReserveOrEvict(
       ConcurrentHashtable.State state, Predicate evictable) {
-    synchronized (ConcurrentHashtable.getTableWriteLock(state)) {
+    ReentrantLock lock = ConcurrentHashtable.getTableWriteLock(state);
+    lock.lock();
+    try {
       return state.sizeManager.tryReserveOrEvict(state.buckets, evictable);
+    } finally {
+      lock.unlock();
     }
   }
 
   /** {@code sizeManager.evictOne}, taking the write lock {@code @GuardedBy} requires. */
   private static TestEntry evictOne(
       ConcurrentHashtable.State state, Predicate evictable) {
-    synchronized (ConcurrentHashtable.getTableWriteLock(state)) {
+    ReentrantLock lock = ConcurrentHashtable.getTableWriteLock(state);
+    lock.lock();
+    try {
       return state.sizeManager.evictOne(state.buckets, evictable);
+    } finally {
+      lock.unlock();
     }
   }
 
   /** {@code sizeManager.evictAll}, taking the write lock {@code @GuardedBy} requires. */
   private static int evictAll(
       ConcurrentHashtable.State state, Predicate evictable) {
-    synchronized (ConcurrentHashtable.getTableWriteLock(state)) {
+    ReentrantLock lock = ConcurrentHashtable.getTableWriteLock(state);
+    lock.lock();
+    try {
       return state.sizeManager.evictAll(state.buckets, evictable);
+    } finally {
+      lock.unlock();
     }
   }
 

From 474241a504faca6e4e1ee4c87eaa4a9829675ce1 Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 10:13:45 -0400
Subject: [PATCH 05/13] Add JMH benchmarks for hashIterable/hashIterator
 allocation and drain lock contention

Addresses the two flag-as-measure findings from the perf review of the
ReentrantLock migration:

- ConcurrentHashtableFindBenchmark confirms (via -Pjmh.profilers=gc) that the
  Iterable/Iterator returned by hashIterable/hashIterator is scalar-replaced
  away rather than allocated on LogCollector's find() hot path.
- ConcurrentHashtableDrainBenchmark measures writer throughput against a
  continuously draining table, giving a repeatable baseline for drain()'s
  per-bucket lock acquisition versus holding the lock for the whole sweep.

Co-Authored-By: Claude Sonnet 5 
---
 .../ConcurrentHashtableDrainBenchmark.java    | 118 ++++++++++++++++++
 .../ConcurrentHashtableFindBenchmark.java     | 101 +++++++++++++++
 2 files changed, 219 insertions(+)
 create mode 100644 internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
 create mode 100644 internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableFindBenchmark.java

diff --git a/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
new file mode 100644
index 00000000000..ae13d0371a9
--- /dev/null
+++ b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
@@ -0,0 +1,118 @@
+package datadog.trace.util;
+
+import static java.util.concurrent.TimeUnit.MICROSECONDS;
+
+import javax.annotation.Nonnull;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Group;
+import org.openjdk.jmh.annotations.GroupThreads;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+/**
+ * Measures writer throughput while a background thread continuously {@link
+ * ConcurrentHashtable#drain}s the same table -- the scenario behind the per-bucket lock acquisition
+ * in {@code drain()} (one {@code lock()}/{@code unlock()} pair per bucket, instead of one held for
+ * the whole sweep), added so writers can interleave between buckets rather than wait out an entire
+ * drain.
+ *
+ * 
    The writer mirrors {@code LogCollector.addLogMessage}'s shape: a lock-free scan over a small,
+ * rotating key set ({@link #N_KEYS} distinct log groups) that hits on the fast path once a key has
+ * been inserted, falling back to {@link ConcurrentHashtable#tryReserve} (which holds the table lock
+ * for the reservation's lifetime) only on a miss. Because {@code drainer} periodically empties the
+ * table, writers keep taking the slow, lock-holding path throughout the run instead of settling
+ * into steady-state lock-free hits -- unlike an unbounded-key writer, this keeps the table's
+ * occupancy well under capacity so writers aren't dominated by the lock-free {@code isFull()}
+ * fast-reject once full, and the throughput actually reflects contention for the table lock against
+ * an in-progress drain.
+ *
+ * 
    One JMH {@code @Group} thread continuously drains; the rest continuously write. JMH reports
+ * separate throughput for the {@code drainer} and {@code writer} roles under the {@code mixed}
+ * group, so a regression in either role under contention is visible without needing to reconstruct
+ * the old whole-sweep-lock strategy separately.
+ *
+ * 
    {@code
+ * ./gradlew :internal-api:jmh -Pjmh.includes=ConcurrentHashtableDrainBenchmark -Pjmh.forks=1
+ * }
    
+ */
+@Fork(2)
+@Warmup(iterations = 2)
+@Measurement(iterations = 3)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(MICROSECONDS)
+public class ConcurrentHashtableDrainBenchmark {
+
+  static final int N_KEYS = 32;
+  static final int CAPACITY = 64;
+
+  static final long[] KEY_HASHES = new long[N_KEYS];
+
+  static {
+    for (int i = 0; i < N_KEYS; ++i) {
+      KEY_HASHES[i] = LongHashingUtils.hash("key-" + i);
+    }
+  }
+
+  static final class DrainEntry extends ConcurrentHashtable.Entry {
+    DrainEntry(long keyHash) {
+      super(keyHash);
+    }
+
+    @Override
+    public boolean matches(@Nonnull DrainEntry other) {
+      // hashIterator already filters candidates by keyHash equality before yielding them.
+      return true;
+    }
+  }
+
+  @State(Scope.Benchmark)
+  public static class SharedState {
+    ConcurrentHashtable.State table;
+
+    @Setup(Level.Iteration)
+    public void setUp() {
+      table = ConcurrentHashtable.createBounded(DrainEntry.class, CAPACITY);
+    }
+  }
+
+  @State(Scope.Thread)
+  public static class WriterState {
+    int cursor;
+
+    int next() {
+      int i = cursor;
+      cursor = (i + 1) & (N_KEYS - 1);
+      return i;
+    }
+  }
+
+  @Benchmark
+  @Group("mixed")
+  @GroupThreads(1)
+  public void drainer(SharedState s) {
+    ConcurrentHashtable.drain(s.table, entry -> {});
+  }
+
+  @Benchmark
+  @Group("mixed")
+  @GroupThreads(3)
+  public void writer(SharedState s, WriterState w) {
+    long keyHash = KEY_HASHES[w.next()];
+    for (DrainEntry entry : ConcurrentHashtable.hashIterable(s.table, keyHash)) {
+      return; // lock-free hit, mirroring LogCollector.find()
+    }
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(s.table)) {
+      if (r.isReserved()) {
+        r.tryGetOrInsertOrNull(new DrainEntry(keyHash));
+      }
+    }
+  }
+}
diff --git a/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableFindBenchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableFindBenchmark.java
new file mode 100644
index 00000000000..3f9ccb79f00
--- /dev/null
+++ b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableFindBenchmark.java
@@ -0,0 +1,101 @@
+package datadog.trace.util;
+
+import static java.util.concurrent.TimeUnit.MICROSECONDS;
+
+import javax.annotation.Nonnull;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+/**
+ * Measures the lock-free hash-bucket scan used by {@link LogCollector#find}: a {@code for} loop
+ * over {@link ConcurrentHashtable#hashIterable}, matching by key hash then a per-entry predicate.
+ *
+ * 
    Run with {@code -Pjmh.profilers=gc} to confirm the {@link Iterable}/{@link java.util.Iterator}
+ * allocated per call (the anonymous instances returned by {@link
+ * ConcurrentHashtable#hashIterable}/{@link ConcurrentHashtable#hashIterator}) are scalar-replaced
+ * away by escape analysis rather than landing on the heap:
+ *
+ * 
    {@code
+ * ./gradlew :internal-api:jmh -Pjmh.includes=ConcurrentHashtableFindBenchmark -Pjmh.profilers=gc -Pjmh.forks=1
+ * }
    
+ */
+@Fork(2)
+@Warmup(iterations = 2)
+@Measurement(iterations = 3)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(MICROSECONDS)
+@Threads(8)
+public class ConcurrentHashtableFindBenchmark {
+
+  static final int N_KEYS = 64;
+  static final int CAPACITY = 128;
+
+  static final long[] KEY_HASHES = new long[N_KEYS];
+
+  static {
+    for (int i = 0; i < N_KEYS; ++i) {
+      KEY_HASHES[i] = LongHashingUtils.hash("key-" + i);
+    }
+  }
+
+  /** Mirrors {@code LogCollector.RawLogMessage}: a keyHash plus a payload compared on match. */
+  static final class FindEntry extends ConcurrentHashtable.Entry {
+    final int payload;
+
+    FindEntry(long keyHash, int payload) {
+      super(keyHash);
+      this.payload = payload;
+    }
+
+    @Override
+    public boolean matches(@Nonnull FindEntry other) {
+      return payload == other.payload;
+    }
+  }
+
+  @State(Scope.Benchmark)
+  public static class SharedState {
+    ConcurrentHashtable.State table;
+
+    @Setup(Level.Iteration)
+    public void setUp() {
+      table = ConcurrentHashtable.createBounded(FindEntry.class, CAPACITY);
+      for (int i = 0; i < N_KEYS; ++i) {
+        ConcurrentHashtable.tryReserve(table).tryGetOrInsertOrNull(new FindEntry(KEY_HASHES[i], i));
+      }
+    }
+  }
+
+  @State(Scope.Thread)
+  public static class ThreadState {
+    int cursor;
+
+    int next() {
+      int i = cursor;
+      cursor = (i + 1) & (N_KEYS - 1);
+      return i;
+    }
+  }
+
+  /** Same loop shape as {@code LogCollector.find}: scan candidates for a keyHash, match, return. */
+  @Benchmark
+  public FindEntry find(SharedState s, ThreadState t) {
+    int i = t.next();
+    for (FindEntry entry : ConcurrentHashtable.hashIterable(s.table, KEY_HASHES[i])) {
+      if (entry.payload == i) {
+        return entry;
+      }
+    }
+    return null;
+  }
+}

From 3349e9dfb36748f5233d611547a01d7b3dc107d2 Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 10:35:17 -0400
Subject: [PATCH 06/13] Add whole-sweep-lock variant to drain benchmark for A/B
 comparison

Reimplements the pre-PR whole-sweep-lock drain() strategy locally (using
only ConcurrentHashtable's public API) so both locking schemes can be
measured side by side in one run, isolating lock granularity as the only
variable (both variants release capacity per-entry).

Measured at capacity=64 under 3-writer contention: per-bucket locking
drains ~32x slower (0.12 vs 3.76 ops/us) but lets writers achieve ~33%
higher combined throughput (262.6 vs 197.1 ops/us) while a drain is in
progress. Accepted tradeoff -- LogCollector's writers are on the
application-thread hot path and drain runs infrequently (once per
telemetry flush), so favoring writer throughput over drain throughput
is the right call for this caller.

Co-Authored-By: Claude Sonnet 5 
---
 .../ConcurrentHashtableDrainBenchmark.java    | 108 +++++++++++++-----
 1 file changed, 81 insertions(+), 27 deletions(-)

diff --git a/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
index ae13d0371a9..6a1e28c875f 100644
--- a/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
+++ b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
@@ -2,6 +2,9 @@
 
 import static java.util.concurrent.TimeUnit.MICROSECONDS;
 
+import java.util.concurrent.atomic.AtomicReferenceArray;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
 import javax.annotation.Nonnull;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -18,26 +21,31 @@
 import org.openjdk.jmh.annotations.Warmup;
 
 /**
- * Measures writer throughput while a background thread continuously {@link
- * ConcurrentHashtable#drain}s the same table -- the scenario behind the per-bucket lock acquisition
- * in {@code drain()} (one {@code lock()}/{@code unlock()} pair per bucket, instead of one held for
- * the whole sweep), added so writers can interleave between buckets rather than wait out an entire
- * drain.
+ * Compares writer throughput against a continuously draining table under two locking schemes:
+ *
+ * 
      
+ *   
    • {@code perBucketLock} -- {@link ConcurrentHashtable#drain}'s current strategy: the table
+ *       lock is acquired and released once per bucket, so a writer can slip in between buckets
+ *       instead of waiting out the whole sweep.
+ *   
    • {@code wholeSweepLock} -- the strategy {@code drain()} used before this PR: the table lock
+ *       is acquired once and held for the entire sweep. Reimplemented locally ({@link
+ *       #drainWholeSweepLock}) using only {@code ConcurrentHashtable}'s public building blocks, so
+ *       this file doesn't depend on checking out an older commit to get the comparison.
+ * 
    
+ *
+ * 
    Both variants release capacity per-entry (one {@link
+ * ConcurrentHashtable.SizeManager#decrement()} per drained entry, mirroring current production
+ * behavior) so the only variable under test is lock granularity, not the capacity-release change
+ * that landed alongside it.
  *
  * 
    The writer mirrors {@code LogCollector.addLogMessage}'s shape: a lock-free scan over a small,
  * rotating key set ({@link #N_KEYS} distinct log groups) that hits on the fast path once a key has
  * been inserted, falling back to {@link ConcurrentHashtable#tryReserve} (which holds the table lock
  * for the reservation's lifetime) only on a miss. Because {@code drainer} periodically empties the
  * table, writers keep taking the slow, lock-holding path throughout the run instead of settling
- * into steady-state lock-free hits -- unlike an unbounded-key writer, this keeps the table's
- * occupancy well under capacity so writers aren't dominated by the lock-free {@code isFull()}
- * fast-reject once full, and the throughput actually reflects contention for the table lock against
- * an in-progress drain.
- *
- * 
    One JMH {@code @Group} thread continuously drains; the rest continuously write. JMH reports
- * separate throughput for the {@code drainer} and {@code writer} roles under the {@code mixed}
- * group, so a regression in either role under contention is visible without needing to reconstruct
- * the old whole-sweep-lock strategy separately.
+ * into steady-state lock-free hits -- this keeps occupancy well under capacity so writers aren't
+ * dominated by the lock-free {@code isFull()} fast-reject once full, and the throughput actually
+ * reflects contention for the table lock against an in-progress drain.
  *
  * 
    {@code
  * ./gradlew :internal-api:jmh -Pjmh.includes=ConcurrentHashtableDrainBenchmark -Pjmh.forks=1
@@ -73,6 +81,34 @@ public boolean matches(@Nonnull DrainEntry other) {
     }
   }
 
+  /**
+   * The pre-PR strategy: one lock acquisition held for the entire bucket-array sweep, still
+   * releasing capacity per-entry (only lock granularity differs from the current {@code drain()}).
+   */
+  private static > void drainWholeSweepLock(
+      @Nonnull ConcurrentHashtable.State state,
+      @Nonnull Consumer drainedEntryConsumer) {
+    ReentrantLock lock = ConcurrentHashtable.getTableWriteLock(state);
+    AtomicReferenceArray buckets = state.buckets;
+    lock.lock();
+    try {
+      for (int i = 0; i < buckets.length(); i++) {
+        TEntry head = buckets.get(i);
+        if (head == null) {
+          continue;
+        }
+        buckets.set(i, null);
+        for (TEntry e = head; e != null; e = e.next()) {
+          state.sizeManager.decrement();
+          drainedEntryConsumer.accept(e);
+        }
+      }
+      state.sizeManager.release(0); // full sweep: reset the scan position
+    } finally {
+      lock.unlock();
+    }
+  }
+
   @State(Scope.Benchmark)
   public static class SharedState {
     ConcurrentHashtable.State table;
@@ -94,25 +130,43 @@ int next() {
     }
   }
 
-  @Benchmark
-  @Group("mixed")
-  @GroupThreads(1)
-  public void drainer(SharedState s) {
-    ConcurrentHashtable.drain(s.table, entry -> {});
-  }
-
-  @Benchmark
-  @Group("mixed")
-  @GroupThreads(3)
-  public void writer(SharedState s, WriterState w) {
+  private static void write(ConcurrentHashtable.State table, WriterState w) {
     long keyHash = KEY_HASHES[w.next()];
-    for (DrainEntry entry : ConcurrentHashtable.hashIterable(s.table, keyHash)) {
+    for (DrainEntry entry : ConcurrentHashtable.hashIterable(table, keyHash)) {
       return; // lock-free hit, mirroring LogCollector.find()
     }
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(s.table)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(table)) {
       if (r.isReserved()) {
         r.tryGetOrInsertOrNull(new DrainEntry(keyHash));
       }
     }
   }
+
+  @Benchmark
+  @Group("perBucketLock")
+  @GroupThreads(1)
+  public void perBucketLockDrainer(SharedState s) {
+    ConcurrentHashtable.drain(s.table, entry -> {});
+  }
+
+  @Benchmark
+  @Group("perBucketLock")
+  @GroupThreads(3)
+  public void perBucketLockWriter(SharedState s, WriterState w) {
+    write(s.table, w);
+  }
+
+  @Benchmark
+  @Group("wholeSweepLock")
+  @GroupThreads(1)
+  public void wholeSweepLockDrainer(SharedState s) {
+    drainWholeSweepLock(s.table, entry -> {});
+  }
+
+  @Benchmark
+  @Group("wholeSweepLock")
+  @GroupThreads(3)
+  public void wholeSweepLockWriter(SharedState s, WriterState w) {
+    write(s.table, w);
+  }
 }

From b632878939582984b58d9cb96cba20b84de8c0fb Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 10:39:54 -0400
Subject: [PATCH 07/13] Prefix hashIterable call with ConcurrentHashtable in
 LogCollector.find

Review feedback: the bare static-imported call read ambiguously at the
call site; qualifying it makes clear which table API is in play.

Co-Authored-By: Claude Sonnet 5 
---
 .../main/java/datadog/trace/api/telemetry/LogCollector.java    | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
index 875ee129f29..b42d8f513de 100644
--- a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
+++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
@@ -1,7 +1,6 @@
 package datadog.trace.api.telemetry;
 
 import static datadog.trace.util.ConcurrentHashtable.estimateSize;
-import static datadog.trace.util.ConcurrentHashtable.hashIterable;
 import static datadog.trace.util.LongHashingUtils.hash;
 
 import datadog.trace.api.internal.VisibleForTesting;
@@ -130,7 +129,7 @@ public Collection drain() {
   private RawLogMessage find(
       long keyHash, String logLevel, String message, @Nullable Throwable throwable) {
     StackTraceElement[] stackTrace = null;
-    for (RawLogMessage entry : hashIterable(rawLogMessages, keyHash)) {
+    for (RawLogMessage entry : ConcurrentHashtable.hashIterable(rawLogMessages, keyHash)) {
       if (!Objects.equals(logLevel, entry.logLevel) || !Objects.equals(message, entry.message)) {
         continue;
       }

From f1d5316a54c3779a69958ae6b4d07ae933b42322 Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 11:04:17 -0400
Subject: [PATCH 08/13] Note removeIf as a candidate for the same per-bucket
 lock tradeoff as drain

No current caller puts removeIf on a contended path, so leave it as a
whole-sweep lock for now, but leave a pointer for future revisit since
the restructuring done for drain() would apply the same way here.

Co-Authored-By: Claude Sonnet 5 
---
 .../main/java/datadog/trace/util/ConcurrentHashtable.java  | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
index ea1c58888e7..82578243681 100644
--- a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
+++ b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
@@ -1733,6 +1733,13 @@ public static > boolean removeIf(
    * {@link #removeIf(AtomicReferenceArray, AtomicInteger, Predicate)} variant for callers tracking
    * occupancy with a {@link State} instead of a bare counter — used by {@link D1#removeIf} and
    * {@link D2#removeIf}.
+   *
+   * 
    TODO: no current caller holds the lock across a full sweep the way {@code drain(State, ...)}
+   * used to before this PR moved it to a per-bucket lock (see that method's Javadoc for the
+   * tradeoff). If a future caller puts {@code removeIf} on a contended path, revisit whether the
+   * same per-bucket restructuring is worth it here — note it would drop the "predicate sees a
+   * stable table" guarantee documented above, so benchmark and weigh that against `drain`'s
+   * measured writer-throughput win before making the change.
    */
   @StrategyConsumer
   public static > boolean removeIf(

From fa44aac4ffd65b710f451afd7871a7f7bed80fdf Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 11:57:31 -0400
Subject: [PATCH 09/13] Trim inline comments per review feedback

- LogCollector.addLogMessage: drop the Reservation/drain interaction
  detail from the slow-path comment; it doesn't belong at this call
  site.
- ConcurrentHashtable.removeIf: move the per-bucket-locking TODO from
  the Javadoc into a shorter body comment.

Co-Authored-By: Claude Sonnet 5 
---
 .../java/datadog/trace/api/telemetry/LogCollector.java   | 6 +-----
 .../java/datadog/trace/util/ConcurrentHashtable.java     | 9 ++-------
 2 files changed, 3 insertions(+), 12 deletions(-)

diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
index b42d8f513de..0490df0cde4 100644
--- a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
+++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
@@ -62,11 +62,7 @@ public void addLogMessage(
       return;
     }
 
-    // Slow path after a miss: tryReserve holds the table write lock for the reservation's whole
-    // lifetime, so concurrent reservations for the same logical duplicate are serialized with each
-    // other, with drain(), and with the locked find-or-insert inside Reservation#finish() -- a
-    // losing reservation cancels immediately instead of transiently inflating size and starving a
-    // genuinely distinct concurrent insert. finish() does its own locked comparison, so there's no
+    // Slow path after a miss: tryGetOrInsertOrNull does its own locked comparison, so there's no
     // need to repeat find() here first.
     try (Reservation reservation = ConcurrentHashtable.tryReserve(rawLogMessages)) {
       if (!reservation.isReserved()) {
diff --git a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
index 82578243681..095c421bf74 100644
--- a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
+++ b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
@@ -1733,17 +1733,12 @@ public static > boolean removeIf(
    * {@link #removeIf(AtomicReferenceArray, AtomicInteger, Predicate)} variant for callers tracking
    * occupancy with a {@link State} instead of a bare counter — used by {@link D1#removeIf} and
    * {@link D2#removeIf}.
-   *
-   * 
    TODO: no current caller holds the lock across a full sweep the way {@code drain(State, ...)}
-   * used to before this PR moved it to a per-bucket lock (see that method's Javadoc for the
-   * tradeoff). If a future caller puts {@code removeIf} on a contended path, revisit whether the
-   * same per-bucket restructuring is worth it here — note it would drop the "predicate sees a
-   * stable table" guarantee documented above, so benchmark and weigh that against `drain`'s
-   * measured writer-throughput win before making the change.
    */
   @StrategyConsumer
   public static > boolean removeIf(
       @Nonnull State state, @Strategy @Nonnull Predicate predicate) {
+    // TODO: no caller contends on this lock today; if one does, consider the same per-bucket
+    // locking drain() uses, trading the "predicate sees a stable table" guarantee for throughput.
     AtomicReferenceArray buckets = state.buckets;
     ReentrantLock lock = getTableWriteLock(state);
     lock.lock();

From 40125241d919ea70918cb38c1b31f6728015ba87 Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 12:56:01 -0400
Subject: [PATCH 10/13] Add test coverage for remaining Reservation overloads

Closes the jacocoTestCoverageVerification gap on
ConcurrentHashtable.Reservation by exercising the previously-untested
2-key tryGetOrInsertOrNull overload, the pre-built-entry
tryGetOrInsertOrNull/tryGetOrInsert overloads, and the 2/3/4-key
tryGetOrInsert (Maybe-wrapping) overloads.

Co-Authored-By: Claude Sonnet 5 
---
 .../ConcurrentHashtableReservationTest.java   | 98 +++++++++++++++++++
 1 file changed, 98 insertions(+)

diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
index f61876b2e66..cae166e38fc 100644
--- a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
+++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
@@ -118,6 +118,22 @@ void closeOnAnAbsentReservationIsANoOp() {
     assertEquals(0, ConcurrentHashtable.estimateSize(state));
   }
 
+  private static final class TwoPartEntry extends ConcurrentHashtable.Entry {
+    final String a;
+    final String b;
+
+    TwoPartEntry(String a, String b) {
+      super(HashingUtils.hash(a, b));
+      this.a = a;
+      this.b = b;
+    }
+
+    @Override
+    public boolean matches(@Nonnull TwoPartEntry other) {
+      return a.equals(other.a) && b.equals(other.b);
+    }
+  }
+
   private static final class ThreePartEntry extends ConcurrentHashtable.Entry {
     final String a;
     final String b;
@@ -181,4 +197,86 @@ void tryGetOrInsertOrNullSupportsUpToFourComponents() {
     assertEquals("y", four.c);
     assertEquals("z", four.d);
   }
+
+  @Test
+  void tryGetOrInsertOrNullSupportsTwoComponents() {
+    ConcurrentHashtable.State state =
+        ConcurrentHashtable.createBounded(TwoPartEntry.class, 2);
+    TwoPartEntry two;
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+      two = r.tryGetOrInsertOrNull("x", "y", TwoPartEntry::new);
+    }
+    assertEquals("x", two.a);
+    assertEquals("y", two.b);
+  }
+
+  @Test
+  void tryGetOrInsertOrNullOnPrebuiltEntry() {
+    ConcurrentHashtable.State state =
+        ConcurrentHashtable.createBounded(TestEntry.class, 1);
+    TestEntry inserted;
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+      inserted = r.tryGetOrInsertOrNull(new TestEntry(1));
+    }
+    assertEquals(1, inserted.value);
+    assertEquals(1, ConcurrentHashtable.estimateSize(state));
+  }
+
+  @Test
+  void tryGetOrInsertOnPrebuiltEntryWrapsResultInMaybe() {
+    ConcurrentHashtable.State state =
+        ConcurrentHashtable.createBounded(TestEntry.class, 1);
+
+    Maybe present;
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+      present = r.tryGetOrInsert(new TestEntry(1));
+    }
+    assertTrue(present.isPresent());
+    assertEquals(1, present.getOrNull().value);
+
+    Maybe absent;
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+      absent = r.tryGetOrInsert(new TestEntry(2));
+    }
+    assertFalse(absent.isPresent());
+    assertNull(absent.getOrNull());
+  }
+
+  @Test
+  void tryGetOrInsertWrapsResultInMaybeForTwoThreeAndFourComponents() {
+    ConcurrentHashtable.State state2 =
+        ConcurrentHashtable.createBounded(TwoPartEntry.class, 2);
+    Maybe two;
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state2)) {
+      two = r.tryGetOrInsert("x", "y", TwoPartEntry::new);
+    }
+    assertTrue(two.isPresent());
+    assertEquals("x", two.getOrNull().a);
+    assertEquals("y", two.getOrNull().b);
+
+    ConcurrentHashtable.State state3 =
+        ConcurrentHashtable.createBounded(ThreePartEntry.class, 2);
+    Maybe three;
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserve(state3)) {
+      three = r.tryGetOrInsert("x", "y", "z", ThreePartEntry::new);
+    }
+    assertTrue(three.isPresent());
+    assertEquals("x", three.getOrNull().a);
+    assertEquals("y", three.getOrNull().b);
+    assertEquals("z", three.getOrNull().c);
+
+    ConcurrentHashtable.State state4 =
+        ConcurrentHashtable.createBounded(FourPartEntry.class, 2);
+    Maybe four;
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserve(state4)) {
+      four = r.tryGetOrInsert("w", "x", "y", "z", FourPartEntry::new);
+    }
+    assertTrue(four.isPresent());
+    assertEquals("w", four.getOrNull().a);
+    assertEquals("x", four.getOrNull().b);
+    assertEquals("y", four.getOrNull().c);
+    assertEquals("z", four.getOrNull().d);
+  }
 }

From 16b56c8c31305c22f63f481fa510fe951eda0ce2 Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 13:25:13 -0400
Subject: [PATCH 11/13] Fix duplicate-undercounting race in ConcurrentHashtable
 reservations

tryReserve/tryReserveSlot could increment the size counter and then
reject the reservation if the atomic check failed, without releasing
the slot; composed via the atomic sizeManager.tryReserve() check
instead so success/failure and the increment happen together.

Add tryReserveFor(state, keyHash), which closes a race where a full
table rejects a reservation without ever checking whether the
concurrent entry that filled the table is actually a match for the
caller's key. It keeps the write lock open on a full table when the
target bucket is non-empty, so Reservation.finish() can still scan
for a concurrent duplicate even though no slot was claimed.
LogCollector.addLogMessage now uses tryReserveFor instead of
hand-rolling its own lock-free duplicate recheck.

Also sync D1/D2 drain() Javadoc with the actual per-bucket locking
behavior (previously described whole-sweep locking).

Co-Authored-By: Claude Sonnet 5 
---
 .../trace/api/telemetry/LogCollector.java     |  16 ++-
 .../trace/util/ConcurrentHashtable.java       | 135 +++++++++++++-----
 .../ConcurrentHashtableReservationTest.java   |  86 +++++++++++
 3 files changed, 198 insertions(+), 39 deletions(-)

diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
index 0490df0cde4..eeaf2bcdc15 100644
--- a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
+++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
@@ -62,19 +62,21 @@ public void addLogMessage(
       return;
     }
 
-    // Slow path after a miss: tryGetOrInsertOrNull does its own locked comparison, so there's no
-    // need to repeat find() here first.
-    try (Reservation reservation = ConcurrentHashtable.tryReserve(rawLogMessages)) {
-      if (!reservation.isReserved()) {
-        // TODO: We could emit a metric for dropped logs.
-        return;
-      }
+    // Slow path after a miss: tryGetOrInsertOrNull does its own locked comparison -- including,
+    // via tryReserveFor, a recheck against a concurrent duplicate even when the table looks full --
+    // so there's no need to repeat find() here first.
+    try (Reservation reservation =
+        ConcurrentHashtable.tryReserveFor(rawLogMessages, keyHash)) {
       // Built zeroed, so this occurrence can be counted uniformly below whether or not
       // tryGetOrInsertOrNull ends up returning this instance or an existing match.
       rawLogMessage =
           reservation.tryGetOrInsertOrNull(
               new RawLogMessage(
                   logLevel, message, throwable, tags, System.currentTimeMillis() / 1000));
+      if (rawLogMessage == null) {
+        // TODO: We could emit a metric for dropped logs.
+        return;
+      }
       rawLogMessage.increment();
     }
   }
diff --git a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
index 095c421bf74..30f55eb1896 100644
--- a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
+++ b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
@@ -395,10 +395,14 @@ public boolean removeIf(@Strategy @Nonnull Predicate predicate)
     /**
      * Removes all entries and invokes {@code drainedEntryConsumer} for each one.
      *
-     * 
    The drain holds the table write lock while detaching buckets and invoking the consumer.
-     * For each removed entry, the consumer is invoked synchronously after its bucket is detached.
-     * Capacity is released only after all invocations return. The consumer should be quick and must
-     * not throw; failures are not rolled back.
+     * 
    The lock is acquired and released bucket by bucket rather than held for the whole sweep,
+     * so a table-level operation (an insert, a reservation, an eviction) can interleave between
+     * buckets instead of waiting out the entire drain. This means the drain is no longer an atomic
+     * snapshot: an entry inserted into a not-yet-visited bucket while the drain is in progress is
+     * swept up too, and each bucket's capacity slot is released immediately before its consumer
+     * invocation runs, rather than once for the whole sweep at the end. The consumer should be
+     * quick and must not throw; entries not yet reached are left in the table with their capacity
+     * still counted, but entries already detached keep their capacity released regardless.
      *
      * 
    Use {@link #drain(Object, BiConsumer)} to avoid a capturing lambda.
      *
@@ -409,10 +413,13 @@ public void drain(@Strategy @Nonnull Consumer drainedEntryConsum
     }
 
     /**
-     * Context-passing {@link #drain(Consumer)}. The drain holds the table write lock while
-     * detaching buckets and invoking {@code drainedEntryConsumer}. For each removed entry, the
-     * consumer is invoked synchronously after its bucket is detached. Capacity is released only
-     * after all invocations return.
+     * Context-passing {@link #drain(Consumer)}. The lock is acquired and released bucket by bucket
+     * rather than held for the whole sweep, so a table-level operation (an insert, a reservation,
+     * an eviction) can interleave between buckets instead of waiting out the entire drain. This
+     * means the drain is no longer an atomic snapshot: an entry inserted into a not-yet-visited
+     * bucket while the drain is in progress is swept up too, and each bucket's capacity slot is
+     * released immediately before its consumer invocation runs, rather than once for the whole
+     * sweep at the end.
      *
      * 
    Pass a non-capturing {@link BiConsumer} (typically a {@code static final}) plus the
      * accumulator as {@code context} (e.g. the target list or event builder) to avoid a
@@ -737,10 +744,14 @@ public boolean removeIf(@Strategy @Nonnull Predicate predicate)
     /**
      * Removes all entries and invokes {@code drainedEntryConsumer} for each one.
      *
-     * 
    The drain holds the table write lock while detaching buckets and invoking the consumer.
-     * For each removed entry, the consumer is invoked synchronously after its bucket is detached.
-     * Capacity is released only after all invocations return. The consumer should be quick and must
-     * not throw; failures are not rolled back.
+     * 
    The lock is acquired and released bucket by bucket rather than held for the whole sweep,
+     * so a table-level operation (an insert, a reservation, an eviction) can interleave between
+     * buckets instead of waiting out the entire drain. This means the drain is no longer an atomic
+     * snapshot: an entry inserted into a not-yet-visited bucket while the drain is in progress is
+     * swept up too, and each bucket's capacity slot is released immediately before its consumer
+     * invocation runs, rather than once for the whole sweep at the end. The consumer should be
+     * quick and must not throw; entries not yet reached are left in the table with their capacity
+     * still counted, but entries already detached keep their capacity released regardless.
      *
      * 
    Use {@link #drain(Object, BiConsumer)} to avoid a capturing lambda.
      *
@@ -751,10 +762,13 @@ public void drain(@Strategy @Nonnull Consumer drainedEntryConsum
     }
 
     /**
-     * Context-passing {@link #drain(Consumer)}. The drain holds the table write lock while
-     * detaching buckets and invoking {@code drainedEntryConsumer}. For each removed entry, the
-     * consumer is invoked synchronously after its bucket is detached. Capacity is released only
-     * after all invocations return.
+     * Context-passing {@link #drain(Consumer)}. The lock is acquired and released bucket by bucket
+     * rather than held for the whole sweep, so a table-level operation (an insert, a reservation,
+     * an eviction) can interleave between buckets instead of waiting out the entire drain. This
+     * means the drain is no longer an atomic snapshot: an entry inserted into a not-yet-visited
+     * bucket while the drain is in progress is swept up too, and each bucket's capacity slot is
+     * released immediately before its consumer invocation runs, rather than once for the whole
+     * sweep at the end.
      *
      * 
    Pass a non-capturing {@link BiConsumer} (typically a {@code static final}) plus the
      * accumulator as {@code context} (e.g. the target list or event builder) to avoid a
@@ -1101,16 +1115,51 @@ public static > boolean tryReserveSlot(
   public static > Reservation tryReserve(
       @Nonnull State state) {
     if (state.sizeManager.isFull()) {
-      return new Reservation<>(null);
+      return new Reservation<>(null, false);
     }
     ReentrantLock lock = state.writeLock;
     lock.lock();
-    if (state.sizeManager.isFull()) {
+    if (!state.sizeManager.tryReserve()) {
       lock.unlock();
-      return new Reservation<>(null);
+      return new Reservation<>(null, false);
+    }
+    return new Reservation<>(state, true);
+  }
+
+  /**
+   * Like {@link #tryReserve}, but lets a rejection on a full table still catch a concurrent
+   * duplicate for {@code keyHash} rather than dropping it outright.
+   *
+   * 
    When the table looks full, {@code keyHash}'s bucket is checked first: if it's empty, nothing
+   * could possibly match this key, so this returns a definitively empty reservation without ever
+   * taking the write lock -- the same fast reject as {@link #tryReserve}. If the bucket is
+   * populated, the write lock is taken and the capacity re-checked; if the table is genuinely still
+   * full, the lock is kept open (rather than released immediately) so {@link Reservation#finish}
+   * can still scan that bucket the next time the caller supplies a candidate entry. This closes the
+   * race where a concurrent insert for the same logical duplicate lands in the exact window between
+   * this caller's own lock-free scan and its reservation attempt -- a race {@link #tryReserve}
+   * can't detect, since it has no key to check the target bucket against.
+   *
+   * 
    The returned reservation still reports {@link Reservation#isReserved()} in the
+   * lock-held-but-full case, but has no slot to insert into: {@code tryGetOrInsertOrNull}/{@code
+   * tryGetOrInsert} either return the concurrent match found under the lock, or {@code null} --
+   * never a newly linked entry.
+   *
+   * @param keyHash hash of the key the caller is about to look up or insert
+   */
+  @Nonnull
+  public static > Reservation tryReserveFor(
+      @Nonnull State state, long keyHash) {
+    if (state.sizeManager.isFull()
+        && bucketAt(state, bucketIndex(state.buckets, keyHash)) == null) {
+      return new Reservation<>(null, false);
     }
-    state.sizeManager.increment();
-    return new Reservation<>(state);
+    ReentrantLock lock = state.writeLock;
+    lock.lock();
+    if (state.sizeManager.tryReserve()) {
+      return new Reservation<>(state, true);
+    }
+    return new Reservation<>(state, false);
   }
 
   /**
@@ -1135,13 +1184,23 @@ public static > Reservation tryReserve(
    */
   public static final class Reservation> implements AutoCloseable {
     @Nullable private final State state;
+    private final boolean slotClaimed;
     private boolean consumed;
 
-    private Reservation(@Nullable State state) {
+    private Reservation(@Nullable State state, boolean slotClaimed) {
       this.state = state;
+      this.slotClaimed = slotClaimed;
     }
 
-    /** {@code true} if this is a real, claimed reservation rather than an empty one. */
+    /**
+     * {@code true} unless this reservation is definitively empty -- i.e. {@link #tryReserveFor}
+     * could rule out a match for the target key without even taking the write lock. A {@code true}
+     * reservation may still have no slot to insert into: {@link #tryReserveFor} keeps the lock open
+     * on a full table when the target bucket is populated, purely so {@link #finish} can still scan
+     * it for a concurrent match. Callers should call {@code tryGetOrInsertOrNull}/{@code
+     * tryGetOrInsert} and check its result rather than assume {@code true} means a slot was
+     * claimed.
+     */
     public boolean isReserved() {
       return state != null;
     }
@@ -1180,10 +1239,12 @@ public  TEntry tryGetOrInsertOrNull(
 
     /**
      * Two key components. Builds {@code factory.apply(...)} (skipped entirely if this reservation
-     * is empty — the table was full) and either links the result as a new entry or discards it in
-     * favor of an existing match found under the write lock. Returns {@code null} only when this
-     * reservation is empty; otherwise always returns a real entry (the newly built one, or the
-     * concurrent match).
+     * is definitively empty — see {@link #isReserved()}) and either links the result as a new entry
+     * or discards it in favor of an existing match found under the write lock. Returns {@code null}
+     * when there is neither a slot to claim nor a match to return -- either because this
+     * reservation is definitively empty, or (see {@link #tryReserveFor}) because the table was
+     * still full even under the lock and no concurrent match was found either; otherwise always
+     * returns a real entry (the newly built one, or the concurrent match).
      *
      * 
    Building the entry here, after the reservation already succeeded, keeps the write lock's
      * critical section limited to the comparison/link/discard decision rather than whatever
@@ -1273,9 +1334,15 @@ public  Maybe tryGetOrInsert(
 
     /**
      * Runs the locked comparison/link decision. The write lock is already held -- acquired by
-     * {@link #tryReserve} and not yet released -- so this needs no {@code synchronized}/{@code
-     * lock()} of its own.
+     * {@link #tryReserve}/{@link #tryReserveFor} and not yet released -- so this needs no {@code
+     * synchronized}/{@code lock()} of its own.
+     *
+     * 
    Always scans for a match first, whether or not a slot was actually claimed: a {@link
+     * #tryReserveFor} reservation with no slot still holds the lock specifically so this scan can
+     * run. Only links {@code newEntry} when a slot was claimed; otherwise a miss here means there
+     * really is nothing to return, so this returns {@code null}.
      */
+    @Nullable
     private TEntry finish(@Nonnull TEntry newEntry) {
       int index = bucketIndex(state.buckets, newEntry.keyHash);
       for (TEntry curEntry = bucketAt(state, index); curEntry != null; curEntry = curEntry.next()) {
@@ -1283,14 +1350,18 @@ private TEntry finish(@Nonnull TEntry newEntry) {
           return curEntry;
         }
       }
+      if (!slotClaimed) {
+        return null;
+      }
       insertHeadEntryAt(state, index, newEntry);
       consumed = true;
       return newEntry;
     }
 
     /**
-     * Releases the write lock {@link #tryReserve} acquired, giving back the slot first if it was
-     * never consumed. A no-op on an empty reservation, which never acquired the lock.
+     * Releases the write lock {@link #tryReserve}/{@link #tryReserveFor} acquired, giving back a
+     * claimed slot first if it was never consumed. A no-op on a definitively empty reservation,
+     * which never acquired the lock.
      */
     @Override
     public void close() {
@@ -1298,7 +1369,7 @@ public void close() {
         return;
       }
       try {
-        if (!consumed) {
+        if (slotClaimed && !consumed) {
           state.sizeManager.decrement();
         }
       } finally {
diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
index cae166e38fc..91fa991f07b 100644
--- a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
+++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
@@ -279,4 +279,90 @@ void tryGetOrInsertWrapsResultInMaybeForTwoThreeAndFourComponents() {
     assertEquals("y", four.getOrNull().c);
     assertEquals("z", four.getOrNull().d);
   }
+
+  @Test
+  void tryReserveForSkipsTheLockWhenTheTargetBucketIsDefinitelyEmpty() {
+    // capacity 3 rounds up to 4 buckets, so one bucket stays empty once the table is full.
+    ConcurrentHashtable.State state =
+        ConcurrentHashtable.createBounded(TestEntry.class, 3);
+    for (int keyHash = 1; keyHash <= 3; keyHash++) {
+      try (ConcurrentHashtable.Reservation r =
+          ConcurrentHashtable.tryReserveFor(state, keyHash)) {
+        r.tryGetOrInsertOrNull(keyHash, TestEntry::new);
+      }
+    }
+    assertTrue(ConcurrentHashtable.isFull(state));
+
+    AtomicInteger factoryCalls = new AtomicInteger();
+    TestEntry result;
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserveFor(state, 4)) {
+      assertFalse(r.isReserved());
+      result =
+          r.tryGetOrInsertOrNull(
+              4,
+              v -> {
+                factoryCalls.incrementAndGet();
+                return new TestEntry(v);
+              });
+    }
+    assertNull(result);
+    assertEquals(0, factoryCalls.get());
+    assertEquals(3, ConcurrentHashtable.estimateSize(state));
+  }
+
+  @Test
+  void tryReserveForFindsAConcurrentDuplicateEvenWhenTheTableLooksFull() {
+    ConcurrentHashtable.State state =
+        ConcurrentHashtable.createBounded(TestEntry.class, 3);
+    TestEntry existing;
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserveFor(state, 1)) {
+      existing = r.tryGetOrInsertOrNull(1, TestEntry::new);
+    }
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserveFor(state, 2)) {
+      r.tryGetOrInsertOrNull(2, TestEntry::new);
+    }
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserveFor(state, 3)) {
+      r.tryGetOrInsertOrNull(3, TestEntry::new);
+    }
+    assertTrue(ConcurrentHashtable.isFull(state));
+
+    // Reserving for the same keyHash again simulates a concurrent duplicate insert landing just
+    // before the caller's own reservation attempt.
+    TestEntry match;
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserveFor(state, 1)) {
+      assertTrue(r.isReserved());
+      match = r.tryGetOrInsertOrNull(new TestEntry(1));
+    }
+    assertSame(existing, match);
+    assertEquals(3, ConcurrentHashtable.estimateSize(state));
+  }
+
+  @Test
+  void tryReserveForReturnsNullWithoutOvercountingWhenGenuinelyFull() {
+    // capacity 3 rounds up to 4 buckets; keyHash 1 and 5 collide on the same bucket (index 1) but
+    // are logically distinct keys (TestEntry.matches compares by value).
+    ConcurrentHashtable.State state =
+        ConcurrentHashtable.createBounded(TestEntry.class, 3);
+    for (int keyHash = 1; keyHash <= 3; keyHash++) {
+      try (ConcurrentHashtable.Reservation r =
+          ConcurrentHashtable.tryReserveFor(state, keyHash)) {
+        r.tryGetOrInsertOrNull(keyHash, TestEntry::new);
+      }
+    }
+    assertTrue(ConcurrentHashtable.isFull(state));
+
+    TestEntry result;
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserveFor(state, 5)) {
+      assertTrue(r.isReserved());
+      result = r.tryGetOrInsertOrNull(new TestEntry(5));
+    }
+    assertNull(result);
+    assertEquals(3, ConcurrentHashtable.estimateSize(state));
+  }
 }

From b226804a3ff7735f58821d8ffce7e3c927023596 Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 13:59:28 -0400
Subject: [PATCH 12/13] Add a near-capacity JMH benchmark for LogCollector

The existing LogCollectorBenchmark cases all hit the lock-free find()
fast path on every call, so they never exercise tryReserveFor's
locked recheck. Add variedKeysNearCapacity, which cycles through more
distinct keys than the table has capacity for, forcing most calls
through the locked path this PR's fix touches.

Co-Authored-By: Claude Sonnet 5 
---
 .../api/telemetry/LogCollectorBenchmark.java  | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/internal-api/src/jmh/java/datadog/trace/api/telemetry/LogCollectorBenchmark.java b/internal-api/src/jmh/java/datadog/trace/api/telemetry/LogCollectorBenchmark.java
index 81d2a9180b6..c305ddc81d5 100644
--- a/internal-api/src/jmh/java/datadog/trace/api/telemetry/LogCollectorBenchmark.java
+++ b/internal-api/src/jmh/java/datadog/trace/api/telemetry/LogCollectorBenchmark.java
@@ -60,4 +60,47 @@ public void unsupportedOperationException(CollectorState state) {
   static void unsupportedOperation() {
     throw new UnsupportedOperationException();
   }
+
+  /**
+   * Exercises the near-capacity path the other benchmarks skip: capacity is well below the number
+   * of distinct keys in play, so once warmed up the table stays full and most calls miss {@code
+   * find()}'s lock-free scan and fall through to {@code tryReserveFor} -- including its locked
+   * recheck for a concurrent duplicate. {@link #duplicateWithoutException} and friends only ever
+   * hit the lock-free fast path, so they don't touch that code at all.
+   */
+  @State(Scope.Benchmark)
+  public static class ContendedCollectorState {
+    static final int N_KEYS = 32;
+    static final String[] MESSAGES = new String[N_KEYS];
+
+    static {
+      for (int i = 0; i < N_KEYS; i++) {
+        MESSAGES[i] = "message-" + i;
+      }
+    }
+
+    LogCollector collector;
+
+    @Setup(Level.Iteration)
+    public void setup() {
+      // Capacity well below N_KEYS keeps the table full/near-full once warmed up.
+      collector = new LogCollector(8);
+    }
+  }
+
+  @State(Scope.Thread)
+  public static class KeyCursorState {
+    int cursor;
+
+    int next() {
+      int i = cursor;
+      cursor = (i + 1) % ContendedCollectorState.N_KEYS;
+      return i;
+    }
+  }
+
+  @Benchmark
+  public void variedKeysNearCapacity(ContendedCollectorState state, KeyCursorState cursor) {
+    state.collector.addLogMessage("error", ContendedCollectorState.MESSAGES[cursor.next()], null);
+  }
 }

From c87c904a258e099fbcf96fa9efdb499c57056db4 Mon Sep 17 00:00:00 2001
From: Douglas Q Hawkins 
Date: Fri, 11 Sep 2026 17:06:43 -0400
Subject: [PATCH 13/13] Merge tryReserve and tryReserveFor into a single
 keyHash-aware method

Every real caller already had a keyHash on hand, so the plain
no-keyHash tryReserve was pure duplication once tryReserveFor
existed to fix the undercounting race. Drop it and keep the single
merged method under the original tryReserve name.

Co-Authored-By: Claude Sonnet 5 
---
 .../api/telemetry/LogCollectorBenchmark.java  |  6 +-
 .../ConcurrentHashtableDrainBenchmark.java    |  3 +-
 .../ConcurrentHashtableFindBenchmark.java     |  3 +-
 .../trace/api/telemetry/LogCollector.java     |  4 +-
 .../trace/util/ConcurrentHashtable.java       | 99 +++++++------------
 .../ConcurrentHashtableReservationTest.java   | 77 ++++++++-------
 6 files changed, 85 insertions(+), 107 deletions(-)

diff --git a/internal-api/src/jmh/java/datadog/trace/api/telemetry/LogCollectorBenchmark.java b/internal-api/src/jmh/java/datadog/trace/api/telemetry/LogCollectorBenchmark.java
index c305ddc81d5..3754ee33fde 100644
--- a/internal-api/src/jmh/java/datadog/trace/api/telemetry/LogCollectorBenchmark.java
+++ b/internal-api/src/jmh/java/datadog/trace/api/telemetry/LogCollectorBenchmark.java
@@ -64,9 +64,9 @@ static void unsupportedOperation() {
   /**
    * Exercises the near-capacity path the other benchmarks skip: capacity is well below the number
    * of distinct keys in play, so once warmed up the table stays full and most calls miss {@code
-   * find()}'s lock-free scan and fall through to {@code tryReserveFor} -- including its locked
-   * recheck for a concurrent duplicate. {@link #duplicateWithoutException} and friends only ever
-   * hit the lock-free fast path, so they don't touch that code at all.
+   * find()}'s lock-free scan and fall through to {@code tryReserve} -- including its locked recheck
+   * for a concurrent duplicate. {@link #duplicateWithoutException} and friends only ever hit the
+   * lock-free fast path, so they don't touch that code at all.
    */
   @State(Scope.Benchmark)
   public static class ContendedCollectorState {
diff --git a/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
index 6a1e28c875f..97ee0914e34 100644
--- a/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
+++ b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableDrainBenchmark.java
@@ -135,7 +135,8 @@ private static void write(ConcurrentHashtable.State table, WriterSta
     for (DrainEntry entry : ConcurrentHashtable.hashIterable(table, keyHash)) {
       return; // lock-free hit, mirroring LogCollector.find()
     }
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(table)) {
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserve(table, keyHash)) {
       if (r.isReserved()) {
         r.tryGetOrInsertOrNull(new DrainEntry(keyHash));
       }
diff --git a/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableFindBenchmark.java b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableFindBenchmark.java
index 3f9ccb79f00..801a79f548b 100644
--- a/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableFindBenchmark.java
+++ b/internal-api/src/jmh/java/datadog/trace/util/ConcurrentHashtableFindBenchmark.java
@@ -71,7 +71,8 @@ public static class SharedState {
     public void setUp() {
       table = ConcurrentHashtable.createBounded(FindEntry.class, CAPACITY);
       for (int i = 0; i < N_KEYS; ++i) {
-        ConcurrentHashtable.tryReserve(table).tryGetOrInsertOrNull(new FindEntry(KEY_HASHES[i], i));
+        ConcurrentHashtable.tryReserve(table, KEY_HASHES[i])
+            .tryGetOrInsertOrNull(new FindEntry(KEY_HASHES[i], i));
       }
     }
   }
diff --git a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
index eeaf2bcdc15..404e8eebb0f 100644
--- a/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
+++ b/internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java
@@ -63,10 +63,10 @@ public void addLogMessage(
     }
 
     // Slow path after a miss: tryGetOrInsertOrNull does its own locked comparison -- including,
-    // via tryReserveFor, a recheck against a concurrent duplicate even when the table looks full --
+    // via tryReserve, a recheck against a concurrent duplicate even when the table looks full --
     // so there's no need to repeat find() here first.
     try (Reservation reservation =
-        ConcurrentHashtable.tryReserveFor(rawLogMessages, keyHash)) {
+        ConcurrentHashtable.tryReserve(rawLogMessages, keyHash)) {
       // Built zeroed, so this occurrence can be counted uniformly below whether or not
       // tryGetOrInsertOrNull ends up returning this instance or an existing match.
       rawLogMessage =
diff --git a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
index 30f55eb1896..48cdecb1c65 100644
--- a/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
+++ b/internal-api/src/main/java/datadog/trace/util/ConcurrentHashtable.java
@@ -1079,12 +1079,12 @@ public static > boolean tryReserveSlot(
   }
 
   /**
-   * Claims one slot in {@code state} and returns a handle for completing the find-or-insert
-   * protocol, or an empty handle if the table is full. Never returns {@code null}, so this always
-   * composes with try-with-resources:
+   * Claims one slot in {@code state} for {@code keyHash} and returns a handle for completing the
+   * find-or-insert protocol, or an empty handle if the table is full. Never returns {@code null},
+   * so this always composes with try-with-resources:
    *
    * 
    {@code
-   * try (Reservation r = ConcurrentHashtable.tryReserve(state)) {
+   * try (Reservation r = ConcurrentHashtable.tryReserve(state, keyHash)) {
    *   return r.tryGetOrInsertOrNull(component1, component2, component3, TEntry::new);
    * }
    * }
    
@@ -1095,13 +1095,19 @@ public static > boolean tryReserveSlot(
    * are the source of truth), it just avoids paying for a lock and a factory call when a hit was
    * already visible lock-free.
    *
-   * 
    Checks {@link SizeManager#isFull()} lock-free first and returns an empty reservation
-   * immediately if the table already looks full, without touching the write lock. Otherwise it
+   * 
    Checks {@link SizeManager#isFull()} lock-free first. If the table looks full, {@code
+   * keyHash}'s bucket is checked too: if it's empty, nothing could possibly match this key, so this
+   * returns a definitively empty reservation without ever taking the write lock. Otherwise it
    * acquires the write lock, re-checks (the lock-free peek may be stale), and either reserves the
    * slot -- returning a real reservation that holds the lock open across this call, to be
-   * released later by {@link Reservation#finish}/{@link Reservation#close} -- or releases the lock
-   * and returns empty. Holding the lock for the reservation's whole lifetime (rather than just
-   * across this call, the way {@link SizeManager#tryReserve()} does) is what lets {@link
+   * released later by {@link Reservation#finish}/{@link Reservation#close} -- or, if the table is
+   * genuinely still full, keeps the lock open anyway (rather than releasing it immediately) so
+   * {@link Reservation#finish} can still scan {@code keyHash}'s bucket the next time the caller
+   * supplies a candidate entry. That's what closes the race where a concurrent insert for the same
+   * logical duplicate lands in the exact window between this caller's own lock-free scan and its
+   * reservation attempt: without the lock held open here, that concurrent duplicate would go
+   * uncounted. Holding the lock for the reservation's whole lifetime (rather than just across this
+   * call, the way {@link SizeManager#tryReserve()} does) is also what lets {@link
    * Reservation#tryGetOrInsertOrNull} skip a second, separately-locked comparison pass: the
    * reserve-then-insert-or-discard sequence is one uninterrupted critical section, so no concurrent
    * reservation for the same logical duplicate can slip in between.
@@ -1109,46 +1115,15 @@ public static > boolean tryReserveSlot(
    * 
    Always returns a non-null handle — even when the table is full — so the caller must check
    * {@link Reservation#isReserved()} (or simply call {@link Reservation#tryGetOrInsertOrNull},
    * which returns {@code null} on an absent reservation) rather than assume every reservation is
-   * real.
-   */
-  @Nonnull
-  public static > Reservation tryReserve(
-      @Nonnull State state) {
-    if (state.sizeManager.isFull()) {
-      return new Reservation<>(null, false);
-    }
-    ReentrantLock lock = state.writeLock;
-    lock.lock();
-    if (!state.sizeManager.tryReserve()) {
-      lock.unlock();
-      return new Reservation<>(null, false);
-    }
-    return new Reservation<>(state, true);
-  }
-
-  /**
-   * Like {@link #tryReserve}, but lets a rejection on a full table still catch a concurrent
-   * duplicate for {@code keyHash} rather than dropping it outright.
-   *
-   * 
    When the table looks full, {@code keyHash}'s bucket is checked first: if it's empty, nothing
-   * could possibly match this key, so this returns a definitively empty reservation without ever
-   * taking the write lock -- the same fast reject as {@link #tryReserve}. If the bucket is
-   * populated, the write lock is taken and the capacity re-checked; if the table is genuinely still
-   * full, the lock is kept open (rather than released immediately) so {@link Reservation#finish}
-   * can still scan that bucket the next time the caller supplies a candidate entry. This closes the
-   * race where a concurrent insert for the same logical duplicate lands in the exact window between
-   * this caller's own lock-free scan and its reservation attempt -- a race {@link #tryReserve}
-   * can't detect, since it has no key to check the target bucket against.
-   *
-   * 
    The returned reservation still reports {@link Reservation#isReserved()} in the
-   * lock-held-but-full case, but has no slot to insert into: {@code tryGetOrInsertOrNull}/{@code
-   * tryGetOrInsert} either return the concurrent match found under the lock, or {@code null} --
-   * never a newly linked entry.
+   * real. A reservation can also report {@link Reservation#isReserved()} {@code true} while still
+   * having no slot to insert into (the lock-held-but-full case above); either way, {@code
+   * tryGetOrInsertOrNull}/{@code tryGetOrInsert} return the concurrent match found under the lock,
+   * or {@code null} -- never a newly linked entry when no slot was claimed.
    *
    * @param keyHash hash of the key the caller is about to look up or insert
    */
   @Nonnull
-  public static > Reservation tryReserveFor(
+  public static > Reservation tryReserve(
       @Nonnull State state, long keyHash) {
     if (state.sizeManager.isFull()
         && bucketAt(state, bucketIndex(state.buckets, keyHash)) == null) {
@@ -1193,11 +1168,11 @@ private Reservation(@Nullable State state, boolean slotClaimed) {
     }
 
     /**
-     * {@code true} unless this reservation is definitively empty -- i.e. {@link #tryReserveFor}
-     * could rule out a match for the target key without even taking the write lock. A {@code true}
-     * reservation may still have no slot to insert into: {@link #tryReserveFor} keeps the lock open
-     * on a full table when the target bucket is populated, purely so {@link #finish} can still scan
-     * it for a concurrent match. Callers should call {@code tryGetOrInsertOrNull}/{@code
+     * {@code true} unless this reservation is definitively empty -- i.e. {@link #tryReserve} could
+     * rule out a match for the target key without even taking the write lock. A {@code true}
+     * reservation may still have no slot to insert into: {@link #tryReserve} keeps the lock open on
+     * a full table when the target bucket is populated, purely so {@link #finish} can still scan it
+     * for a concurrent match. Callers should call {@code tryGetOrInsertOrNull}/{@code
      * tryGetOrInsert} and check its result rather than assume {@code true} means a slot was
      * claimed.
      */
@@ -1211,7 +1186,7 @@ public boolean isReserved() {
      * a {@code Function}'s type argument:
      *
      * 
    {@code
-     * try (Reservation r = ConcurrentHashtable.tryReserve(state)) {
+     * try (Reservation r = ConcurrentHashtable.tryReserve(state, keyHash)) {
      *   if (!r.isReserved()) {
      *     return null;
      *   }
@@ -1242,9 +1217,9 @@ public  TEntry tryGetOrInsertOrNull(
      * is definitively empty — see {@link #isReserved()}) and either links the result as a new entry
      * or discards it in favor of an existing match found under the write lock. Returns {@code null}
      * when there is neither a slot to claim nor a match to return -- either because this
-     * reservation is definitively empty, or (see {@link #tryReserveFor}) because the table was
-     * still full even under the lock and no concurrent match was found either; otherwise always
-     * returns a real entry (the newly built one, or the concurrent match).
+     * reservation is definitively empty, or (see {@link #tryReserve}) because the table was still
+     * full even under the lock and no concurrent match was found either; otherwise always returns a
+     * real entry (the newly built one, or the concurrent match).
      *
      * 
    Building the entry here, after the reservation already succeeded, keeps the write lock's
      * critical section limited to the comparison/link/discard decision rather than whatever
@@ -1334,13 +1309,13 @@ public  Maybe tryGetOrInsert(
 
     /**
      * Runs the locked comparison/link decision. The write lock is already held -- acquired by
-     * {@link #tryReserve}/{@link #tryReserveFor} and not yet released -- so this needs no {@code
-     * synchronized}/{@code lock()} of its own.
+     * {@link #tryReserve} and not yet released -- so this needs no {@code synchronized}/{@code
+     * lock()} of its own.
      *
      * 
    Always scans for a match first, whether or not a slot was actually claimed: a {@link
-     * #tryReserveFor} reservation with no slot still holds the lock specifically so this scan can
-     * run. Only links {@code newEntry} when a slot was claimed; otherwise a miss here means there
-     * really is nothing to return, so this returns {@code null}.
+     * #tryReserve} reservation with no slot still holds the lock specifically so this scan can run.
+     * Only links {@code newEntry} when a slot was claimed; otherwise a miss here means there really
+     * is nothing to return, so this returns {@code null}.
      */
     @Nullable
     private TEntry finish(@Nonnull TEntry newEntry) {
@@ -1359,9 +1334,9 @@ private TEntry finish(@Nonnull TEntry newEntry) {
     }
 
     /**
-     * Releases the write lock {@link #tryReserve}/{@link #tryReserveFor} acquired, giving back a
-     * claimed slot first if it was never consumed. A no-op on a definitively empty reservation,
-     * which never acquired the lock.
+     * Releases the write lock {@link #tryReserve} acquired, giving back a claimed slot first if it
+     * was never consumed. A no-op on a definitively empty reservation, which never acquired the
+     * lock.
      */
     @Override
     public void close() {
diff --git a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
index 91fa991f07b..723ebed66fa 100644
--- a/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
+++ b/internal-api/src/test/java/datadog/trace/util/ConcurrentHashtableReservationTest.java
@@ -33,7 +33,7 @@ void tryGetOrInsertOrNullInsertsOnMissAndFindsOnHit() {
         ConcurrentHashtable.createBounded(TestEntry.class, 4);
 
     TestEntry first;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       assertTrue(r.isReserved());
       first = r.tryGetOrInsertOrNull(1, TestEntry::new);
     }
@@ -43,7 +43,7 @@ void tryGetOrInsertOrNullInsertsOnMissAndFindsOnHit() {
     // Reserving again for a key that already exists should discard the reservation and return the
     // existing entry, not double-insert or leak the claimed slot.
     TestEntry second;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       second = r.tryGetOrInsertOrNull(1, TestEntry::new);
     }
     assertSame(first, second);
@@ -52,20 +52,25 @@ void tryGetOrInsertOrNullInsertsOnMissAndFindsOnHit() {
 
   @Test
   void reserveOnFullTableIsAbsentAndSkipsTheFactory() {
+    // capacity 3 rounds up to 4 buckets, so one bucket stays empty once the table is full --
+    // that's what lets tryReserve rule out keyHash 4 without ever taking the write lock.
     ConcurrentHashtable.State state =
-        ConcurrentHashtable.createBounded(TestEntry.class, 1);
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
-      r.tryGetOrInsertOrNull(1, TestEntry::new);
+        ConcurrentHashtable.createBounded(TestEntry.class, 3);
+    for (int keyHash = 1; keyHash <= 3; keyHash++) {
+      try (ConcurrentHashtable.Reservation r =
+          ConcurrentHashtable.tryReserve(state, keyHash)) {
+        r.tryGetOrInsertOrNull(keyHash, TestEntry::new);
+      }
     }
     assertTrue(ConcurrentHashtable.isFull(state));
 
     AtomicInteger factoryCalls = new AtomicInteger();
     TestEntry result;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 4)) {
       assertFalse(r.isReserved());
       result =
           r.tryGetOrInsertOrNull(
-              2,
+              4,
               v -> {
                 factoryCalls.incrementAndGet();
                 return new TestEntry(v);
@@ -73,14 +78,14 @@ void reserveOnFullTableIsAbsentAndSkipsTheFactory() {
     }
     assertNull(result);
     assertEquals(0, factoryCalls.get());
-    assertEquals(1, ConcurrentHashtable.estimateSize(state));
+    assertEquals(3, ConcurrentHashtable.estimateSize(state));
   }
 
   @Test
   void closeCancelsAnUnconsumedReservation() {
     ConcurrentHashtable.State state =
         ConcurrentHashtable.createBounded(TestEntry.class, 1);
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       assertTrue(r.isReserved());
       // Deliberately not consuming the reservation.
     }
@@ -94,14 +99,14 @@ void tryGetOrInsertWrapsResultInMaybe() {
         ConcurrentHashtable.createBounded(TestEntry.class, 1);
 
     Maybe present;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       present = r.tryGetOrInsert(1, TestEntry::new);
     }
     assertTrue(present.isPresent());
     assertEquals(1, present.getOrNull().value);
 
     Maybe absent;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 2)) {
       absent = r.tryGetOrInsert(2, TestEntry::new);
     }
     assertFalse(absent.isPresent());
@@ -112,7 +117,7 @@ void tryGetOrInsertWrapsResultInMaybe() {
   void closeOnAnAbsentReservationIsANoOp() {
     ConcurrentHashtable.State state =
         ConcurrentHashtable.createBounded(TestEntry.class, 0);
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       assertFalse(r.isReserved());
     }
     assertEquals(0, ConcurrentHashtable.estimateSize(state));
@@ -178,7 +183,7 @@ void tryGetOrInsertOrNullSupportsUpToFourComponents() {
         ConcurrentHashtable.createBounded(ThreePartEntry.class, 2);
     ThreePartEntry three;
     try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserve(state3)) {
+        ConcurrentHashtable.tryReserve(state3, HashingUtils.hash("x", "y", "z"))) {
       three = r.tryGetOrInsertOrNull("x", "y", "z", ThreePartEntry::new);
     }
     assertEquals("x", three.a);
@@ -189,7 +194,7 @@ void tryGetOrInsertOrNullSupportsUpToFourComponents() {
         ConcurrentHashtable.createBounded(FourPartEntry.class, 2);
     FourPartEntry four;
     try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserve(state4)) {
+        ConcurrentHashtable.tryReserve(state4, HashingUtils.hash("w", "x", "y", "z"))) {
       four = r.tryGetOrInsertOrNull("w", "x", "y", "z", FourPartEntry::new);
     }
     assertEquals("w", four.a);
@@ -203,7 +208,8 @@ void tryGetOrInsertOrNullSupportsTwoComponents() {
     ConcurrentHashtable.State state =
         ConcurrentHashtable.createBounded(TwoPartEntry.class, 2);
     TwoPartEntry two;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserve(state, HashingUtils.hash("x", "y"))) {
       two = r.tryGetOrInsertOrNull("x", "y", TwoPartEntry::new);
     }
     assertEquals("x", two.a);
@@ -215,7 +221,7 @@ void tryGetOrInsertOrNullOnPrebuiltEntry() {
     ConcurrentHashtable.State state =
         ConcurrentHashtable.createBounded(TestEntry.class, 1);
     TestEntry inserted;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       inserted = r.tryGetOrInsertOrNull(new TestEntry(1));
     }
     assertEquals(1, inserted.value);
@@ -228,14 +234,14 @@ void tryGetOrInsertOnPrebuiltEntryWrapsResultInMaybe() {
         ConcurrentHashtable.createBounded(TestEntry.class, 1);
 
     Maybe present;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       present = r.tryGetOrInsert(new TestEntry(1));
     }
     assertTrue(present.isPresent());
     assertEquals(1, present.getOrNull().value);
 
     Maybe absent;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 2)) {
       absent = r.tryGetOrInsert(new TestEntry(2));
     }
     assertFalse(absent.isPresent());
@@ -247,7 +253,8 @@ void tryGetOrInsertWrapsResultInMaybeForTwoThreeAndFourComponents() {
     ConcurrentHashtable.State state2 =
         ConcurrentHashtable.createBounded(TwoPartEntry.class, 2);
     Maybe two;
-    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state2)) {
+    try (ConcurrentHashtable.Reservation r =
+        ConcurrentHashtable.tryReserve(state2, HashingUtils.hash("x", "y"))) {
       two = r.tryGetOrInsert("x", "y", TwoPartEntry::new);
     }
     assertTrue(two.isPresent());
@@ -258,7 +265,7 @@ void tryGetOrInsertWrapsResultInMaybeForTwoThreeAndFourComponents() {
         ConcurrentHashtable.createBounded(ThreePartEntry.class, 2);
     Maybe three;
     try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserve(state3)) {
+        ConcurrentHashtable.tryReserve(state3, HashingUtils.hash("x", "y", "z"))) {
       three = r.tryGetOrInsert("x", "y", "z", ThreePartEntry::new);
     }
     assertTrue(three.isPresent());
@@ -270,7 +277,7 @@ void tryGetOrInsertWrapsResultInMaybeForTwoThreeAndFourComponents() {
         ConcurrentHashtable.createBounded(FourPartEntry.class, 2);
     Maybe four;
     try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserve(state4)) {
+        ConcurrentHashtable.tryReserve(state4, HashingUtils.hash("w", "x", "y", "z"))) {
       four = r.tryGetOrInsert("w", "x", "y", "z", FourPartEntry::new);
     }
     assertTrue(four.isPresent());
@@ -281,13 +288,13 @@ void tryGetOrInsertWrapsResultInMaybeForTwoThreeAndFourComponents() {
   }
 
   @Test
-  void tryReserveForSkipsTheLockWhenTheTargetBucketIsDefinitelyEmpty() {
+  void tryReserveSkipsTheLockWhenTheTargetBucketIsDefinitelyEmpty() {
     // capacity 3 rounds up to 4 buckets, so one bucket stays empty once the table is full.
     ConcurrentHashtable.State state =
         ConcurrentHashtable.createBounded(TestEntry.class, 3);
     for (int keyHash = 1; keyHash <= 3; keyHash++) {
       try (ConcurrentHashtable.Reservation r =
-          ConcurrentHashtable.tryReserveFor(state, keyHash)) {
+          ConcurrentHashtable.tryReserve(state, keyHash)) {
         r.tryGetOrInsertOrNull(keyHash, TestEntry::new);
       }
     }
@@ -295,8 +302,7 @@ void tryReserveForSkipsTheLockWhenTheTargetBucketIsDefinitelyEmpty() {
 
     AtomicInteger factoryCalls = new AtomicInteger();
     TestEntry result;
-    try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserveFor(state, 4)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 4)) {
       assertFalse(r.isReserved());
       result =
           r.tryGetOrInsertOrNull(
@@ -312,20 +318,17 @@ void tryReserveForSkipsTheLockWhenTheTargetBucketIsDefinitelyEmpty() {
   }
 
   @Test
-  void tryReserveForFindsAConcurrentDuplicateEvenWhenTheTableLooksFull() {
+  void tryReserveFindsAConcurrentDuplicateEvenWhenTheTableLooksFull() {
     ConcurrentHashtable.State state =
         ConcurrentHashtable.createBounded(TestEntry.class, 3);
     TestEntry existing;
-    try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserveFor(state, 1)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       existing = r.tryGetOrInsertOrNull(1, TestEntry::new);
     }
-    try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserveFor(state, 2)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 2)) {
       r.tryGetOrInsertOrNull(2, TestEntry::new);
     }
-    try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserveFor(state, 3)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 3)) {
       r.tryGetOrInsertOrNull(3, TestEntry::new);
     }
     assertTrue(ConcurrentHashtable.isFull(state));
@@ -333,8 +336,7 @@ void tryReserveForFindsAConcurrentDuplicateEvenWhenTheTableLooksFull() {
     // Reserving for the same keyHash again simulates a concurrent duplicate insert landing just
     // before the caller's own reservation attempt.
     TestEntry match;
-    try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserveFor(state, 1)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 1)) {
       assertTrue(r.isReserved());
       match = r.tryGetOrInsertOrNull(new TestEntry(1));
     }
@@ -343,22 +345,21 @@ void tryReserveForFindsAConcurrentDuplicateEvenWhenTheTableLooksFull() {
   }
 
   @Test
-  void tryReserveForReturnsNullWithoutOvercountingWhenGenuinelyFull() {
+  void tryReserveReturnsNullWithoutOvercountingWhenGenuinelyFull() {
     // capacity 3 rounds up to 4 buckets; keyHash 1 and 5 collide on the same bucket (index 1) but
     // are logically distinct keys (TestEntry.matches compares by value).
     ConcurrentHashtable.State state =
         ConcurrentHashtable.createBounded(TestEntry.class, 3);
     for (int keyHash = 1; keyHash <= 3; keyHash++) {
       try (ConcurrentHashtable.Reservation r =
-          ConcurrentHashtable.tryReserveFor(state, keyHash)) {
+          ConcurrentHashtable.tryReserve(state, keyHash)) {
         r.tryGetOrInsertOrNull(keyHash, TestEntry::new);
       }
     }
     assertTrue(ConcurrentHashtable.isFull(state));
 
     TestEntry result;
-    try (ConcurrentHashtable.Reservation r =
-        ConcurrentHashtable.tryReserveFor(state, 5)) {
+    try (ConcurrentHashtable.Reservation r = ConcurrentHashtable.tryReserve(state, 5)) {
       assertTrue(r.isReserved());
       result = r.tryGetOrInsertOrNull(new TestEntry(5));
     }