ThreadSafeMapBenchmark.java

package datadog.trace.util;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.function.Supplier;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Threads;
import org.openjdk.jmh.annotations.Warmup;

/**
 * <ul>
 *   Benchmark comparing different approaches to filling and reading a Map in a multi-thread
 *   context.
 *   <li>ConcurrentMap - only when there are simultaneously readers & writers in multiple threads
 *   <li>HashMap via volatile - preferred for background thread updates
 *   <li>synchronized HashMap - when simultaneous readers & writers are uncommon (e.g. tags)
 *   <li>FlatHashtable - lock-free reads (no lock, no volatile; benign-race) of a fixed, once-built
 *       keyed set; a find-or-create table, not a general concurrent Map (no arbitrary put/remove)
 * </ul>
 *
 * <p>
 *
 * <p>In most situations in dd-java-agent, ConcurrentMaps are not necessarily needed and incur
 * additional overhead. ConcurrentMaps make sense when concurrent writers are likely.
 *
 * <p>If a Map can be created atomically in one thread and then stored into a volatile, that is the
 * preferred solution. For example, requesting an update from agent / API and then exposing to the
 * rest of the tracer via a global.
 *
 * <p>If a Map needs to be written in a thread-safe manner, but is primarily accessed from one
 * thread at a time, then a synchronized HashMap is usually the best option. <code>
 * MacBook M1 with 1 thread (Java 21)
 *
 * Benchmark                                            Mode  Cnt          Score          Error  Units
 * ThreadSafeMapBenchmark.create_concHashMap           thrpt    6    8081979.153 ±   261559.222  ops/s
 * ThreadSafeMapBenchmark.create_concSkipListMap       thrpt    6    2998832.124 ±   103708.038  ops/s
 * ThreadSafeMapBenchmark.create_hashMap               thrpt    6   24938311.610 ±   673725.902  ops/s
 * ThreadSafeMapBenchmark.create_hashMap_synchronized  thrpt    6    7971740.607 ±   121986.296  ops/s
 *
 * ThreadSafeMapBenchmark.get_concHashMap              thrpt    6  173942565.340 ± 12003493.448  ops/s
 * ThreadSafeMapBenchmark.get_concSkipListMap          thrpt    6   79230298.061 ± 13007895.765  ops/s
 * ThreadSafeMapBenchmark.get_hashMap_synchronized     thrpt    6   98056657.832 ±  3413815.061  ops/s
 * ThreadSafeMapBenchmark.get_hashMap_volatile         thrpt    6  210511753.596 ±  5017502.317  ops/s
 * </code> <code>
 * MacBook M1 with 8 threads (Java 21)
 *
 * Benchmark                                            Mode  Cnt          Score           Error  Units
 * ThreadSafeMapBenchmark.create_concHashMap           thrpt    6   58015351.219 ±   6201384.867  ops/s
 * ThreadSafeMapBenchmark.create_concSkipListMap       thrpt    6   19296105.790 ±   4516587.751  ops/s
 * ThreadSafeMapBenchmark.create_hashMap               thrpt    6  147917381.815 ±  22901897.589  ops/s
 * ThreadSafeMapBenchmark.create_hashMap_synchronized  thrpt    6   56466354.962 ±  13202034.783  ops/s
 *
 * ThreadSafeMapBenchmark.get_concHashMap              thrpt    6  849986442.797 ±  14499355.893  ops/s
 * ThreadSafeMapBenchmark.get_concSkipListMap          thrpt    6   26828246.629 ±   2772377.532  ops/s
 * ThreadSafeMapBenchmark.get_hashMap_synchronized     thrpt    6   20123419.604 ±   4858466.787  ops/s
 * ThreadSafeMapBenchmark.get_hashMap_volatile         thrpt    6  286024211.995 ± 114449056.603  ops/s
 * </code>
 */
@Fork(2)
@Warmup(iterations = 2)
@Measurement(iterations = 3)
@Threads(8)
@State(Scope.Thread)
public class ThreadSafeMapBenchmark {
  static final String[] INSERTION_KEYS = {
    "foo", "bar", "baz", "quux", "foobar", "foobaz", "key0", "key1", "key2", "key3"
  };

  static final String[] EQUAL_KEYS =
      init(
          () -> {
            String[] keys = new String[INSERTION_KEYS.length];
            for (int i = 0; i < INSERTION_KEYS.length; ++i) {
              keys[i] = new String(INSERTION_KEYS[i]);
            }
            return keys;
          });

  static <T> T init(Supplier<T> supplier) {
    return supplier.get();
  }

  // Per-thread (@State(Scope.Thread)) so cycling the lookup key doesn't contend a shared counter.
  // The maps below stay static/shared (the point — concurrent reads of one map); only the index is
  // per-thread. A shared counter's cache-line ping-pong would otherwise floor the fastest reads
  // (e.g. FlatHashtable's lock-free probe), hiding exactly the differences this benchmark compares.
  int lookupIndex = 0;

  String nextLookupKey() {
    return nextLookupKey(EQUAL_KEYS);
  }

  String nextLookupKey(String[] keys) {
    int localIndex = ++lookupIndex;
    if (localIndex >= keys.length) {
      lookupIndex = localIndex = 0;
    }
    return keys[localIndex];
  }

  static void fill(Map<String, Integer> map) {
    for (int i = 0; i < INSERTION_KEYS.length; ++i) {
      map.put(INSERTION_KEYS[i], i);
    }
  }

  // FlatHashtable's contribution here is the lock-free concurrent read: get() is a plain array
  // probe
  // with no lock and no volatile — safe under concurrency because the table is published once (a
  // final static field) and each entry's identity fields are final. (Fixture mirrors the one in
  // SingleThreadedMapBenchmark; the benchmarks are self-contained.)
  static final class IntEntry {
    final String key;
    final int value;

    IntEntry(String key, int value) {
      this.key = key;
      this.value = value;
    }
  }

  static final class IntEntryKeyStrategy extends FlatHashtable.EntryStrategy<IntEntry, String> {
    static final IntEntryKeyStrategy INSTANCE = new IntEntryKeyStrategy();

    private IntEntryKeyStrategy() {}

    @Override
    public boolean matches(IntEntry entry, String key) {
      return key.equals(entry.key);
    }

    @Override
    public long hashOf(IntEntry entry) {
      return entry.key.hashCode(); // consistent with the default hashKey
    }
  }

  // --- CHA-defeat decoys ---------------------------------------------------------------------
  // These are never used to build a table; they exist only to be *loaded* (see CHA_DEFEAT), so
  // MatchingStrategy.matches and .hashKey each have >=2 concrete implementors. That denies C2 the
  // single-implementor CHA devirtualization of matchStrat.hashKey/matches inside get(). If the
  // strategy calls still inline afterward, the win is structural (the constant INSTANCE's exact
  // type propagated through the inlined get), not a CHA bet that would deopt on a second subclass.

  // Second matches impl -> MatchingStrategy.matches is polymorphic.
  static final class DecoyMatchStrategy extends FlatHashtable.EntryStrategy<IntEntry, String> {
    static final DecoyMatchStrategy INSTANCE = new DecoyMatchStrategy();

    private DecoyMatchStrategy() {}

    @Override
    public boolean matches(IntEntry entry, String key) {
      return key == entry.key; // deliberately different body from IntEntryKeyStrategy
    }

    @Override
    public long hashOf(IntEntry entry) {
      return entry.key.hashCode();
    }
  }

  // Overrides hashKey -> MatchingStrategy.hashKey is polymorphic too (default + this override).
  static final class DecoyHashKeyStrategy extends FlatHashtable.EntryStrategy<IntEntry, String> {
    static final DecoyHashKeyStrategy INSTANCE = new DecoyHashKeyStrategy();

    private DecoyHashKeyStrategy() {}

    @Override
    public long hashKey(String key) {
      return key.length();
    }

    @Override
    public boolean matches(IntEntry entry, String key) {
      return key.equals(entry.key);
    }

    @Override
    public long hashOf(IntEntry entry) {
      return entry.key.length();
    }
  }

  // Referenced only so these three concrete implementors load at benchmark class-init, before the
  // hot method compiles — see the CHA-defeat note above.
  @SuppressWarnings("unused")
  static final Object[] CHA_DEFEAT = {
    IntEntryKeyStrategy.INSTANCE, DecoyMatchStrategy.INSTANCE, DecoyHashKeyStrategy.INSTANCE
  };

  static IntEntry[] _create_flat() {
    // Sized to the key count (FlatHashtable is fixed-capacity, no resize): load factor <= 0.5.
    IntEntry[] table = FlatHashtable.create(IntEntry.class, INSERTION_KEYS.length);
    for (int i = 0; i < INSERTION_KEYS.length; ++i) {
      FlatHashtable.insert(table, new IntEntry(INSERTION_KEYS[i], i), IntEntryKeyStrategy.INSTANCE);
    }
    return table;
  }

  static final HashMap<String, Integer> _create_hashMap() {
    HashMap<String, Integer> map = new HashMap<>();
    fill(map);
    return map;
  }

  @Benchmark
  public Map<String, Integer> create_hashMap() {
    return _create_hashMap();
  }

  static volatile HashMap<String, Integer> VOLATILE_HASH_MAP = _create_hashMap();

  @Benchmark
  public Integer get_hashMap_volatile() {
    Map<String, Integer> map = VOLATILE_HASH_MAP;
    return map.get(nextLookupKey());
  }

  static final Map<String, Integer> _create_hashMap_synchronized() {
    Map<String, Integer> map = Collections.synchronizedMap(new HashMap<>());
    fill(map);
    return map;
  }

  @Benchmark
  public Map<String, Integer> create_hashMap_synchronized() {
    return _create_hashMap_synchronized();
  }

  static final Map<String, Integer> SYNC_HASH_MAP = _create_hashMap_synchronized();

  @Benchmark
  public Integer get_hashMap_synchronized() {
    return SYNC_HASH_MAP.get(nextLookupKey());
  }

  static ConcurrentHashMap<String, Integer> _create_concHashMap() {
    ConcurrentHashMap<String, Integer> map = new ConcurrentHashMap<>();
    fill(map);
    return map;
  }

  @Benchmark
  public ConcurrentHashMap<String, Integer> create_concHashMap() {
    return _create_concHashMap();
  }

  static final ConcurrentHashMap<String, Integer> CONC_HASH_MAP = _create_concHashMap();

  @Benchmark
  public Integer get_concHashMap() {
    return CONC_HASH_MAP.get(nextLookupKey());
  }

  static ConcurrentSkipListMap<String, Integer> _create_concSkipListMap() {
    ConcurrentSkipListMap<String, Integer> map = new ConcurrentSkipListMap<>();
    fill(map);
    return map;
  }

  @Benchmark
  public ConcurrentSkipListMap<String, Integer> create_concSkipListMap() {
    return _create_concSkipListMap();
  }

  static final ConcurrentSkipListMap<String, Integer> CONC_SKIP_LIST_MAP =
      _create_concSkipListMap();

  @Benchmark
  public Integer get_concSkipListMap() {
    return CONC_SKIP_LIST_MAP.get(nextLookupKey());
  }

  @Benchmark
  public IntEntry[] create_flatHashtable() {
    return _create_flat();
  }

  static final IntEntry[] FLAT_TABLE = _create_flat();

  @Benchmark
  public IntEntry get_flatHashtable() {
    // Lock-free concurrent read of the shared, once-published table.
    return FlatHashtable.get(FLAT_TABLE, nextLookupKey(), IntEntryKeyStrategy.INSTANCE);
  }
}