diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java index 896342058a2..a1adf022ae3 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java @@ -1,5 +1,7 @@ package datadog.trace.bootstrap; +import datadog.instrument.fieldinject.GlobalObjectStore; + /** * {@link ContextStore} that attempts to store context in its keys by using bytecode-injected * fields. Delegates to a lazy {@link WeakMap} for keys that don't have a field for this store. @@ -16,7 +18,7 @@ public Object get(final Object key) { if (key instanceof FieldBackedContextAccessor) { return ((FieldBackedContextAccessor) key).$get$__datadogContext$(storeId); } else { - return weakStore().get(key); + return GlobalObjectStore.get(key, storeId); } } @@ -25,7 +27,7 @@ public void put(final Object key, final Object context) { if (key instanceof FieldBackedContextAccessor) { ((FieldBackedContextAccessor) key).$put$__datadogContext$(storeId, context); } else { - weakStore().put(key, context); + GlobalObjectStore.put(key, storeId, context); } } @@ -45,7 +47,7 @@ public Object putIfAbsent(final Object key, final Object context) { } return existingContext; } else { - return weakStore().putIfAbsent(key, context); + return GlobalObjectStore.getOrPut(key, storeId, context); } } @@ -71,7 +73,7 @@ public Object computeIfAbsent( } return existingContext; } else { - return weakStore().computeIfAbsent(key, contextFactory); + return GlobalObjectStore.getOrCompute(key, storeId, contextFactory::create); } } @@ -90,22 +92,7 @@ public Object remove(Object key) { } return existingContext; } else { - return weakStore().remove(key); - } - } - - // only create WeakMap-based fall-back when we need it - private volatile WeakMapContextStore weakStore; - private final Object synchronizationInstance = new Object(); - - WeakMapContextStore weakStore() { - if (null == weakStore) { - synchronized (synchronizationInstance) { - if (null == weakStore) { - weakStore = new WeakMapContextStore<>(); - } - } + return GlobalObjectStore.remove(key, storeId); } - return weakStore; } } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java index 529e475ab87..de3bbc68c3c 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java @@ -1,5 +1,6 @@ package datadog.trace.bootstrap; +import datadog.instrument.fieldinject.GlobalObjectStore; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.Arrays; import java.util.concurrent.ConcurrentHashMap; @@ -137,11 +138,11 @@ private static FieldBackedContextStore createStore(final int storeId) { /** Injection helper that immediately delegates to the weak-map for the given context store. */ public static Object weakGet(final Object key, final int storeId) { - return getContextStore(storeId).weakStore().get(key); + return GlobalObjectStore.get(key, storeId); } /** Injection helper that immediately delegates to the weak-map for the given context store. */ public static void weakPut(final Object key, final int storeId, final Object context) { - getContextStore(storeId).weakStore().put(key, context); + GlobalObjectStore.put(key, storeId, context); } } diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java index 3a8c7065362..49416afa03b 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java @@ -6,6 +6,7 @@ import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer; import datadog.environment.SystemProperties; +import datadog.instrument.fieldinject.GlobalObjectStore; import datadog.trace.agent.tooling.bytebuddy.SharedTypePools; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableRedefinitionStrategyListener; import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers; @@ -59,7 +60,8 @@ public class AgentInstaller { enableByteBuddyRawTypes(); disableByteBuddyNexus(); // register weak map supplier as early as possible - WeakMaps.registerAsSupplier(); + // WeakMaps.registerAsSupplier(); + AgentTaskScheduler.get().scheduleAtFixedRate(GlobalObjectStore::removeStaleEntries, 1, 1, TimeUnit.SECONDS); circularityErrorWorkaround(); } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java index 1392a08fbfa..25d179ac134 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java @@ -3,6 +3,7 @@ import static datadog.trace.bootstrap.FieldBackedContextStores.getContextStoreId; import static datadog.trace.util.Strings.getInternalName; +import datadog.instrument.fieldinject.GlobalObjectStore; import datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers; import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.Pair; @@ -35,8 +36,8 @@ public final class FieldBackedContextInjector implements AsmVisitorWrapper { private static final Logger log = LoggerFactory.getLogger(FieldBackedContextInjector.class); - static final String FIELD_BACKED_CONTEXT_STORES_CLASS = - getInternalName(FieldBackedContextStores.class.getName()); + static final String GLOBAL_OBJECT_STORE_CLASS = + getInternalName(GlobalObjectStore.class.getName()); static final String FIELD_BACKED_CONTEXT_ACCESSOR_CLASS = getInternalName(FieldBackedContextAccessor.class.getName()); @@ -51,12 +52,12 @@ public final class FieldBackedContextInjector implements AsmVisitorWrapper { static final String PUTTER_METHOD_DESCRIPTOR = Type.getMethodDescriptor(Type.VOID_TYPE, Type.INT_TYPE, Type.getType(Object.class)); - static final String WEAK_GET_METHOD = "weakGet"; + static final String WEAK_GET_METHOD = "get"; static final String WEAK_GET_METHOD_DESCRIPTOR = Type.getMethodDescriptor( Type.getType(Object.class), Type.getType(Object.class), Type.INT_TYPE); - static final String WEAK_PUT_METHOD = "weakPut"; + static final String WEAK_PUT_METHOD = "put"; static final String WEAK_PUT_METHOD_DESCRIPTOR = Type.getMethodDescriptor( Type.VOID_TYPE, Type.getType(Object.class), Type.INT_TYPE, Type.getType(Object.class)); @@ -445,7 +446,7 @@ private void invokeWeakGet(final MethodVisitor mv) { mv.visitIntInsn(Opcodes.ILOAD, 1); mv.visitMethodInsn( Opcodes.INVOKESTATIC, - FIELD_BACKED_CONTEXT_STORES_CLASS, + GLOBAL_OBJECT_STORE_CLASS, WEAK_GET_METHOD, WEAK_GET_METHOD_DESCRIPTOR, false); @@ -458,7 +459,7 @@ private void invokeWeakPut(final MethodVisitor mv) { mv.visitIntInsn(Opcodes.ALOAD, 2); mv.visitMethodInsn( Opcodes.INVOKESTATIC, - FIELD_BACKED_CONTEXT_STORES_CLASS, + GLOBAL_OBJECT_STORE_CLASS, WEAK_PUT_METHOD, WEAK_PUT_METHOD_DESCRIPTOR, false); diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index ca00719d2b1..83178dd416a 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -34,7 +34,7 @@ scala33 = "3.3.0" autoservice = "1.1.1" asm = "9.10.1" byte-buddy = "1.18.10" -instrument-java = "0.0.4" +instrument-java = "0.0.5-SNAPSHOT" # Benchmarks jmh = "1.37"