Fix a bug when FieldBackedProvider uses a map and value is set to null

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 V computeIfAbsent(K key, ValueSupplier<? super K, ? extends V> supplier)
       }
     }
 
+    @Override
+    public V remove(K key) {
+      return map.remove(key);
+    }
+
     @Override
     public String toString() {
       return map.toString();

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

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/WeakMapSuppliers.java

@@ -121,6 +121,11 @@ public V computeIfAbsent(K key, ValueSupplier<? super K, ? extends V> supplier)
           }
         }
       }
+
+      @Override
+      public V remove(K key) {
+        return map.remove(key);
+      }
     }
 
     static class Inline implements WeakMap.Implementation {

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/context/FieldBackedProvider.java

@@ -936,7 +936,11 @@ private Object mapGet(Object key) {
     }
 
     private void mapPut(Object key, Object value) {
-      map.put(key, value);
+      if (value == null) {
+        map.remove(key);
+      } else {
+        map.put(key, value);
+      }
     }
 
     private Object mapSynchronizeInstance(Object key) {

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()

testing-common/src/test/java/context/ContextTestInstrumentation.java

@@ -57,6 +57,7 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
         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 static void methodExit(
         @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 static void methodExit(@Advice.This KeyClass thiz, @Advice.Argument(0) in
     }
   }
 
+  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 int getContextCount() {
     public void putContextCount(int value) {
       // implementation replaced with test instrumentation
     }
+
+    public void removeContextCount() {
+      // implementation replaced with test instrumentation
+    }
   }
 
   /**