From 65a61a8fa9e2245577e1db153e9c45089ee60ee3 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Fri, 9 Aug 2024 01:00:05 +0100 Subject: [PATCH] Avoid collected store ids from leaking to later field-injections when an instrumentation fails Co-authored-by: luneo7 --- .../context/FieldBackedContextInjector.java | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) 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 45e717ad49b..276ac6daf1b 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 @@ -5,6 +5,7 @@ import datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers; import datadog.trace.api.InstrumenterConfig; +import datadog.trace.api.Pair; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.FieldBackedContextAccessor; import datadog.trace.bootstrap.FieldBackedContextStores; @@ -70,7 +71,7 @@ public final class FieldBackedContextInjector implements AsmVisitorWrapper { Type.getType(FieldBackedContextAccessor.class); /** Keeps track of injection requests for the class being transformed by the current thread. */ - static final ThreadLocal INJECTED_STORE_IDS = new ThreadLocal<>(); + static final ThreadLocal> INJECTED_STORE_IDS = new ThreadLocal<>(); final boolean serialVersionUIDFieldInjection = InstrumenterConfig.get().isSerialVersionUIDFieldInjection(); @@ -129,7 +130,7 @@ public void visit( // keep track of all injection requests for the class currently being transformed // because we need to switch between them in the generated getter/putter methods - int storeId = injectContextStore(keyClassName, contextClassName); + int storeId = injectContextStore(name, keyClassName, contextClassName); storeFieldName = CONTEXT_STORE_ACCESS_PREFIX + storeId; if (interfaces == null) { @@ -484,26 +485,30 @@ private void invokeSuperPut(final MethodVisitor mv, final String superName) { } /** Requests injection of a context store for a key and context. */ - static int injectContextStore(final String keyClassName, final String contextClassName) { + static int injectContextStore( + final String target, final String keyClassName, final String contextClassName) { int storeId = getContextStoreId(keyClassName, contextClassName); - BitSet injectedStoreIds = INJECTED_STORE_IDS.get(); - if (null == injectedStoreIds) { - injectedStoreIds = new BitSet(); + // collect a new set of store ids every time we see a new target + Pair injectedStoreIds = INJECTED_STORE_IDS.get(); + if (null == injectedStoreIds || !target.equals(injectedStoreIds.getLeft())) { + injectedStoreIds = Pair.of(target, new BitSet()); INJECTED_STORE_IDS.set(injectedStoreIds); } - injectedStoreIds.set(storeId); + injectedStoreIds.getRight().set(storeId); return storeId; } /** Returns all context store injection requests for the class being transformed. */ static BitSet getInjectedContextStores() { - BitSet injectedStoreIds = INJECTED_STORE_IDS.get(); + Pair injectedStoreIds = INJECTED_STORE_IDS.get(); if (null != injectedStoreIds) { INJECTED_STORE_IDS.remove(); + return injectedStoreIds.getRight(); + } else { + return null; } - return injectedStoreIds; } private static final class SerialVersionUIDInjector