Factor out reference counting to explicit API

apm-agent-core/src/main/java/co/elastic/apm/agent/collections/SpanConcurrentHashMap.java → apm-agent-core/src/main/java/co/elastic/apm/agent/collections/ReferencedCountedConcurrentHashMap.java

@@ -18,24 +18,24 @@
  */
 package co.elastic.apm.agent.collections;
 
-import co.elastic.apm.agent.tracer.AbstractSpan;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounted;
 
 import javax.annotation.Nullable;
 import java.util.Iterator;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
- * Hash map dedicated to storage of in-flight spans and transactions, reference count is being incremented/decremented
+ * Hash map dedicated to storage of reference counted objects where counts are being incremented/decremented
  * when entry is added/removed. Usage of this map is intended for providing GC-based storage of context associated
  * to a framework-level object key, when the latter is collected by GC it allows to decrement and then recycle the
  * span/transaction.
  *
  * @param <K> key type
  * @param <V> context type
  */
-public class SpanConcurrentHashMap<K, V extends AbstractSpan<?>> extends ConcurrentHashMap<K, V> {
+public class ReferencedCountedConcurrentHashMap<K, V extends ReferenceCounted> extends ConcurrentHashMap<K, V> {
 
-    SpanConcurrentHashMap() {
+    ReferencedCountedConcurrentHashMap() {
     }
 
     @Nullable
@@ -71,7 +71,7 @@ public void clear() {
         super.clear();
     }
 
-    private void onPut(@Nullable AbstractSpan<?> previous, AbstractSpan<?> value) {
+    private void onPut(@Nullable ReferenceCounted previous, ReferenceCounted value) {
         if (previous == null) {
             // new entry
             value.incrementReferences();
@@ -82,7 +82,7 @@ private void onPut(@Nullable AbstractSpan<?> previous, AbstractSpan<?> value) {
         }
     }
 
-    private void onRemove(@Nullable AbstractSpan<?> removed) {
+    private void onRemove(@Nullable ReferenceCounted removed) {
         if (removed == null) {
             return;
         }

apm-agent-core/src/main/java/co/elastic/apm/agent/collections/WeakConcurrentProviderImpl.java

@@ -22,7 +22,7 @@
 import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
 import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
 import co.elastic.apm.agent.sdk.weakconcurrent.WeakSet;
-import co.elastic.apm.agent.tracer.AbstractSpan;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounted;
 import com.blogspot.mydailyjava.weaklockfree.AbstractWeakConcurrentMap;
 import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentSet;
 
@@ -37,8 +37,8 @@ public class WeakConcurrentProviderImpl implements WeakConcurrent.WeakConcurrent
 
     private static final WeakConcurrentSet<AbstractWeakConcurrentMap<?, ?, ?>> registeredMaps = new WeakConcurrentSet<>(WeakConcurrentSet.Cleaner.INLINE);
 
-    public static <K, V extends AbstractSpan<?>> WeakMap<K, V> createWeakSpanMap() {
-        SpanConcurrentHashMap<AbstractWeakConcurrentMap.WeakKey<K>, V> map = new SpanConcurrentHashMap<>();
+    public static <K, V extends ReferenceCounted> WeakMap<K, V> createWeakReferenceCountedMap() {
+        ReferencedCountedConcurrentHashMap<AbstractWeakConcurrentMap.WeakKey<K>, V> map = new ReferencedCountedConcurrentHashMap<>();
         NullSafeWeakConcurrentMap<K, V> result = new NullSafeWeakConcurrentMap<>(map);
         registeredMaps.add(result);
         return result;

apm-agent-core/src/main/java/co/elastic/apm/agent/collections/WeakMapReferenceCounter.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package co.elastic.apm.agent.collections;
+
+import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounted;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounter;
+
+import javax.annotation.Nullable;
+
+public class WeakMapReferenceCounter<K, V extends ReferenceCounted> implements ReferenceCounter<K, V> {
+
+    private final WeakMap<K, V> map;
+
+    public WeakMapReferenceCounter(WeakMap<K, V> map) {
+        this.map = map;
+    }
+
+    @Override
+    @Nullable
+    public V get(K key) {
+        return map.get(key);
+    }
+
+    @Override
+    public boolean contains(K key) {
+        return map.containsKey(key);
+    }
+
+    @Override
+    public void put(K key, V value) {
+        map.put(key, value);
+    }
+
+    @Override
+    @Nullable
+    public V remove(K key) {
+        return map.remove(key);
+    }
+}

apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java

@@ -18,6 +18,8 @@
  */
 package co.elastic.apm.agent.impl;
 
+import co.elastic.apm.agent.collections.WeakConcurrentProviderImpl;
+import co.elastic.apm.agent.collections.WeakMapReferenceCounter;
 import co.elastic.apm.agent.common.JvmRuntimeInfo;
 import co.elastic.apm.agent.common.util.WildcardMatcher;
 import co.elastic.apm.agent.configuration.CoreConfiguration;
@@ -49,6 +51,10 @@
 import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
 import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
 import co.elastic.apm.agent.tracer.GlobalTracer;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounted;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounter;
+import co.elastic.apm.agent.util.DependencyInjectingServiceLoader;
+import co.elastic.apm.agent.util.ExecutorUtils;
 import co.elastic.apm.agent.tracer.Scope;
 import co.elastic.apm.agent.tracer.dispatch.BinaryHeaderGetter;
 import co.elastic.apm.agent.tracer.dispatch.TextHeaderGetter;
@@ -597,6 +603,11 @@ public ObjectPoolFactory getObjectPoolFactory() {
         return objectPoolFactory;
     }
 
+    @Override
+    public <K, V extends ReferenceCounted> ReferenceCounter<K, V> createReferenceCounter() {
+        return new WeakMapReferenceCounter<>(WeakConcurrentProviderImpl.<K, V>createWeakReferenceCountedMap());
+    }
+
     @Override
     @Nullable
     public AbstractSpan<?> getActive() {
@@ -938,5 +949,4 @@ public <T extends co.elastic.apm.agent.tracer.Tracer> T require(Class<T> type) {
     public Set<String> getTraceHeaderNames() {
         return TraceContext.TRACE_TEXTUAL_HEADERS;
     }
-
 }

apm-agent-core/src/test/java/co/elastic/apm/agent/collections/SpanConcurrentHashMapTest.java → apm-agent-core/src/test/java/co/elastic/apm/agent/collections/ReferencedCountedConcurrentHashMapTest.java

@@ -34,7 +34,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.awaitility.Awaitility.await;
 
-class SpanConcurrentHashMapTest {
+class ReferencedCountedConcurrentHashMapTest {
 
     @Nullable
     private Object key;
@@ -45,7 +45,7 @@ void putRemove() {
         checkRefCount(testSpan, 0);
 
         key = new Object();
-        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakSpanMap();
+        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakReferenceCountedMap();
         map.put(key, testSpan);
 
         checkRefCount(testSpan, 1);
@@ -63,7 +63,7 @@ void putRemove() {
     void putTwice(PutOperation operation) {
         TestSpan testSpan = new TestSpan();
         key = new Object();
-        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakSpanMap();
+        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakReferenceCountedMap();
 
         checkRefCount(testSpan, 0);
 
@@ -82,7 +82,7 @@ void swapValues() {
         TestSpan ts2 = new TestSpan();
 
         key = new Object();
-        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakSpanMap();
+        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakReferenceCountedMap();
 
         map.put(key, ts1);
         map.put(key, ts2);
@@ -99,7 +99,7 @@ void testPutIfAbsent() {
         TestSpan ts2 = new TestSpan();
 
         key = new Object();
-        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakSpanMap();
+        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakReferenceCountedMap();
 
         map.putIfAbsent(key, ts1);
         map.putIfAbsent(key, ts2);
@@ -130,7 +130,7 @@ void execute(WeakMap<Object, TestSpan> map, Object key, TestSpan value) {
 
     @Test
     void clear() {
-        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakSpanMap();
+        WeakMap<Object, TestSpan> map = WeakConcurrentProviderImpl.createWeakReferenceCountedMap();
 
         List<AbstractSpan<?>> list = new ArrayList<>();
         for (int i = 0; i < 5; i++) {
@@ -153,7 +153,7 @@ void weakMapDecrementOnStaleKeyGC() {
         key = new Object();
         TestSpan span = new TestSpan();
 
-        WeakMap<Object, AbstractSpan<?>> map = WeakConcurrentProviderImpl.createWeakSpanMap();
+        WeakMap<Object, AbstractSpan<?>> map = WeakConcurrentProviderImpl.createWeakReferenceCountedMap();
 
         map.put(key, span);
 

apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/AbstractAsyncHttpClientInstrumentation.java

@@ -19,14 +19,13 @@
 package co.elastic.apm.agent.asynchttpclient;
 
 import co.elastic.apm.agent.bci.TracerAwareInstrumentation;
-import co.elastic.apm.agent.collections.WeakConcurrentProviderImpl;
 import co.elastic.apm.agent.httpclient.HttpClientHelper;
 import co.elastic.apm.agent.tracer.AbstractSpan;
 import co.elastic.apm.agent.tracer.Outcome;
 import co.elastic.apm.agent.tracer.Span;
 import co.elastic.apm.agent.sdk.DynamicTransformer;
 import co.elastic.apm.agent.sdk.ElasticApmInstrumentation;
-import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounter;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
@@ -54,7 +53,7 @@ public abstract class AbstractAsyncHttpClientInstrumentation extends TracerAware
 
     public static class Helper {
 
-        static final WeakMap<AsyncHandler<?>, Span> handlerSpanMap = WeakConcurrentProviderImpl.createWeakSpanMap();
+        static final ReferenceCounter<AsyncHandler<?>, Span> handlerSpanMap = tracer.createReferenceCounter();
 
         public static final List<Class<? extends ElasticApmInstrumentation>> ASYNC_HANDLER_INSTRUMENTATIONS = Arrays.<Class<? extends ElasticApmInstrumentation>>asList(
             AsyncHandlerOnCompletedInstrumentation.class,

apm-agent-plugins/apm-dubbo-plugin/src/main/java/co/elastic/apm/agent/dubbo/AlibabaCallbackHolder.java

@@ -18,11 +18,11 @@
  */
 package co.elastic.apm.agent.dubbo;
 
-import co.elastic.apm.agent.collections.WeakConcurrentProviderImpl;
 import co.elastic.apm.agent.tracer.AbstractSpan;
-import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
+import co.elastic.apm.agent.tracer.GlobalTracer;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounter;
 import com.alibaba.dubbo.remoting.exchange.ResponseCallback;
 
 public class AlibabaCallbackHolder {
-    public static final WeakMap<ResponseCallback, AbstractSpan<?>> callbackSpanMap = WeakConcurrentProviderImpl.createWeakSpanMap();
+    public static final ReferenceCounter<ResponseCallback, AbstractSpan<?>> callbackSpanMap = GlobalTracer.get().createReferenceCounter();
 }

apm-agent-plugins/apm-grpc/apm-grpc-plugin/src/main/java/co/elastic/apm/agent/grpc/GrpcHelper.java

@@ -18,7 +18,6 @@
  */
 package co.elastic.apm.agent.grpc;
 
-import co.elastic.apm.agent.collections.WeakConcurrentProviderImpl;
 import co.elastic.apm.agent.tracer.GlobalTracer;
 import co.elastic.apm.agent.tracer.AbstractSpan;
 import co.elastic.apm.agent.tracer.Outcome;
@@ -31,6 +30,7 @@
 import co.elastic.apm.agent.tracer.dispatch.HeaderUtils;
 import co.elastic.apm.agent.tracer.dispatch.TextHeaderGetter;
 import co.elastic.apm.agent.tracer.dispatch.TextHeaderSetter;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounter;
 import io.grpc.CallOptions;
 import io.grpc.ClientCall;
 import io.grpc.Metadata;
@@ -58,27 +58,27 @@ public static GrpcHelper getInstance() {
     /**
      * Map of all in-flight {@link Span} with {@link ClientCall} instance as key.
      */
-    private final WeakMap<ClientCall<?, ?>, Span<?>> clientCallSpans;
+    private final ReferenceCounter<ClientCall<?, ?>, Span<?>> clientCallSpans;
 
     /**
      * Map of all in-flight {@link Span} with {@link ClientCall} instance as key.
      */
-    private final WeakMap<ClientCall<?, ?>, Span<?>> delayedClientCallSpans;
+    private final ReferenceCounter<ClientCall<?, ?>, Span<?>> delayedClientCallSpans;
 
     /**
      * Map of all in-flight {@link Span} with {@link ClientCall.Listener} instance as key.
      */
-    private final WeakMap<ClientCall.Listener<?>, Span<?>> clientCallListenerSpans;
+    private final ReferenceCounter<ClientCall.Listener<?>, Span<?>> clientCallListenerSpans;
 
     /**
      * Map of all in-flight {@link Transaction} with {@link ServerCall.Listener} instance as key.
      */
-    private final WeakMap<ServerCall.Listener<?>, Transaction<?>> serverListenerTransactions;
+    private final ReferenceCounter<ServerCall.Listener<?>, Transaction<?>> serverListenerTransactions;
 
     /**
      * Map of all in-flight {@link Transaction} with {@link ServerCall} instance as key.
      */
-    private final WeakMap<ServerCall<?, ?>, Transaction<?>> serverCallTransactions;
+    private final ReferenceCounter<ServerCall<?, ?>, Transaction<?>> serverCallTransactions;
 
     /**
      * gRPC header cache used to minimize allocations
@@ -91,19 +91,18 @@ public static GrpcHelper getInstance() {
     private final Tracer tracer;
 
     public GrpcHelper() {
-        clientCallSpans = WeakConcurrentProviderImpl.createWeakSpanMap();
-        delayedClientCallSpans = WeakConcurrentProviderImpl.createWeakSpanMap();
-        clientCallListenerSpans = WeakConcurrentProviderImpl.createWeakSpanMap();
+        tracer = GlobalTracer.get();
+        clientCallSpans = tracer.createReferenceCounter();
+        delayedClientCallSpans = tracer.createReferenceCounter();
+        clientCallListenerSpans = tracer.createReferenceCounter();
 
-        serverListenerTransactions = WeakConcurrentProviderImpl.createWeakSpanMap();
-        serverCallTransactions = WeakConcurrentProviderImpl.createWeakSpanMap();
+        serverListenerTransactions = tracer.createReferenceCounter();
+        serverCallTransactions = tracer.createReferenceCounter();
 
         headerCache = WeakConcurrent.buildMap();
 
         headerSetter = new GrpcHeaderSetter();
         headerGetter = new GrpcHeaderGetter();
-
-        tracer = GlobalTracer.get();
     }
 
     // transaction management (server part)
@@ -480,7 +479,7 @@ public void clientCallStartExit(@Nullable Span<?> spanFromEntry, ClientCall.List
     }
 
     public void cancelCall(ClientCall<?, ?> clientCall, @Nullable Throwable cause) {
-        WeakMap<ClientCall<?, ?>, Span<?>> clientCallMap = (isDelayedClientCall(clientCall)) ? delayedClientCallSpans : clientCallSpans;
+        ReferenceCounter<ClientCall<?, ?>, Span<?>> clientCallMap = (isDelayedClientCall(clientCall)) ? delayedClientCallSpans : clientCallSpans;
         // we can't remove yet, in order to avoid reference decrement prematurely
         Span<?> span = clientCallMap.get(clientCall);
         if (span != null) {

apm-agent-plugins/apm-grpc/apm-grpc-plugin/src/test/java/co/elastic/apm/agent/grpc/GrpcHelperTest.java

@@ -18,6 +18,7 @@
  */
 package co.elastic.apm.agent.grpc;
 
+import co.elastic.apm.agent.AbstractInstrumentationTest;
 import co.elastic.apm.agent.tracer.Outcome;
 import io.grpc.Status;
 import org.junit.jupiter.api.Test;
@@ -26,7 +27,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
-class GrpcHelperTest {
+class GrpcHelperTest extends AbstractInstrumentationTest {
 
     @ParameterizedTest
     @EnumSource(Status.Code.class)

apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java

@@ -18,14 +18,14 @@
  */
 package co.elastic.apm.agent.concurrent;
 
-import co.elastic.apm.agent.collections.WeakConcurrentProviderImpl;
 import co.elastic.apm.agent.common.ThreadUtils;
 import co.elastic.apm.agent.tracer.AbstractSpan;
 import co.elastic.apm.agent.sdk.DynamicTransformer;
 import co.elastic.apm.agent.sdk.ElasticApmInstrumentation;
 import co.elastic.apm.agent.sdk.state.GlobalState;
-import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
+import co.elastic.apm.agent.tracer.GlobalTracer;
 import co.elastic.apm.agent.tracer.Tracer;
+import co.elastic.apm.agent.tracer.reference.ReferenceCounter;
 
 import javax.annotation.Nullable;
 import java.util.ArrayList;
@@ -42,7 +42,7 @@
 @GlobalState
 public class JavaConcurrent {
 
-    private static final WeakMap<Object, AbstractSpan<?>> contextMap = WeakConcurrentProviderImpl.createWeakSpanMap();
+    private static final ReferenceCounter<Object, AbstractSpan<?>> contextMap = GlobalTracer.get().createReferenceCounter();
 
     private static final List<Class<? extends ElasticApmInstrumentation>> RUNNABLE_CALLABLE_FJTASK_INSTRUMENTATION = Collections.
         <Class<? extends ElasticApmInstrumentation>>singletonList(RunnableCallableForkJoinTaskInstrumentation.class);