Fix empty context activation behaviour (#3227)

Author: JonasKunz

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

@@ -58,8 +58,11 @@ class ActiveStack {
      */
     private final Deque<ElasticContext<?>> activeContextStack = new ArrayDeque<ElasticContext<?>>();
 
-    ActiveStack(int stackMaxDepth) {
+    private final EmptyElasticContext emptyContext;
+
+    ActiveStack(int stackMaxDepth, EmptyElasticContext emptyContextForTracer) {
         this.stackMaxDepth = stackMaxDepth;
+        this.emptyContext = emptyContextForTracer;
     }
 
     @Nullable
@@ -80,7 +83,7 @@ public ElasticContext<?> currentContext() {
         if (current instanceof ElasticContextWrapper) {
             return ((ElasticContextWrapper<?>) current).getWrappedContext();
         }
-        return current != null ? current : EmptyElasticContext.INSTANCE;
+        return current != null ? current : emptyContext;
     }
 
     boolean activate(ElasticContext<?> context, List<ActivationListener> activationListeners) {

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

@@ -107,10 +107,12 @@ public class ElasticApmTracer implements Tracer {
     private final Reporter reporter;
     private final ObjectPoolFactory objectPoolFactory;
 
+    private final EmptyElasticContext emptyContext;
+
     private final ThreadLocal<ActiveStack> activeStack = new ThreadLocal<ActiveStack>() {
         @Override
         protected ActiveStack initialValue() {
-            return new ActiveStack(transactionMaxSpans);
+            return new ActiveStack(transactionMaxSpans, emptyContext);
         }
     };
 
@@ -186,6 +188,7 @@ private static void checkClassloader() {
 
     ElasticApmTracer(ConfigurationRegistry configurationRegistry, MetricRegistry metricRegistry, Reporter reporter, ObjectPoolFactory poolFactory,
                      ApmServerClient apmServerClient, final String ephemeralId, MetaDataFuture metaDataFuture) {
+        this.emptyContext = new EmptyElasticContext(this);
         this.metricRegistry = metricRegistry;
         this.configurationRegistry = configurationRegistry;
         this.reporter = reporter;

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

@@ -20,15 +20,13 @@
 
 import co.elastic.apm.agent.impl.transaction.AbstractSpan;
 import co.elastic.apm.agent.impl.transaction.ElasticContext;
-import co.elastic.apm.agent.tracer.Scope;
 
 import javax.annotation.Nullable;
 
 class EmptyElasticContext extends ElasticContext<EmptyElasticContext> {
 
-    static final ElasticContext<?> INSTANCE = new EmptyElasticContext();
-
-    private EmptyElasticContext() {
+    EmptyElasticContext(ElasticApmTracer tracer) {
+        super(tracer);
     }
 
     @Nullable
@@ -37,21 +35,6 @@ public AbstractSpan<?> getSpan() {
         return null;
     }
 
-    @Override
-    public EmptyElasticContext activate() {
-        return this;
-    }
-
-    @Override
-    public EmptyElasticContext deactivate() {
-        return this;
-    }
-
-    @Override
-    public Scope activateInScope() {
-        return NoopScope.INSTANCE;
-    }
-
     @Override
     public void incrementReferences() {
 

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

@@ -28,7 +28,6 @@
 import co.elastic.apm.agent.sdk.logging.LoggerFactory;
 import co.elastic.apm.agent.sdk.internal.util.LoggerUtils;
 import co.elastic.apm.agent.tracer.Outcome;
-import co.elastic.apm.agent.tracer.Scope;
 import co.elastic.apm.agent.tracer.dispatch.BinaryHeaderGetter;
 import co.elastic.apm.agent.tracer.dispatch.HeaderGetter;
 import co.elastic.apm.agent.tracer.dispatch.TextHeaderGetter;
@@ -55,7 +54,6 @@ public abstract class AbstractSpan<T extends AbstractSpan<T>> extends ElasticCon
      */
     protected final StringBuilder name = new StringBuilder();
     protected final boolean collectBreakdownMetrics;
-    protected final ElasticApmTracer tracer;
     protected final AtomicLong timestamp = new AtomicLong();
     protected final AtomicLong endTimestamp = new AtomicLong();
 
@@ -197,7 +195,7 @@ public long getDuration() {
     }
 
     public AbstractSpan(ElasticApmTracer tracer) {
-        this.tracer = tracer;
+        super(tracer);
         traceContext = TraceContext.with64BitId(this.tracer);
         boolean selfTimeCollectionEnabled = !WildcardMatcher.isAnyMatch(tracer.getConfig(ReporterConfiguration.class).getDisableMetrics(), "span.self_time");
         boolean breakdownMetricsEnabled = tracer.getConfig(CoreConfiguration.class).isBreakdownMetricsEnabled();
@@ -625,23 +623,6 @@ private boolean hasChildId(Id spanId) {
      */
     public abstract Transaction getParentTransaction();
 
-    @Override
-    public T activate() {
-        tracer.activate(this);
-        return thiz();
-    }
-
-    @Override
-    public T deactivate() {
-        tracer.deactivate(this);
-        return thiz();
-    }
-
-    @Override
-    public Scope activateInScope() {
-        return tracer.activateInScope(this);
-    }
-
     /**
      * Set start timestamp
      *

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

@@ -18,6 +18,8 @@
  */
 package co.elastic.apm.agent.impl.transaction;
 
+import co.elastic.apm.agent.impl.ElasticApmTracer;
+import co.elastic.apm.agent.tracer.Scope;
 import co.elastic.apm.agent.tracer.dispatch.BinaryHeaderSetter;
 import co.elastic.apm.agent.tracer.dispatch.HeaderUtils;
 import co.elastic.apm.agent.tracer.dispatch.TextHeaderGetter;
@@ -27,6 +29,12 @@
 
 public abstract class ElasticContext<T extends ElasticContext<T>> implements co.elastic.apm.agent.tracer.ElasticContext<T> {
 
+    protected final ElasticApmTracer tracer;
+
+    protected ElasticContext(ElasticApmTracer tracer) {
+        this.tracer = tracer;
+    }
+
     @Nullable
     public abstract AbstractSpan<?> getSpan();
 
@@ -39,6 +47,25 @@ public final Transaction getTransaction() {
         return contextSpan != null ? contextSpan.getParentTransaction() : null;
     }
 
+    @Override
+    @SuppressWarnings("unchecked")
+    public T activate() {
+        tracer.activate(this);
+        return (T) this;
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public T deactivate() {
+        tracer.deactivate(this);
+        return (T) this;
+    }
+
+    @Override
+    public Scope activateInScope() {
+        return tracer.activateInScope(this);
+    }
+
     @Nullable
     @Override
     public co.elastic.apm.agent.impl.transaction.Span createSpan() {

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

@@ -59,6 +59,7 @@ public class ElasticContextWrapper<T extends ElasticContext<T>> extends ElasticC
     private final Map<Class<?>, ElasticContext<?>> contextWrappers;
 
     public ElasticContextWrapper(int initialSize, ElasticContext<T> context) {
+        super(context.tracer);
         this.contextWrappers = new HashMap<>(initialSize, 1.0f);
         this.context = context;
     }

apm-agent-core/src/test/java/co/elastic/apm/agent/impl/ElasticApmTracerTest.java

@@ -544,6 +544,24 @@ void testActivateDeactivateTwice() {
         transaction.end();
     }
 
+    @Test
+    void testEmptyContextActivation() {
+        final Transaction transaction = startTestRootTransaction();
+        assertThat(tracerImpl.currentContext().getTransaction()).isNull();
+        tracerImpl.activate(transaction);
+        assertThat(tracerImpl.currentContext().getTransaction()).isEqualTo(transaction);
+
+        EmptyElasticContext empty = new EmptyElasticContext(tracerImpl);
+        empty.activate();
+        assertThat(tracerImpl.currentContext().getTransaction()).isNull();
+
+        empty.deactivate();
+        assertThat(tracerImpl.currentContext().getTransaction()).isEqualTo(transaction);
+        tracerImpl.deactivate(transaction);
+        assertThat(tracerImpl.currentContext().getTransaction()).isNull();
+        transaction.end();
+    }
+
     @Test
     void testOverrideServiceNameWithoutExplicitServiceName() {
         final ElasticApmTracer tracer = new ElasticApmTracerBuilder()
@@ -714,19 +732,8 @@ void testUnknownConfiguration() {
 
     private static final class TestContext extends ElasticContext<TestContext> {
 
-        @Override
-        public TestContext activate() {
-            return null;
-        }
-
-        @Override
-        public TestContext deactivate() {
-            return null;
-        }
-
-        @Override
-        public Scope activateInScope() {
-            return null;
+        private TestContext() {
+            super(null);
         }
 
         @Nullable

apm-agent-plugins/apm-opentelemetry/apm-opentelemetry-plugin/src/main/java/co/elastic/apm/agent/opentelemetry/tracing/OTelBridgeContext.java

@@ -51,10 +51,8 @@ public class OTelBridgeContext extends ElasticContext<OTelBridgeContext> impleme
      */
     private final Context otelContext;
 
-    private final ElasticApmTracer tracer;
-
     private OTelBridgeContext(ElasticApmTracer tracer, Context otelContext) {
-        this.tracer = tracer;
+        super(tracer);
         this.otelContext = otelContext;
     }
 
@@ -99,23 +97,6 @@ public static OTelBridgeContext wrapElasticActiveSpan(ElasticApmTracer tracer, A
         return new OTelBridgeContext(tracer, originalRootContext.with(otelSpan));
     }
 
-    @Override
-    public OTelBridgeContext activate() {
-        tracer.activate(this);
-        return this;
-    }
-
-    @Override
-    public co.elastic.apm.agent.tracer.Scope activateInScope() {
-        return tracer.activateInScope(this);
-    }
-
-    @Override
-    public OTelBridgeContext deactivate() {
-        tracer.deactivate(this);
-        return this;
-    }
-
     @Nullable
     @Override
     public AbstractSpan<?> getSpan() {