Fix a bug when FieldBackedProvider uses a map and value is set to null
MODIFY
auto-api/src/main/java/io/opentelemetry/instrumentation/auto/api/WeakMap.java
auto-api/src/main/java/io/opentelemetry/instrumentation/auto/api/WeakMap.java
@@ -37,6 +37,8 @@ public interface WeakMap<K, V> {
V computeIfAbsent(K key, ValueSupplier<? super K, ? extends V> supplier);
+ V remove(K key);
+
class Provider {
private static final Logger log = LoggerFactory.getLogger(Provider.class);
@@ -145,6 +147,11 @@ public interface WeakMap<K, V> {
}
}
+ @Override
+ public V remove(K key) {
+ return map.remove(key);
+ }
+
@Override
public String toString() {
return map.toString();

MODIFY
auto-api/src/test/groovy/io/opentelemetry/instrumentation/auto/api/WeakMapTest.groovy
auto-api/src/test/groovy/io/opentelemetry/instrumentation/auto/api/WeakMapTest.groovy
@@ -55,6 +55,19 @@ class WeakMapTest extends Specification {
supplier.counter == 2
}
+ def "remove a value"() {
+ given:
+ weakMap.put('key', 42)
+
+ when:
+ def removed = weakMap.remove('key')
+
+ then:
+ removed == 42
+ weakMap.get('key') == null
+ weakMap.size() == 0
+ }
+
class CounterSupplier implements WeakMap.ValueSupplier<String, Integer> {
def counter = 0

MODIFY
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/WeakMapSuppliers.java
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/WeakMapSuppliers.java
@@ -121,6 +121,11 @@ class WeakMapSuppliers {
}
}
}
+
+ @Override
+ public V remove(K key) {
+ return map.remove(key);
+ }
}
static class Inline implements WeakMap.Implementation {

MODIFY
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/FieldBackedProvider.java
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/FieldBackedProvider.java
@@ -936,8 +936,12 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
}
private void mapPut(Object key, Object value) {
+ if (value == null) {
+ map.remove(key);
+ } else {
map.put(key, value);
}
+ }
private Object mapSynchronizeInstance(Object key) {
return map;

MODIFY
testing-common/src/test/groovy/context/FieldBackedProviderTest.groovy
testing-common/src/test/groovy/context/FieldBackedProviderTest.groovy
@@ -126,6 +126,22 @@ class FieldBackedProviderTest extends AgentTestRunner {
new UntransformableKeyClass() | _
}
+ def "remove test"() {
+ given:
+ instance1.putContextCount(10)
+
+ when:
+ instance1.removeContextCount()
+
+ then:
+ instance1.getContextCount() == 0
+
+ where:
+ instance1 | _
+ new KeyClass() | _
+ new UntransformableKeyClass() | _
+ }
+
def "works with cglib enhanced instances which duplicates context getter and setter methods"() {
setup:
Enhancer enhancer = new Enhancer()

MODIFY
testing-common/src/test/java/context/ContextTestInstrumentation.java
testing-common/src/test/java/context/ContextTestInstrumentation.java
@@ -57,6 +57,7 @@ public class ContextTestInstrumentation extends Instrumenter.Default {
named("incrementContextCount"), StoreAndIncrementApiUsageAdvice.class.getName());
transformers.put(named("getContextCount"), GetApiUsageAdvice.class.getName());
transformers.put(named("putContextCount"), PutApiUsageAdvice.class.getName());
+ transformers.put(named("removeContextCount"), RemoveApiUsageAdvice.class.getName());
transformers.put(
named("incorrectKeyClassUsage"), IncorrectKeyClassContextApiUsageAdvice.class.getName());
transformers.put(
@@ -116,7 +117,8 @@ public class ContextTestInstrumentation extends Instrumenter.Default {
@Advice.This KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) {
ContextStore<KeyClass, Context> contextStore =
InstrumentationContext.get(KeyClass.class, Context.class);
- contextCount = contextStore.get(thiz).count;
+ Context context = contextStore.get(thiz);
+ contextCount = context == null ? 0 : context.count;
}
}
@@ -131,6 +133,15 @@ public class ContextTestInstrumentation extends Instrumenter.Default {
}
}
+ public static class RemoveApiUsageAdvice {
+ @Advice.OnMethodExit
+ public static void methodExit(@Advice.This KeyClass thiz) {
+ ContextStore<KeyClass, Context> contextStore =
+ InstrumentationContext.get(KeyClass.class, Context.class);
+ contextStore.put(thiz, null);
+ }
+ }
+
public static class IncorrectKeyClassContextApiUsageAdvice {
@Advice.OnMethodExit
public static void methodExit() {
@@ -191,6 +202,10 @@ public class ContextTestInstrumentation extends Instrumenter.Default {
public void putContextCount(int value) {
// implementation replaced with test instrumentation
}
+
+ public void removeContextCount() {
+ // implementation replaced with test instrumentation
+ }
}
/**

Fix a bug when FieldBackedProvider uses a map and value is set to null (#1177)