Skip to content

Commit

Permalink
Avoid collected store ids from leaking to later field-injections when…
Browse files Browse the repository at this point in the history
… an instrumentation fails (#7403)

Co-authored-by: luneo7 <[email protected]>
  • Loading branch information
mcculls and luneo7 authored Aug 9, 2024
1 parent 2f32a35 commit bb44d60
Showing 1 changed file with 14 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<BitSet> INJECTED_STORE_IDS = new ThreadLocal<>();
static final ThreadLocal<Pair<String, BitSet>> INJECTED_STORE_IDS = new ThreadLocal<>();

final boolean serialVersionUIDFieldInjection =
InstrumenterConfig.get().isSerialVersionUIDFieldInjection();
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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<String, BitSet> 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<String, BitSet> 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
Expand Down

0 comments on commit bb44d60

Please sign in to comment.