From 4cca40daa9d7bbadc06f5a2c47e474d6ec5d538c Mon Sep 17 00:00:00 2001
From: eyalkoren <41850454+eyalkoren@users.noreply.github.com>
Date: Wed, 8 Dec 2021 11:56:48 +0200
Subject: [PATCH] Switching to use Span concurrent map in Scala plugin

---
 .../scalaconcurrent/FutureInstrumentation.java  | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/apm-agent-plugins/apm-scala-concurrent-plugin/src/main/java/co/elastic/apm/agent/scalaconcurrent/FutureInstrumentation.java b/apm-agent-plugins/apm-scala-concurrent-plugin/src/main/java/co/elastic/apm/agent/scalaconcurrent/FutureInstrumentation.java
index a1b7903c77..2e4660a040 100644
--- a/apm-agent-plugins/apm-scala-concurrent-plugin/src/main/java/co/elastic/apm/agent/scalaconcurrent/FutureInstrumentation.java
+++ b/apm-agent-plugins/apm-scala-concurrent-plugin/src/main/java/co/elastic/apm/agent/scalaconcurrent/FutureInstrumentation.java
@@ -19,8 +19,8 @@
 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 @@
 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 static void onExit(@Advice.This Object thiz) {
                 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 static class AdviceClass {
             @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;
             }