diff --git a/field-inject/src/jmh/java/datadog/instrument/fieldinject/ObjectStoreBenchmark.java b/field-inject/src/jmh/java/datadog/instrument/fieldinject/ObjectStoreBenchmark.java index 2b5a25e..2b30fa5 100644 --- a/field-inject/src/jmh/java/datadog/instrument/fieldinject/ObjectStoreBenchmark.java +++ b/field-inject/src/jmh/java/datadog/instrument/fieldinject/ObjectStoreBenchmark.java @@ -1,5 +1,7 @@ package datadog.instrument.fieldinject; +import static datadog.instrument.fieldinject.GlobalObjectStore.SHARD_COUNT; +import static datadog.instrument.fieldinject.GlobalObjectStore.SHARD_HARD_LIMIT; import static java.util.concurrent.TimeUnit.MICROSECONDS; import static java.util.concurrent.TimeUnit.SECONDS; import static java.util.stream.Collectors.toList; @@ -44,7 +46,7 @@ public class ObjectStoreBenchmark { @State(Scope.Benchmark) public static class TrialData { - @Param({"1000", "5000", "10000", "50000", "100000", "150000"}) + @Param({"1000", "5000", "10000", "50000", "100000", "200000", "300000"}) public int targetGlobalOccupancy; public List seedRequestIds; @@ -85,12 +87,14 @@ public void multiThreaded(Blackhole blackhole, ThreadData threadData) { batchedCompute(blackhole, threadData); } - private static final int NUM_STORES = 10; + // multiple of GlobalObjectStore's shard count, for even spread + private static final int NUM_STORES = SHARD_COUNT * 4; private static final ObjectStore[] stores = new ObjectStore[NUM_STORES]; static { - Arrays.setAll(stores, store -> ObjectStore.of("KeyType", "ValueType$" + store)); + // distinct key types per store spread store-ids evenly across shards + Arrays.setAll(stores, store -> ObjectStore.of("KeyType$" + store, "ValueType")); } private static final WeakObjectMap[] maps = new WeakObjectMap[NUM_STORES]; @@ -100,7 +104,7 @@ public void multiThreaded(Blackhole blackhole, ThreadData threadData) { } // create enough keys to cover max target occupancy per-store - private static final Object[] keys = new Object[150_000 / NUM_STORES]; + private static final Object[] keys = new Object[300_000 / NUM_STORES]; static { generateKeys(); @@ -174,7 +178,7 @@ private Object compute(int requestId) { */ static final class WeakObjectMap { // total capacity over all per-map stores should equal GlobalObjectStore's hard limit - private static final int MAX_SIZE = 100_000 / NUM_STORES; + private static final int MAX_SIZE = (SHARD_HARD_LIMIT * SHARD_COUNT) / NUM_STORES; private final WeakConcurrentMap map = new WeakConcurrentMap<>(false, true); diff --git a/field-inject/src/jmh/java/datadog/instrument/fieldinject/ObjectStoreContentionBenchmark.java b/field-inject/src/jmh/java/datadog/instrument/fieldinject/ObjectStoreContentionBenchmark.java new file mode 100644 index 0000000..7039f57 --- /dev/null +++ b/field-inject/src/jmh/java/datadog/instrument/fieldinject/ObjectStoreContentionBenchmark.java @@ -0,0 +1,234 @@ +/* + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). + * Copyright 2025-Present Datadog, Inc. + */ + +package datadog.instrument.fieldinject; + +import static java.util.concurrent.TimeUnit.SECONDS; + +import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap; +import java.util.concurrent.atomic.AtomicInteger; +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.Threads; +import org.openjdk.jmh.annotations.Warmup; + +/** + * Measures whether contention on unrelated {@link ObjectStore}s spills over onto a store nobody + * else is touching. A per-store map shares nothing, so it should show no difference between running + * alone and running alongside a contended pool. {@link ObjectStore} is instead backed by {@link + * GlobalObjectStore}, which picks a shard by fibonacci-hashing {@code storeId}; a contended store + * landing on the same shard as the isolated store (via a storeId hash collision, or a shared + * ageing/eviction pass) could degrade its throughput. + * + *

The contended pool uses a distinct key type per store; since the shard hash mixes both halves + * of {@code storeId}, varying either key type or value type spreads store-ids across shards. + * + *

Each variant runs the isolated store alone as a baseline ({@code isolatedStoreAlone_*}), then + * again alongside a disjoint contended pool hammered by more threads ({@code + * isolatedStoreUnderLoad_*} paired with {@code contendedPool_*} in a JMH {@code @Group}). Re-run + * this when tuning {@code GlobalObjectStore}'s shard count/selection or backing map capacity: less + * spillover should shrink the gap between {@code isolatedStoreAlone_globalObjectStore} and {@code + * isolatedStoreUnderLoad_globalObjectStore}. + * + * 

+ *   ./gradlew :field-inject:jmh -Pjmh.includes=ObjectStoreContentionBenchmark
+ *
+ */ +@BenchmarkMode(Mode.Throughput) +@OutputTimeUnit(SECONDS) +@Fork(3) +@Warmup(iterations = 2) +@Measurement(iterations = 5) +@State(Scope.Benchmark) +@SuppressWarnings("unused") +public class ObjectStoreContentionBenchmark { + + private static final int CONTENDED_STORE_COUNT = 4; + private static final int CONTENDED_GROUP_THREADS = 7; + private static final int ISOLATED_GROUP_THREADS = 1; + + // Well below GlobalObjectStore's shard limit so this measures hashing/locking, not eviction. + private static final int CONTENDED_KEY_BUDGET = 6400; + private static final int ISOLATED_KEY_POOL_SIZE = 4096; + + private static final Object CONTEXT = new Object(); + + private ObjectStore[] contendedGlobalStores; + private ObjectStore isolatedGlobalStore; + + private PerStoreMap[] contendedPerStoreMaps; + private PerStoreMap isolatedPerStoreMap; + + private final AtomicInteger nextContendedThreadIndex = new AtomicInteger(); + + @Setup(Level.Trial) + @SuppressWarnings("unchecked") + public void setupTrial() { + contendedGlobalStores = new ObjectStore[CONTENDED_STORE_COUNT]; + contendedPerStoreMaps = new PerStoreMap[CONTENDED_STORE_COUNT]; + for (int i = 0; i < CONTENDED_STORE_COUNT; i++) { + // Distinct key type per store, mirroring real shard spread (see class javadoc). + contendedGlobalStores[i] = ObjectStore.of("ContendedKey" + i, "ContendedValue" + i); + contendedPerStoreMaps[i] = new PerStoreMap(); + } + isolatedGlobalStore = ObjectStore.of("IsolatedKey", "IsolatedValue"); + isolatedPerStoreMap = new PerStoreMap(); + } + + private static KeyCursor newKeyCursor(int size) { + Object[] keys = new Object[size]; + for (int i = 0; i < size; i++) { + keys[i] = new Object(); + } + return new KeyCursor(keys); + } + + private static final class KeyCursor { + private final Object[] keys; + private int cursor; + + KeyCursor(Object[] keys) { + this.keys = keys; + } + + Object next() { + cursor = (cursor + 1) % keys.length; + return keys[cursor]; + } + } + + @State(Scope.Thread) + public static class ContendedGlobalState { + private ObjectStore store; + private KeyCursor keyCursor; + + @Setup(Level.Trial) + public void setup(ObjectStoreContentionBenchmark benchmark) { + int threadIndex = benchmark.nextContendedThreadIndex.getAndIncrement(); + store = benchmark.contendedGlobalStores[threadIndex % CONTENDED_STORE_COUNT]; + keyCursor = newKeyCursor(Math.max(1, CONTENDED_KEY_BUDGET / CONTENDED_GROUP_THREADS)); + } + } + + @State(Scope.Thread) + public static class ContendedPerStoreState { + private PerStoreMap store; + private KeyCursor keyCursor; + + @Setup(Level.Trial) + public void setup(ObjectStoreContentionBenchmark benchmark) { + int threadIndex = benchmark.nextContendedThreadIndex.getAndIncrement(); + store = benchmark.contendedPerStoreMaps[threadIndex % CONTENDED_STORE_COUNT]; + keyCursor = newKeyCursor(Math.max(1, CONTENDED_KEY_BUDGET / CONTENDED_GROUP_THREADS)); + } + } + + @State(Scope.Thread) + public static class IsolatedGlobalState { + private KeyCursor keyCursor; + + @Setup(Level.Trial) + public void setup() { + keyCursor = newKeyCursor(ISOLATED_KEY_POOL_SIZE); + } + } + + @State(Scope.Thread) + public static class IsolatedPerStoreState { + private KeyCursor keyCursor; + + @Setup(Level.Trial) + public void setup() { + keyCursor = newKeyCursor(ISOLATED_KEY_POOL_SIZE); + } + } + + private static void putGetRemove(ObjectStore store, KeyCursor keyCursor) { + Object key = keyCursor.next(); + store.put(key, CONTEXT); + store.get(key); + store.remove(key); + } + + private static void putGetRemove(PerStoreMap store, KeyCursor keyCursor) { + Object key = keyCursor.next(); + store.put(key, CONTEXT); + store.get(key); + store.remove(key); + } + + // --- Baseline: the isolated store with zero interference from any other thread. --- + + @Benchmark + @Threads(ISOLATED_GROUP_THREADS) + public void isolatedStoreAlone_globalObjectStore(IsolatedGlobalState state) { + putGetRemove(isolatedGlobalStore, state.keyCursor); + } + + @Benchmark + @Threads(ISOLATED_GROUP_THREADS) + public void isolatedStoreAlone_mapPerStore(IsolatedPerStoreState state) { + putGetRemove(isolatedPerStoreMap, state.keyCursor); + } + + // --- Mixed: the isolated store, plus a disjoint contended pool hammered concurrently. --- + // ISOLATED_GROUP_THREADS + CONTENDED_GROUP_THREADS must sum to the benchmark's total threads. + + @Benchmark + @Group("mixedGlobalObjectStore") + @GroupThreads(ISOLATED_GROUP_THREADS) + public void isolatedStoreUnderLoad_globalObjectStore(IsolatedGlobalState state) { + putGetRemove(isolatedGlobalStore, state.keyCursor); + } + + @Benchmark + @Group("mixedGlobalObjectStore") + @GroupThreads(CONTENDED_GROUP_THREADS) + public void contendedPool_globalObjectStore(ContendedGlobalState state) { + putGetRemove(state.store, state.keyCursor); + } + + @Benchmark + @Group("mixedMapPerStore") + @GroupThreads(ISOLATED_GROUP_THREADS) + public void isolatedStoreUnderLoad_mapPerStore(IsolatedPerStoreState state) { + putGetRemove(isolatedPerStoreMap, state.keyCursor); + } + + @Benchmark + @Group("mixedMapPerStore") + @GroupThreads(CONTENDED_GROUP_THREADS) + public void contendedPool_mapPerStore(ContendedPerStoreState state) { + putGetRemove(state.store, state.keyCursor); + } + + /** Genuine per-store map: nothing here is ever shared with any other store. */ + private static final class PerStoreMap { + private final WeakConcurrentMap map = new WeakConcurrentMap<>(false, true); + + Object get(Object key) { + return map.get(key); + } + + void put(Object key, Object value) { + map.put(key, value); + } + + Object remove(Object key) { + return map.remove(key); + } + } +} diff --git a/field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java b/field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java index 4b53b3b..2737119 100644 --- a/field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java +++ b/field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java @@ -13,46 +13,64 @@ import java.util.Iterator; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.atomic.AtomicIntegerArray; +import java.util.concurrent.atomic.AtomicReferenceArray; import java.util.function.Function; import java.util.function.Supplier; import javax.annotation.Nullable; /** - * Global key-value store used when field-injection is not possible. Since the same object may - * participate in multiple stores each global key captures the store identity along with a weak + * Global sharded key-value store used when field-injection is not possible. Since the same object + * may participate in multiple stores each global key captures the store identity along with a weak * reference to the owning key object. * - *

The store is split into two maps with separate reference queues: young and old. Ageing the - * store by one generation creates a new young map; the previous young map becomes the old map. + *

Each shard is split into two maps with separate reference queues: young and old. Ageing a + * shard by one generation creates a new young map; the previous young map becomes the old map. */ public final class GlobalObjectStore { - /** Target ceiling for the total number of objects in the global store, young and old. */ - private static final int GLOBAL_HARD_LIMIT = 100_000; + /** Target ceiling for the total number of objects in a shard, young and old. */ + static final int SHARD_HARD_LIMIT = 32_000; - /** Threshold at which we age the current store by one generation. */ - private static final int AGEING_THRESHOLD = GLOBAL_HARD_LIMIT / 2; + /** Threshold at which we age a shard by one generation. */ + static final int AGEING_THRESHOLD = SHARD_HARD_LIMIT / 2; - /** Target ceiling for the total number of objects allowed after background eviction. */ - private static final int GLOBAL_SOFT_LIMIT = (GLOBAL_HARD_LIMIT + AGEING_THRESHOLD) / 2; + /** Target ceiling for total number of objects allowed in a shard after background eviction. */ + static final int SHARD_SOFT_LIMIT = (SHARD_HARD_LIMIT + AGEING_THRESHOLD) / 2; + + /** Shift used to pick a shard from a store-id after fibonacci-hashing. */ + private static final int SHARD_BITS = 3; + + /** Number of independent shards; store-ids are spread across shards. */ + static final int SHARD_COUNT = 1 << SHARD_BITS; /** Threshold at which we start doing limited cleanup at the same time as put operations. */ - private static final int INLINE_CLEANUP_THRESHOLD = 5_000; + private static final int INLINE_CLEANUP_THRESHOLD = 2_000; - /** Sample underlying map size periodically, driven by misses when writing. */ - private static final int MAP_SIZE_SAMPLE_RATE = (1 << 10) - 1; // sample every 1k misses + /** Randomly sample underlying map size, approximately once every 1024 writes per-thread. */ + private static final int SIZE_SAMPLE_RATE = 1 << 10; - /** Constant supplier used when nothing in the store is considered old. */ + /** Constant supplier used when nothing in a shard is considered old. */ private static final Supplier NO_OLD_STALE_KEYS = () -> null; - /** The current generation of the global object store. */ - private static volatile GlobalObjectStore store = new GlobalObjectStore(); + /** The current generation of each shard of the global object store. */ + private static final AtomicReferenceArray shards = initShards(); + + /** Per-shard token used to decide which thread gets to age that shard. */ + private static final AtomicIntegerArray ageing = new AtomicIntegerArray(SHARD_COUNT); + + private static AtomicReferenceArray initShards() { + AtomicReferenceArray shards = new AtomicReferenceArray<>(SHARD_COUNT); + for (int shardIndex = 0; shardIndex < SHARD_COUNT; shardIndex++) { + shards.set(shardIndex, new GlobalObjectStore(shardIndex)); + } + return shards; + } - /** Token used to decide which thread gets to age the store. */ - private static final AtomicBoolean ageing = new AtomicBoolean(); + // the following fields represent a generation of each shard in the global object store - // the following fields represent a generation of the object store + private final int shardIndex; /** Supplies store keys where the key object is unused and eligible for collection. */ private final ReferenceQueue staleKeys = new ReferenceQueue<>(); @@ -68,13 +86,15 @@ public final class GlobalObjectStore { private transient volatile int sampledYoungSize; - private GlobalObjectStore() { + private GlobalObjectStore(int shardIndex) { + this.shardIndex = shardIndex; this.map = new ConcurrentHashMap<>(); this.oldStaleKeys = NO_OLD_STALE_KEYS; this.oldMap = Collections.emptyMap(); } private GlobalObjectStore(GlobalObjectStore oldStore) { + this.shardIndex = oldStore.shardIndex; this.map = new ConcurrentHashMap<>(INLINE_CLEANUP_THRESHOLD); this.oldStaleKeys = oldStore.staleKeys::poll; this.oldMap = oldStore.map; @@ -89,7 +109,11 @@ private GlobalObjectStore(GlobalObjectStore oldStore) { * @return the estimated remaining size of the global object-store */ public static int removeStaleEntries() { - return store.doRemoveStaleEntries(); + int estimatedSize = 0; + for (int shardIndex = 0; shardIndex < SHARD_COUNT; shardIndex++) { + estimatedSize += shards.get(shardIndex).doRemoveStaleEntries(); + } + return estimatedSize; } private int doRemoveStaleEntries() { @@ -115,7 +139,7 @@ private int doRemoveStaleEntries() { // randomly evict old content to keep us below the soft limit Iterator itr = oldMap.keySet().iterator(); - while (estimatedTotal >= GLOBAL_SOFT_LIMIT && itr.hasNext()) { + while (estimatedTotal >= SHARD_SOFT_LIMIT && itr.hasNext()) { itr.next(); itr.remove(); estimatedTotal--; @@ -135,7 +159,7 @@ private int doRemoveStaleEntries() { public static Object get(Object key, int storeId) { LookupKey lookupKey = LookupKey.with(key, storeId); try { - return store.doGet(lookupKey); + return shard(storeId).doGet(lookupKey); } finally { lookupKey.reset(); } @@ -158,7 +182,7 @@ private Object doGet(LookupKey lookupKey) { */ public static void put(Object key, int storeId, @Nullable Object value) { if (value != null) { - store.checkCapacity(LookupKey.skip()).doPut(key, storeId, value); + shard(storeId).checkCapacity().doPut(key, storeId, value); } else { remove(key, storeId); } @@ -180,12 +204,12 @@ private void doPut(Object key, int storeId, Object value) { public static Object getOrPut(Object key, int storeId, @Nullable Object value) { LookupKey lookupKey = LookupKey.with(key, storeId); try { - GlobalObjectStore s = store; + GlobalObjectStore s = shard(storeId); Object existing = s.doGet(lookupKey); // avoids creating unnecessary store key if (existing != null || value == null) { return existing; } else { - return s.checkCapacity(lookupKey.miss()).doGetOrPut(key, storeId, value); + return s.checkCapacity().doGetOrPut(key, storeId, value); } } finally { lookupKey.reset(); @@ -210,12 +234,12 @@ private Object doGetOrPut(Object key, int storeId, Object value) { public static Object getOrCompute(Object key, int storeId, Function valueFunction) { LookupKey lookupKey = LookupKey.with(key, storeId); try { - GlobalObjectStore s = store; + GlobalObjectStore s = shard(storeId); Object existing = s.doGet(lookupKey); // avoids creating unnecessary store key if (existing != null) { return existing; } else { - return s.checkCapacity(lookupKey.miss()).doGetOrCompute(key, storeId, valueFunction); + return s.checkCapacity().doGetOrCompute(key, storeId, valueFunction); } } finally { lookupKey.reset(); @@ -239,7 +263,7 @@ private Object doGetOrCompute(Object key, int storeId, Function valueFunction) { public static Object remove(Object key, int storeId) { LookupKey lookupKey = LookupKey.with(key, storeId); try { - return store.doRemove(lookupKey); + return shard(storeId).doRemove(lookupKey); } finally { lookupKey.reset(); } @@ -255,15 +279,16 @@ private Object doRemove(LookupKey lookupKey) { } /** - * Checks store capacity, performing inline eviction or ageing if appropriate. + * Checks shard capacity, performing inline eviction or ageing if appropriate. * - * @param misses lookup misses on this thread when writing - * @return the latest generation of the global store + * @return the latest generation of the shard */ - private GlobalObjectStore checkCapacity(int misses) { - int youngSize = sampledYoungSize; - if ((misses & MAP_SIZE_SAMPLE_RATE) == 1) { // sample on first miss and every RATE after - youngSize = sampledYoungSize = map.size(); + private GlobalObjectStore checkCapacity() { + int youngSize; + if (ThreadLocalRandom.current().nextInt(SIZE_SAMPLE_RATE) == 0) { + sampledYoungSize = youngSize = map.size(); + } else { + youngSize = sampledYoungSize; } if (youngSize >= INLINE_CLEANUP_THRESHOLD) { Object staleKey = staleKeys.poll(); @@ -279,27 +304,32 @@ private GlobalObjectStore checkCapacity(int misses) { } /** - * Attempts to age this store by one generation; if already ageing don't block, use latest. + * Attempts to age this shard by one generation; if already ageing don't block, use latest. * - * @return the latest generation of the global store + * @return the latest generation of the shard */ - @SuppressFBWarnings("ST") // we want to update the global object store private GlobalObjectStore maybeAgeStore() { - // first try to get the token that allows us to age the global store - boolean attemptAgeing = ageing.compareAndSet(false, true); - // only after this get the latest generation of the store - GlobalObjectStore s = store; + // first try to get the token that allows us to age this shard + boolean attemptAgeing = ageing.compareAndSet(shardIndex, 0, 1); + // only after this get the latest generation of the shard + GlobalObjectStore s = shards.get(shardIndex); if (attemptAgeing) { try { if (s == this) { - // our store is still the latest; go ahead and age it - s = store = new GlobalObjectStore(this); + // our shard generation is still the latest; go ahead and age it + shards.set(shardIndex, s = new GlobalObjectStore(this)); } } finally { - ageing.set(false); // relinquish the token + ageing.set(shardIndex, 0); // relinquish the token } } - return s; // always return the latest generation of the store + return s; // always return the latest generation of the shard + } + + /** Returns the shard for the given store-id. */ + static GlobalObjectStore shard(int storeId) { + // use fibonacci-hashing to spread store-ids evenly, then take top bits + return shards.get((storeId * 0x9E3779B9) >>> (32 - SHARD_BITS)); } /** Key used to weakly associate a non-injected key and store-id with a value. */ @@ -345,9 +375,6 @@ private static final class LookupKey { int hash; int storeId; - /** Number of times a lookup missed when writing. */ - int misses; - /** * Returns a temporary lookup key for the current thread with the given object key and store-id. * This key must be reset by calling {@link #reset} as soon as the get/remove request completes. @@ -364,16 +391,6 @@ static LookupKey with(Object key, int storeId) { return thiz; } - /** Record the lookup was skipped when writing. */ - static int skip() { - return LOOKUP_KEY_CACHE.get().miss(); - } - - /** Record the lookup missed when writing. */ - int miss() { - return ++misses; - } - /** Resets this temporary lookup key so it can be reused in a future get/remove request. */ void reset() { this.key = null; // only need to clear the object key so it can be collected diff --git a/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreShardingTest.java b/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreShardingTest.java new file mode 100644 index 0000000..1ffdb27 --- /dev/null +++ b/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreShardingTest.java @@ -0,0 +1,76 @@ +package datadog.instrument.fieldinject; + +import static datadog.instrument.fieldinject.GlobalObjectStore.SHARD_COUNT; +import static datadog.instrument.fieldinject.ObjectStoreIds.objectStoreId; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.IdentityHashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; + +/** + * Checks that store-ids spread evenly across shards for common usage patterns. + * + *

Store-ids are built from key/value type-ids handed out sequentially, so a shard formula that + * just masks the low bits can badly cluster or even collapse onto a single shard; these tests guard + * against that regressing. Type names are unique per test to avoid interference from other tests + * sharing the same type-id cache. + */ +class ObjectStoreShardingTest { + + private static void assertSpreadAcrossShards(int[] storeIds) { + Map countsByShard = new IdentityHashMap<>(); + for (int storeId : storeIds) { + GlobalObjectStore shard = GlobalObjectStore.shard(storeId); + countsByShard.merge(shard, 1, Integer::sum); + } + + int used = countsByShard.size(); + int max = countsByShard.values().stream().mapToInt(Integer::intValue).max().orElse(0); + + assertTrue( + used >= SHARD_COUNT / 2, + "expected at least half the shards to be used, got " + used + "/" + SHARD_COUNT); + + int average = storeIds.length / SHARD_COUNT; + assertTrue( + max <= average * 3, + "expected no shard to receive far more than its share, got max=" + + max + + " for average=" + + average); + } + + @Test + void distinctKeyAndValueTypesPerStoreSpreadAcrossShards() { + int n = 64; + int[] storeIds = new int[n]; + for (int i = 0; i < n; i++) { + storeIds[i] = objectStoreId("DistinctKey" + i, "DistinctValue" + i); + } + assertSpreadAcrossShards(storeIds); + } + + @Test + void sharedKeyTypeWithManyValueTypesSpreadsAcrossShards() { + int n = 64; + int[] storeIds = new int[n]; + for (int i = 0; i < n; i++) { + storeIds[i] = objectStoreId("SharedKeyType", "InjectedFieldType" + i); + } + assertSpreadAcrossShards(storeIds); + } + + @Test + void smallKeyTypePoolWithGrowingValueTypePoolSpreadsAcrossShards() { + String[] keyTypes = { + "PoolSpan", "PoolAgentSpan", "PoolHttpRequest", "PoolHttpResponse", "PoolDBStatement" + }; + int n = 60; + int[] storeIds = new int[n]; + for (int i = 0; i < n; i++) { + storeIds[i] = objectStoreId(keyTypes[i % keyTypes.length], "PoolInjectedField" + i); + } + assertSpreadAcrossShards(storeIds); + } +} diff --git a/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java b/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java index a3982c4..3a632df 100644 --- a/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java +++ b/field-inject/src/test/java/datadog/instrument/fieldinject/ObjectStoreTest.java @@ -1,5 +1,10 @@ package datadog.instrument.fieldinject; +import static datadog.instrument.fieldinject.GlobalObjectStore.AGEING_THRESHOLD; +import static datadog.instrument.fieldinject.GlobalObjectStore.SHARD_COUNT; +import static datadog.instrument.fieldinject.GlobalObjectStore.SHARD_HARD_LIMIT; +import static datadog.instrument.fieldinject.GlobalObjectStore.SHARD_SOFT_LIMIT; +import static datadog.instrument.fieldinject.ObjectStoreIds.objectStoreId; import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; @@ -221,21 +226,19 @@ void removeStaleEntriesIsIdempotent() { // --- generational capacity / eviction --- - // Mirrors the private thresholds in GlobalObjectStore; kept here so the test intent is clear - // without exposing internals. If those thresholds change this test should be revisited. - private static final int AGEING_THRESHOLD = 50_000; - private static final int GLOBAL_SOFT_LIMIT = 75_000; - private static final int GLOBAL_HARD_LIMIT = 100_000; - @Test void sustainedInsertionIsBoundedByAgeingAndSoftLimitTrim() { ObjectStore capStore = ObjectStore.of("test.Capacity.Key", "test.Capacity.Value"); + // removeStaleEntries() sums estimated size across all shards; only this store's shard is + // populated here, so the sum reflects that one shard directly. + int otherShardsBaseline = ObjectStore.removeStaleEntries(); + // Insert enough distinct, strongly-referenced keys to drive the young generation past // AGEING_THRESHOLD several times over, landing mid-cycle (comfortably above the soft // limit) so both inline ageing and removeStaleEntries' soft-limit trim get exercised. - int totalInserts = (AGEING_THRESHOLD * 4) + 40_000; + int totalInserts = (AGEING_THRESHOLD * 4) + 8_000; List keys = new ArrayList<>(totalInserts); for (int i = 0; i < totalInserts; i++) { Object key = new Object(); @@ -243,19 +246,156 @@ void sustainedInsertionIsBoundedByAgeingAndSoftLimitTrim() { capStore.put(key, i); } - // Inline enforceCapacity keeps young+old from ever exceeding the hard limit by ageing + // Inline enforceCapacity keeps young+old from ever exceeding the shard hard limit by ageing // young into old before that point is reached, so recently inserted keys must still be - // retrievable even after hundreds of thousands of insertions. + // retrievable even after many multiples of the shard's capacity have been inserted. Object lastKey = keys.get(keys.size() - 1); assertEquals(totalInserts - 1, capStore.get(lastKey)); - int finalSize = ObjectStore.removeStaleEntries(); + int finalSize = ObjectStore.removeStaleEntries() - otherShardsBaseline; assertTrue( - finalSize < GLOBAL_HARD_LIMIT, - "Sustained insertion should have triggered eviction rather than unbounded growth"); + finalSize < SHARD_HARD_LIMIT, + "Sustained insertion should have triggered eviction rather than unbounded growth, observed " + + finalSize); assertTrue( - finalSize <= GLOBAL_SOFT_LIMIT, - "removeStaleEntries should trim content back to the soft limit, observed " + finalSize); + finalSize <= SHARD_SOFT_LIMIT, + "removeStaleEntries should trim content back to the shard soft limit, observed " + + finalSize); + } + + @Test + void survivingKeyIsShadowedAfterAgeingAndRemoveClearsBothGenerations() { + ObjectStore genStore = + ObjectStore.of("test.GenerationBoundary.Key", "test.GenerationBoundary.Value"); + + Object survivorKey = new Object(); + genStore.put(survivorKey, -1); + + // Insert enough further distinct, strongly-referenced keys to push the young generation past + // AGEING_THRESHOLD (ageing survivorKey's entry into the old generation), while staying well + // short of a second ageing cycle that would drop it again. + int keysToForceAgeing = AGEING_THRESHOLD + 10_000; + List keys = new ArrayList<>(keysToForceAgeing); + for (int i = 0; i < keysToForceAgeing; i++) { + Object key = new Object(); + keys.add(key); + genStore.put(key, i); + } + + Object lastBulkKey = keys.get(keys.size() - 1); + assertEquals(keysToForceAgeing - 1, genStore.get(lastBulkKey)); + + // survivorKey's original entry should now live in the old generation, but must still resolve. + assertEquals(-1, genStore.get(survivorKey)); + + // Overwriting writes into the (new) young generation, shadowing the stale old-generation entry. + genStore.put(survivorKey, 42); + assertEquals(42, genStore.get(survivorKey), "put should shadow the stale old-generation entry"); + + // remove() must clear both generations; otherwise the old entry would resurrect the old value. + assertEquals(42, genStore.remove(survivorKey)); + assertNull( + genStore.get(survivorKey), + "remove must clear the old-generation copy too, or it would resurrect the shadowed value"); + } + + @Test + void singleThreadCyclingAcrossAllShardsStillSamplesEachShard() { + // A thread writing to several stores in a fixed round-robin cycle whose length divides + // SIZE_SAMPLE_RATE (1024) should sample capacity on every shard it visits, not just one. + GlobalObjectStore[] chosenShards = new GlobalObjectStore[SHARD_COUNT]; + String[] chosenKeyTypes = new String[SHARD_COUNT]; + int found = 0; + int suffix = 0; + while (found < SHARD_COUNT) { + String candidate = "test.ShardCycle.Key" + suffix++; + GlobalObjectStore shard = + GlobalObjectStore.shard(objectStoreId(candidate, "test.ShardCycle.Value")); + boolean alreadyChosen = false; + for (int i = 0; i < found; i++) { + if (chosenShards[i] == shard) { + alreadyChosen = true; + break; + } + } + if (!alreadyChosen) { + chosenShards[found] = shard; + chosenKeyTypes[found] = candidate; + found++; + } + } + + ObjectStore[] cycleStores = new ObjectStore[SHARD_COUNT]; + for (int i = 0; i < SHARD_COUNT; i++) { + cycleStores[i] = ObjectStore.of(chosenKeyTypes[i], "test.ShardCycle.Value"); + } + + int baseline = ObjectStore.removeStaleEntries(); + + int perStoreInserts = SHARD_HARD_LIMIT + AGEING_THRESHOLD; + int totalInserts = perStoreInserts * SHARD_COUNT; + List keys = new ArrayList<>(totalInserts); + for (int i = 0; i < totalInserts; i++) { + Object key = new Object(); + keys.add(key); + cycleStores[i % SHARD_COUNT].put(key, i); + } + + Object lastKey = keys.get(keys.size() - 1); + assertEquals(totalInserts - 1, cycleStores[(totalInserts - 1) % SHARD_COUNT].get(lastKey)); + + // If any shard were starved of sampling it would never age or trim, so its young map would + // grow to hold roughly its entire share of inserts (perStoreInserts) instead of settling near + // the soft limit; that dwarfs a healthy aggregate across all shards. + int aggregate = ObjectStore.removeStaleEntries() - baseline; + int healthyCeiling = SHARD_COUNT * SHARD_SOFT_LIMIT; + assertTrue( + aggregate <= healthyCeiling, + "Every shard touched by this cyclic pattern should sample and age independently; " + + "observed aggregate size " + + aggregate + + " exceeds " + + healthyCeiling + + " (a starved shard would grow unbounded instead of ageing)"); + } + + @Test + void oneShardsSustainedLoadDoesNotAgeOrEvictAnotherShard() { + // Pick two key types landing on different shards, matching ObjectStoreShardingTest's approach. + String loadedKeyType = "test.ShardIsolation.LoadedKey"; + int loadedStoreId = objectStoreId(loadedKeyType, "test.ShardIsolation.Value"); + String quietKeyType = null; + int suffix = 0; + while (true) { + String candidate = "test.ShardIsolation.QuietKey" + suffix; + if (GlobalObjectStore.shard(objectStoreId(candidate, "test.ShardIsolation.Value")) + != GlobalObjectStore.shard(loadedStoreId)) { + quietKeyType = candidate; + break; + } + suffix++; + } + + ObjectStore quietStore = + ObjectStore.of(quietKeyType, "test.ShardIsolation.Value"); + Object quietKey = new Object(); + quietStore.put(quietKey, "still here"); + + ObjectStore loadedStore = + ObjectStore.of(loadedKeyType, "test.ShardIsolation.Value"); + int totalInserts = (AGEING_THRESHOLD * 4) + 8_000; + List keys = new ArrayList<>(totalInserts); + for (int i = 0; i < totalInserts; i++) { + Object key = new Object(); + keys.add(key); + loadedStore.put(key, i); + } + + Object lastKey = keys.get(keys.size() - 1); + assertEquals(totalInserts - 1, loadedStore.get(lastKey)); + + // The loaded shard aged/evicted repeatedly, but the quiet shard's entry must be untouched. + assertEquals("still here", quietStore.get(quietKey)); } // --- concurrency --- @@ -306,6 +446,58 @@ void concurrentPutsOnDistinctKeysAreAllVisible() throws InterruptedException { } } + @Test + void concurrentSustainedInsertionAcrossAgeingIsRaceFree() throws InterruptedException { + // maybeAgeStore() uses a per-shard CAS token so only one thread ages a shard at a time; this + // drives several threads through repeated ageing on the same shard to check that races there + // don't lose recently written entries. + ObjectStore concCapStore = + ObjectStore.of("test.ConcurrentCapacity.Key", "test.ConcurrentCapacity.Value"); + + int threads = 8; + int perThread = (AGEING_THRESHOLD * 4 + 8_000) / threads; + + // Tracks whichever write actually finishes last. A thread's own last write is not a safe + // thing to check here: other threads may still have thousands of writes left, which can + // legitimately age it out by design. + AtomicReference lastWrite = new AtomicReference<>(); + + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + ExecutorService executor = Executors.newFixedThreadPool(threads); + try { + for (int t = 0; t < threads; t++) { + final int threadIdx = t; + executor.submit( + () -> { + try { + start.await(); + for (int i = 0; i < perThread; i++) { + Object key = new Object(); + int value = (threadIdx * perThread) + i; + concCapStore.put(key, value); + lastWrite.set(new Object[] {key, value}); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + done.countDown(); + } + }); + } + start.countDown(); + assertThat(done.await(30, TimeUnit.SECONDS)).isTrue(); + } finally { + executor.shutdown(); + } + + // The write that actually finished last has nothing written after it: if a concurrent ageing + // race ever let two threads swap generations at once, it could still land in a generation + // that gets discarded instead of becoming the new old generation. + Object[] lastKeyValuePair = lastWrite.get(); + assertEquals(lastKeyValuePair[1], concCapStore.get(lastKeyValuePair[0])); + } + // --- helpers --- /**