Add find-or-insert Reservation API to ConcurrentHashtable

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.
  *
- * <p>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.
+ * <p>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.
  *
  * <p>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 @@
  *   <li>{@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.
  *   <li>{@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.
  * </ul>
  */
 @Fork(2)
@@ -73,25 +72,21 @@ 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<String> {
-    private static final AtomicLongFieldUpdater<CounterEntry> COUNT =
+    static final AtomicLongFieldUpdater<CounterEntry> COUNT =
         AtomicLongFieldUpdater.newUpdater(CounterEntry.class, "count");
 
     volatile long count;
 
     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<String, CounterEntry> 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

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<String> {
+  /**
+   * 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<String> {
     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<String, D1Entry> table;
+    ConcurrentHashtable.D1<String, LongEntry> table;
     ConcurrentHashMap<String, Long> concurrentHashMap;
     ConcurrentSkipListMap<String, Long> skipListMap;
     Map<String, Long> 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));
   }
 
   /**

internal-api/src/jmh/java/datadog/trace/util/ThreadSafeMapD2Benchmark.java

@@ -108,7 +108,7 @@ static final class D2Entry extends ConcurrentHashtable.D2.Entry<String, Integer>
    * 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<SupportEntry> {
     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. */

internal-api/src/main/java/datadog/trace/api/telemetry/LogCollector.java

@@ -1,16 +1,15 @@
 package datadog.trace.api.telemetry;
 
+import datadog.trace.util.ConcurrentHashtable;
 import datadog.trace.util.HashingUtils;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.slf4j.Marker;
 import org.slf4j.MarkerFactory;
@@ -20,8 +19,7 @@ public class LogCollector {
   public static final Marker EXCLUDE_TELEMETRY = MarkerFactory.getMarker("EXCLUDE_TELEMETRY");
   private static final int DEFAULT_MAX_CAPACITY = 10;
   private static final LogCollector INSTANCE = new LogCollector();
-  private final Map<RawLogMessage, AtomicInteger> rawLogMessages;
-  private final int maxCapacity;
+  private final ConcurrentHashtable.State<RawLogMessage> rawLogMessages;
 
   public static LogCollector get() {
     return INSTANCE;
@@ -35,8 +33,7 @@ private LogCollector() {
       value = "SING_SINGLETON_HAS_NONPRIVATE_CONSTRUCTOR",
       justification = "Usage in tests")
   LogCollector(int maxCapacity) {
-    this.maxCapacity = maxCapacity;
-    this.rawLogMessages = new ConcurrentHashMap<>(maxCapacity);
+    this.rawLogMessages = ConcurrentHashtable.createBounded(RawLogMessage.class, maxCapacity);
   }
 
   public void addLogMessage(String logLevel, String message, @Nullable Throwable throwable) {
@@ -54,57 +51,70 @@ public void addLogMessage(String logLevel, String message, @Nullable Throwable t
    */
   public void addLogMessage(
       String logLevel, String message, @Nullable Throwable throwable, @Nullable String tags) {
-    if (rawLogMessages.size() >= maxCapacity) {
-      // TODO: We could emit a metric for dropped logs.
-      return;
+    long keyHash = RawLogMessage.hash(logLevel, message, throwable);
+
+    // Memoized once per call and shared across every candidate below, rather than letting
+    // matchesKey call throwable.getStackTrace() (a defensive-copy allocation) on each comparison.
+    StackTraceElement[] throwableStackTrace = null;
+
+    // Lock-free scan first: most calls are re-observations of an already-seen message, so this
+    // avoids paying for a reservation and a RawLogMessage allocation on the common path.
+    for (RawLogMessage existing : ConcurrentHashtable.hashIterable(rawLogMessages, keyHash)) {
+      if (throwable != null && existing.throwable != null && existing.throwable != throwable) {
+        if (throwableStackTrace == null) {
+          throwableStackTrace = throwable.getStackTrace();
+        }
+      }
+      if (existing.matchesKey(logLevel, message, throwable, throwableStackTrace)) {
+        existing.count.incrementAndGet();
+        return;
+      }
+    }
+
+    try (ConcurrentHashtable.Reservation<RawLogMessage> reservation =
+        ConcurrentHashtable.tryReserve(rawLogMessages)) {
+      // TODO: We could emit a metric for dropped logs when the reservation is empty (table full).
+      RawLogMessage rawLogMessage =
+          reservation.tryGetOrInsertOrNull(logLevel, message, throwable, tags, RawLogMessage::new);
+      if (rawLogMessage != null) {
+        rawLogMessage.count.incrementAndGet();
+      }
     }
-    RawLogMessage rawLogMessage =
-        new RawLogMessage(logLevel, message, throwable, tags, System.currentTimeMillis() / 1000);
-    AtomicInteger count = rawLogMessages.computeIfAbsent(rawLogMessage, k -> new AtomicInteger());
-    count.incrementAndGet();
   }
 
   public Collection<RawLogMessage> drain() {
-    if (rawLogMessages.isEmpty()) {
+    int size = ConcurrentHashtable.estimateSize(rawLogMessages);
+    if (size == 0) {
       return Collections.emptyList();
     }
 
-    List<RawLogMessage> list = new ArrayList<>(rawLogMessages.size());
-    Iterator<Map.Entry<RawLogMessage, AtomicInteger>> iterator =
-        rawLogMessages.entrySet().iterator();
-
-    while (iterator.hasNext()) {
-      Map.Entry<RawLogMessage, AtomicInteger> entry = iterator.next();
-      RawLogMessage logMessage = entry.getKey();
-      // XXX: There might be lost writers to the counters under concurrency if another thread
-      // increments it
-      //      while we are reading it here. At the moment, we are not overdoing this to prevent some
-      // counter losses.
-      logMessage.count = entry.getValue().get();
-      iterator.remove();
-      list.add(logMessage);
-    }
-
+    List<RawLogMessage> list = new ArrayList<>(size);
+    ConcurrentHashtable.drain(rawLogMessages, list::add);
     return list;
   }
 
-  public static final class RawLogMessage {
+  public static final class RawLogMessage extends ConcurrentHashtable.Entry<RawLogMessage> {
     public final String message;
     public final String logLevel;
     public final Throwable throwable;
     public final String tags;
     public final long timestamp;
-    public int count;
+    public final AtomicInteger count = new AtomicInteger();
 
     private StackTraceElement[] cachedStackTrace = null;
 
     public RawLogMessage(
-        String logLevel, String message, Throwable throwable, String tags, long timestamp) {
+        String logLevel, String message, Throwable throwable, @Nullable String tags) {
+      super(hash(logLevel, message, throwable));
       this.logLevel = logLevel;
       this.message = message;
       this.throwable = throwable;
       this.tags = tags;
-      this.timestamp = timestamp;
+      this.timestamp = System.currentTimeMillis() / 1000;
+    }
+
+    static long hash(String logLevel, String message, @Nullable Throwable throwable) {
+      return HashingUtils.hash(logLevel, message, throwable == null ? null : throwable.getClass());
     }
 
     public StackTraceElement[] stackTrace() {
@@ -122,34 +132,49 @@ public StackTraceElement[] stackTrace() {
       return stackTrace;
     }
 
-    @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;
-
-      if (throwable == that.throwable) {
+    /**
+     * @param throwableStackTrace {@code throwable.getStackTrace()}, memoized once by the caller and
+     *     shared across every candidate scanned for a given {@code addLogMessage} call -- avoids
+     *     paying {@code getStackTrace()}'s defensive-copy allocation on each comparison. Only
+     *     non-null when {@code throwable} needs a deep comparison against some candidate.
+     */
+    private boolean matchesKey(
+        String logLevel,
+        String message,
+        @Nullable Throwable throwable,
+        @Nullable StackTraceElement[] throwableStackTrace) {
+      if (!Objects.equals(this.logLevel, logLevel)) return false;
+      if (!Objects.equals(this.message, message)) return false;
+
+      if (this.throwable == throwable) {
         // DQH - While this path may seem unlikely, it does happen if the JVM fast
         // throws optimization kicks-in (for NPE, etc), so this case is worth optimizing.
 
         // This also covers the case where both throwables are null
         return true;
-      } else if (throwable != null && that.throwable != null) {
+      } else if (this.throwable != null && throwable != null) {
         // Both have a throwable perform a deeper comparison
-        return throwable.getClass().equals(that.throwable.getClass())
-            && Objects.deepEquals(stackTrace(), that.stackTrace());
+        return this.throwable.getClass().equals(throwable.getClass())
+            && Objects.deepEquals(stackTrace(), throwableStackTrace);
       } else {
         // One has an exception & the other doesn't, not equal
         return false;
       }
     }
 
     @Override
-    public int hashCode() {
-      return HashingUtils.hash(logLevel, message, throwable == null ? null : throwable.getClass());
+    public boolean matches(@Nonnull RawLogMessage other) {
+      if (!Objects.equals(logLevel, other.logLevel)) return false;
+      if (!Objects.equals(message, other.message)) return false;
+
+      if (throwable == other.throwable) {
+        return true;
+      } else if (throwable != null && other.throwable != null) {
+        return throwable.getClass().equals(other.throwable.getClass())
+            && Objects.deepEquals(stackTrace(), other.stackTrace());
+      } else {
+        return false;
+      }
     }
   }
 }