Context propagation refactorings

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

@@ -69,9 +69,9 @@ Transaction currentTransaction() {
     }
 
     /**
-     * @return the currently active context, {@literal null} if there is none.
+     * @return the current context, potentially empty when no span, transaction or baggage is currently active.
      */
-    @Nullable
+
     public ElasticContext<?> currentContext() {
         ElasticContext<?> current = activeContextStack.peek();
 
@@ -80,7 +80,7 @@ public ElasticContext<?> currentContext() {
         if (current instanceof ElasticContextWrapper) {
             return ((ElasticContextWrapper<?>) current).getWrappedContext();
         }
-        return current;
+        return current != null ? current : EmptyElasticContext.INSTANCE;
     }
 
     boolean activate(ElasticContext<?> context, List<ActivationListener> activationListeners) {
@@ -99,9 +99,9 @@ boolean activate(ElasticContext<?> context, List<ActivationListener> activationL
             return false;
         }
 
+        context.incrementReferences();
         AbstractSpan<?> span = context.getSpan();
         if (span != null) {
-            span.incrementReferences();
             triggerActivationListeners(span, true, activationListeners);
         }
 
@@ -122,22 +122,15 @@ boolean deactivate(ElasticContext<?> context, List<ActivationListener> activatio
         ElasticContext<?> activeContext = currentContext();
         activeContextStack.remove();
 
-        AbstractSpan<?> span = context.getSpan();
-
-        if (activeContext != context && activeContext instanceof ElasticContextWrapper) {
-            // when context has been wrapped, we need to get the underlying context
-            activeContext = ((ElasticContextWrapper<?>) activeContext).getWrappedContext();
-        }
-
         try {
             assertIsActive(context, activeContext, assertionsEnabled);
+
+            AbstractSpan<?> span = context.getSpan();
             if (null != span) {
                 triggerActivationListeners(span, false, activationListeners);
             }
         } finally {
-            if (null != span) {
-                span.decrementReferences();
-            }
+            context.decrementReferences();
         }
         return true;
     }

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

@@ -18,8 +18,8 @@
  */
 package co.elastic.apm.agent.impl;
 
-import co.elastic.apm.agent.collections.WeakReferenceCountedMap;
 import co.elastic.apm.agent.bci.IndyBootstrap;
+import co.elastic.apm.agent.collections.WeakReferenceCountedMap;
 import co.elastic.apm.agent.common.JvmRuntimeInfo;
 import co.elastic.apm.agent.common.util.WildcardMatcher;
 import co.elastic.apm.agent.configuration.CoreConfiguration;
@@ -51,6 +51,9 @@
 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.Scope;
+import co.elastic.apm.agent.tracer.dispatch.BinaryHeaderGetter;
+import co.elastic.apm.agent.tracer.dispatch.TextHeaderGetter;
 import co.elastic.apm.agent.tracer.reference.ReferenceCounted;
 import co.elastic.apm.agent.tracer.reference.ReferenceCountedMap;
 import co.elastic.apm.agent.util.DependencyInjectingServiceLoader;
@@ -614,33 +617,6 @@ public <K, V extends ReferenceCounted> ReferenceCountedMap<K, V> newReferenceCou
         return new WeakReferenceCountedMap<>();
     }
 
-    @Override
-    @Nullable
-    public AbstractSpan<?> getActive() {
-        ElasticContext<?> active = currentContext();
-        return active != null ? active.getSpan() : null;
-    }
-
-    @Nullable
-    @Override
-    public Span getActiveSpan() {
-        final AbstractSpan<?> active = getActive();
-        if (active instanceof Span) {
-            return (Span) active;
-        }
-        return null;
-    }
-
-    @Nullable
-    @Override
-    public Span getActiveExitSpan() {
-        final Span span = getActiveSpan();
-        if (span != null && span.isExit()) {
-            return span;
-        }
-        return null;
-    }
-
     public void registerSpanListener(ActivationListener activationListener) {
         this.activationListeners.add(activationListener);
     }
@@ -825,11 +801,17 @@ public <T> T getLifecycleListener(Class<T> listenerClass) {
     /**
      * @return the currently active context, {@literal null} if there is none.
      */
-    @Nullable
+
     public ElasticContext<?> currentContext() {
         return activeStack.get().currentContext();
     }
 
+    @Nullable
+    @Override
+    public AbstractSpan<?> getActive() {
+        return currentContext().getSpan();
+    }
+
     /**
      * Lazily wraps the currently active context if required, wrapper instance is cached with wrapperClass as key.
      * Wrapping is transparently handled by {@link #currentContext()}.

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

@@ -0,0 +1,65 @@
+/*
+ * 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.impl;
+
+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() {
+    }
+
+    @Nullable
+    @Override
+    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() {
+
+    }
+
+    @Override
+    public void decrementReferences() {
+
+    }
+
+}

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

@@ -22,6 +22,7 @@
 import co.elastic.apm.agent.impl.error.ErrorCapture;
 import co.elastic.apm.agent.impl.sampling.Sampler;
 import co.elastic.apm.agent.impl.transaction.AbstractSpan;
+import co.elastic.apm.agent.impl.transaction.ElasticContext;
 import co.elastic.apm.agent.impl.transaction.Span;
 import co.elastic.apm.agent.impl.transaction.Transaction;
 import co.elastic.apm.agent.tracer.dispatch.BinaryHeaderGetter;
@@ -100,15 +101,16 @@ <C> Transaction startChildTransaction(@Nullable C headerCarrier, BinaryHeaderGet
                                           Sampler sampler, long epochMicros, @Nullable ClassLoader initiatingClassLoader);
 
     @Override
-    @Nullable
-    Transaction currentTransaction();
+    ElasticContext<?> currentContext();
 
     @Override
     @Nullable
     AbstractSpan<?> getActive();
 
+    @Override
     @Nullable
-    Span getActiveSpan();
+    Transaction currentTransaction();
+
 
     /**
      * Captures an exception without providing an explicit reference to a parent {@link AbstractSpan}
@@ -124,9 +126,6 @@ <C> Transaction startChildTransaction(@Nullable C headerCarrier, BinaryHeaderGet
     @Nullable
     ErrorCapture captureException(@Nullable Throwable e, @Nullable AbstractSpan<?> parent, @Nullable ClassLoader initiatingClassLoader);
 
-    @Nullable
-    Span getActiveExitSpan();
-
     TracerState getState();
 
     @Nullable

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

@@ -30,10 +30,8 @@
 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.BinaryHeaderSetter;
 import co.elastic.apm.agent.tracer.dispatch.HeaderGetter;
 import co.elastic.apm.agent.tracer.dispatch.TextHeaderGetter;
-import co.elastic.apm.agent.tracer.dispatch.TextHeaderSetter;
 import co.elastic.apm.agent.tracer.pooling.Recyclable;
 
 import javax.annotation.Nullable;
@@ -44,7 +42,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 
-public abstract class AbstractSpan<T extends AbstractSpan<T>> implements Recyclable, ElasticContext<T>, co.elastic.apm.agent.tracer.AbstractSpan<T> {
+public abstract class AbstractSpan<T extends AbstractSpan<T>> extends ElasticContext<T> implements Recyclable, co.elastic.apm.agent.tracer.AbstractSpan<T> {
     private static final Logger logger = LoggerFactory.getLogger(AbstractSpan.class);
     private static final Logger oneTimeDuplicatedEndLogger = LoggerUtils.logOnce(logger);
     private static final Logger oneTimeMaxSpanLinksLogger = LoggerUtils.logOnce(logger);
@@ -619,6 +617,14 @@ private boolean hasChildId(Id spanId) {
         return false;
     }
 
+    /**
+     * Returns this, if this AbstractSpan is a {@link co.elastic.apm.agent.tracer.Transaction}.
+     * Otherwise returns the parent transaction of this span.
+     *
+     * @return the transaction.
+     */
+    public abstract Transaction getParentTransaction();
+
     @Override
     public T activate() {
         tracer.activate(this);
@@ -693,20 +699,6 @@ public void decrementReferences() {
 
     protected abstract void recycle();
 
-    @Override
-    public <C> void propagateTraceContext(C carrier, TextHeaderSetter<C> headerSetter) {
-        // the context of this span is propagated downstream so we can't discard it even if it's faster than span_min_duration
-        setNonDiscardable();
-        getTraceContext().propagateTraceContext(carrier, headerSetter);
-    }
-
-    @Override
-    public <C> boolean propagateTraceContext(C carrier, BinaryHeaderSetter<C> headerSetter) {
-        // the context of this span is propagated downstream so we can't discard it even if it's faster than span_min_duration
-        setNonDiscardable();
-        return getTraceContext().propagateTraceContext(carrier, headerSetter);
-    }
-
     @Override
     public void setNonDiscardable() {
         getTraceContext().setNonDiscardable();

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

@@ -18,20 +18,85 @@
  */
 package co.elastic.apm.agent.impl.transaction;
 
+import co.elastic.apm.agent.tracer.dispatch.BinaryHeaderSetter;
+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 javax.annotation.Nullable;
 
-public interface ElasticContext<T extends ElasticContext<T>> extends co.elastic.apm.agent.tracer.ElasticContext<T> {
+public abstract class ElasticContext<T extends ElasticContext<T>> implements co.elastic.apm.agent.tracer.ElasticContext<T> {
 
-    /**
-     * @return the span/transaction that is associated to this context, {@literal null} if there is none
-     */
     @Nullable
-    AbstractSpan<?> getSpan();
+    public abstract AbstractSpan<?> getSpan();
 
     /**
      * @return transaction associated to this context, {@literal null} if there is none
      */
     @Nullable
-    Transaction getTransaction();
+    public final Transaction getTransaction() {
+        AbstractSpan<?> contextSpan = getSpan();
+        return contextSpan != null ? contextSpan.getParentTransaction() : null;
+    }
+
+    @Nullable
+    @Override
+    public co.elastic.apm.agent.impl.transaction.Span createSpan() {
+        AbstractSpan<?> contextSpan = getSpan();
+        return contextSpan != null ? contextSpan.createSpan() : null;
+    }
+
+    @Nullable
+    @Override
+    public co.elastic.apm.agent.impl.transaction.Span createExitSpan() {
+        AbstractSpan<?> contextSpan = getSpan();
+        return contextSpan != null ? contextSpan.createExitSpan() : null;
+    }
+
+    public boolean isEmpty() {
+        return getSpan() == null && !containsBaggage();
+    }
+
+    //TODO: make abstract and implement correctly in subclasses
+    protected final boolean containsBaggage() {
+        return false;
+    }
+
+    @Override
+    public final <C> boolean propagateContext(C carrier, BinaryHeaderSetter<C> headerSetter) {
+        AbstractSpan<?> contextSpan = getSpan();
+        if (contextSpan != null) {
+            contextSpan.setNonDiscardable();
+            return contextSpan.getTraceContext().propagateTraceContext(carrier, headerSetter);
+        }
+        return false;
+    }
+
+    @Override
+    public final <C> void propagateContext(C carrier, TextHeaderSetter<C> headerSetter, @Nullable TextHeaderGetter<C> headerGetter) {
+        propagateContext(carrier, headerSetter, carrier, headerGetter);
+    }
+
+    @Override
+    public <C1, C2> void propagateContext(C1 carrier, TextHeaderSetter<C1> headerSetter, @Nullable C2 carrier2, @Nullable TextHeaderGetter<C2> headerGetter) {
+        AbstractSpan<?> contextSpan = getSpan();
+        if (contextSpan != null) {
+            if (headerGetter == null || carrier2 == null || !HeaderUtils.containsAny(TraceContext.TRACE_TEXTUAL_HEADERS, carrier2, headerGetter)) {
+                contextSpan.setNonDiscardable();
+                contextSpan.getTraceContext().propagateTraceContext(carrier, headerSetter);
+            }
+        }
+    }
+
+    @Override
+    public <C> boolean isPropagationRequired(C carrier, TextHeaderGetter<C> headerGetter) {
+        AbstractSpan<?> contextSpan = getSpan();
+        return contextSpan != null && !HeaderUtils.containsAny(TraceContext.TRACE_TEXTUAL_HEADERS, carrier, headerGetter);
+    }
 
+    @Override
+    public final boolean shouldSkipChildSpanCreation() {
+        Transaction contextTransaction = getTransaction();
+        return contextTransaction == null || contextTransaction.checkSkipChildSpanCreation();
+    }
 }