Switching to use Span concurrent map in Scala plugin
package co.elastic.apm.agent.scalaconcurrent;
import co.elastic.apm.agent.bci.TracerAwareInstrumentation;
+import co.elastic.apm.agent.collections.WeakConcurrentProviderImpl;
import co.elastic.apm.agent.impl.transaction.AbstractSpan;
-import co.elastic.apm.agent.sdk.weakconcurrent.WeakConcurrent;
import co.elastic.apm.agent.sdk.weakconcurrent.WeakMap;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
@@ -39,8 +39,7 @@ import static net.bytebuddy.matcher.ElementMatchers.returns;
public abstract class FutureInstrumentation extends TracerAwareInstrumentation {
@SuppressWarnings("WeakerAccess")
- public static final WeakMap<Object, AbstractSpan<?>> promisesToContext =
- WeakConcurrent.buildMap();
+ public static final WeakMap<Object, AbstractSpan<?>> promisesToContext = WeakConcurrentProviderImpl.createWeakSpanMap();
@Nonnull
@Override
@@ -66,9 +65,6 @@ public abstract class FutureInstrumentation extends TracerAwareInstrumentation {
final AbstractSpan<?> context = tracer.getActive();
if (context != null) {
promisesToContext.put(thiz, context);
- // this span might be ended before the Promise$Transformation#run method starts
- // we have to avoid that this span gets recycled, even in the above mentioned case
- context.incrementReferences();
}
}
}
@@ -91,12 +87,13 @@ public abstract class FutureInstrumentation extends TracerAwareInstrumentation {
@Nullable
@Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
public static Object onEnter(@Advice.This Object thiz) {
- AbstractSpan<?> context = promisesToContext.remove(thiz);
+ // We cannot remove yet, as this may decrement the ref count of the span to 0 if it has already ended,
+ // thus causing it to be recycled just before we activate it on the current thread. So we first get().
+ AbstractSpan<?> context = promisesToContext.get(thiz);
if (context != null) {
context.activate();
- // decrements the reference we incremented to avoid that the parent context gets recycled before the promise is run
- // because we have activated it, we can be sure it doesn't get recycled until we deactivate in the OnMethodExit advice
- context.decrementReferences();
+ // Now it's safe to remove, as ref count is at least 2
+ promisesToContext.remove(thiz);
}
return context;
}
