Add span links to messaging system tracing (#2610)

Author: eyalkoren

CHANGELOG.asciidoc

@@ -39,6 +39,7 @@ the agent will now include the database name in the dependency, thus `mysql/my-d
 * Added support for downloading the latest agent version through the attach CLI by setting `--download-agent-version latest`. In
 addition, when using the `apm-agent-attach-cli-slim.jar`, which does not contain a bundled agent, the latest version will be downloaded
 from maven at runtime unless configured otherwise through  `--download-agent-version` - {pull}2659[#2659]
+* Added span-links to messaging systems instrumentation (supported by APM Server 8.3+ only) - {pull}2610[#2610]
 
 [float]
 ===== Bug fixes

apm-agent-core/src/main/java/co/elastic/apm/agent/configuration/MessagingConfiguration.java

@@ -30,8 +30,9 @@
 public class MessagingConfiguration extends ConfigurationOptionProvider {
     private static final String MESSAGING_CATEGORY = "Messaging";
     private static final String MESSAGE_POLLING_TRANSACTION_STRATEGY = "message_polling_transaction_strategy";
+    private static final String MESSAGE_BATCH_STRATEGY = "message_batch_strategy";
 
-    private ConfigurationOption<Strategy> messagePollingTransactionStrategy = ConfigurationOption.enumOption(Strategy.class)
+    private ConfigurationOption<JmsStrategy> messagePollingTransactionStrategy = ConfigurationOption.enumOption(JmsStrategy.class)
         .key(MESSAGE_POLLING_TRANSACTION_STRATEGY)
         .configurationCategory(MESSAGING_CATEGORY)
         .tags("internal")
@@ -41,7 +42,19 @@ public class MessagingConfiguration extends ConfigurationOptionProvider {
             "\n" +
             "This option is case-insensitive and is only relevant for JMS.")
         .dynamic(true)
-        .buildWithDefault(Strategy.HANDLING);
+        .buildWithDefault(JmsStrategy.HANDLING);
+
+    private ConfigurationOption<BatchStrategy> messageBatchStrategy = ConfigurationOption.enumOption(BatchStrategy.class)
+        .key(MESSAGE_BATCH_STRATEGY)
+        .configurationCategory(MESSAGING_CATEGORY)
+        .tags("internal")
+        .description("Determines whether Spring messaging system libraries should create a batch for the processing of the entire \n" +
+            "message/record batch, or one transaction for each message/record processing, typically by wrapping the message batch data \n" +
+            "structure. Valid options are `SINGLE_HANDLING` and `BATCH_HANDLING`. \n" +
+            "\n" +
+            "This option is case-insensitive and is only relevant for Spring messaging system libraries that support batch processing.")
+        .dynamic(true)
+        .buildWithDefault(BatchStrategy.BATCH_HANDLING);
 
     private ConfigurationOption<Boolean> collectQueueAddress = ConfigurationOption.booleanOption()
         .key("collect_queue_address")
@@ -79,10 +92,14 @@ public class MessagingConfiguration extends ConfigurationOptionProvider {
         .dynamic(true)
         .buildWithDefault(Boolean.TRUE);
 
-    public MessagingConfiguration.Strategy getMessagePollingTransactionStrategy() {
+    public JmsStrategy getMessagePollingTransactionStrategy() {
         return messagePollingTransactionStrategy.get();
     }
 
+    public BatchStrategy getMessageBatchStrategy() {
+        return messageBatchStrategy.get();
+    }
+
     public List<WildcardMatcher> getIgnoreMessageQueues() {
         return ignoreMessageQueues.get();
     }
@@ -95,9 +112,35 @@ public boolean shouldEndMessagingTransactionOnPoll() {
         return endMessagingTransactionOnPoll.get();
     }
 
-    public enum Strategy {
+    public enum JmsStrategy {
+        /**
+         * Create a transaction capturing JMS {@code receive} invocations
+         */
         POLLING,
+        /**
+         * Use heuristics to create a transaction that captures the JMS message handling execution. This strategy requires heuristics
+         * when JMS {@code receive} APIs are used (rather than {@code onMessage}), as there is no API representing message handling start
+         * and end. Even though this is riskier and less deterministic, it is the default JMS tracing strategy otherwise all
+         * "interesting" subsequent events that follow message receive will be missed because there will be no active transaction.
+         */
         HANDLING,
+        /**
+         * Create a transaction both for the polling ({@code receive}) action AND the subsequent message handling.
+         */
         BOTH
     }
+
+    /**
+     * Only relevant for Spring wrappers around supported messaging clients, such as AMQP.
+     */
+    public enum BatchStrategy {
+        /**
+         * Create a transaction for each received message/record, typically by wrapping the message batch data structure
+         */
+        SINGLE_HANDLING,
+        /**
+         * Create a single transaction encapsulating the entire message/record batch-processing.
+         */
+        BATCH_HANDLING
+    }
 }

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

@@ -84,6 +84,7 @@ public class ElasticApmTracer implements Tracer {
     private final ObjectPool<Transaction> transactionPool;
     private final ObjectPool<Span> spanPool;
     private final ObjectPool<ErrorCapture> errorPool;
+    private final ObjectPool<TraceContext> spanLinkPool;
     private final Reporter reporter;
     private final ObjectPoolFactory objectPoolFactory;
     // Maintains a stack of all the activated spans/contexts
@@ -145,6 +146,9 @@ public void onChange(ConfigurationOption<?> configurationOption, Boolean oldValu
         // we are assuming that we don't need as many errors as spans or transactions
         errorPool = poolFactory.createErrorPool(maxPooledElements / 2, this);
 
+        // span links pool allows for 10X the maximum allowed span links per span
+        spanLinkPool = poolFactory.createSpanLinkPool(AbstractSpan.MAX_ALLOWED_SPAN_LINKS * 10, this);
+
         sampler = ProbabilitySampler.of(coreConfiguration.getSampleRate().get());
         coreConfiguration.getSampleRate().addChangeListener(new ConfigurationOption.ChangeListener<Double>() {
             @Override
@@ -447,6 +451,10 @@ public void endError(ErrorCapture error) {
         reporter.report(error);
     }
 
+    public TraceContext createSpanLink() {
+        return spanLinkPool.createInstance();
+    }
+
     public void recycle(Transaction transaction) {
         transactionPool.recycle(transaction);
     }
@@ -459,6 +467,10 @@ public void recycle(ErrorCapture error) {
         errorPool.recycle(error);
     }
 
+    public void recycle(TraceContext traceContext) {
+        spanLinkPool.recycle(traceContext);
+    }
+
     public synchronized void stop() {
         if (tracerState == TracerState.STOPPED) {
             // may happen if explicitly stopped in a unit test and executed again within a shutdown hook

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

@@ -32,6 +32,7 @@
 
 import javax.annotation.Nullable;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -46,6 +47,7 @@ public abstract class AbstractSpan<T extends AbstractSpan<T>> implements Recycla
     public static final int PRIO_DEFAULT = 0;
     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);
 
     protected static final double MS_IN_MICROS = TimeUnit.MILLISECONDS.toMicros(1);
     protected final TraceContext traceContext;
@@ -117,6 +119,10 @@ public abstract class AbstractSpan<T extends AbstractSpan<T>> implements Recycla
 
     protected final AtomicReference<Span> bufferedSpan = new AtomicReference<>();
 
+    // Span links handling
+    public static final int MAX_ALLOWED_SPAN_LINKS = 1000;
+    private final List<TraceContext> spanLinks = new UniqueSpanLinkArrayList();
+
     @Nullable
     private OTelSpanKind otelKind = null;
 
@@ -369,6 +375,50 @@ public TraceContext getTraceContext() {
         return traceContext;
     }
 
+    /**
+     * Adds a span link based on the tracecontext header retrieved from the provided {@code carrier} through the provided {@code
+     * headerGetter}.
+     * @param childContextCreator the proper tracecontext inference implementation, corresponding on the header and types
+     * @param headerGetter the proper header extractor, corresponding the header and carrier types
+     * @param carrier the object from which the tracecontext header is to be retrieved
+     * @param <H> the tracecontext header type - either binary ({@code byte[]}) or textual ({@code String})
+     * @param <C> the tracecontext header carrier type, e.g. Kafka record or JMS message
+     * @return {@code true} if added, {@code false} otherwise
+     */
+    public <H, C> boolean addSpanLink(TraceContext.ChildContextCreatorTwoArg<C, HeaderGetter<H, C>> childContextCreator,
+                                    HeaderGetter<H, C> headerGetter, @Nullable C carrier) {
+        if (spanLinks.size() == MAX_ALLOWED_SPAN_LINKS) {
+            oneTimeMaxSpanLinksLogger.warn("Span links for {} has reached the allowed maximum ({}). No more spans will be linked.",
+                this, MAX_ALLOWED_SPAN_LINKS);
+            return false;
+        }
+        boolean added = false;
+        try {
+            TraceContext childTraceContext = tracer.createSpanLink();
+            if (childContextCreator.asChildOf(childTraceContext, carrier, headerGetter)) {
+                added = spanLinks.add(childTraceContext);
+            }
+            if (!added) {
+                tracer.recycle(childTraceContext);
+            }
+        } catch (Exception e) {
+            logger.error(String.format("Failed to add span link to %s from header carrier %s and %s", this, carrier,
+                headerGetter.getClass().getName()), e);
+        }
+        return added;
+    }
+
+    /**
+     * Returns a list of links from this span to other spans in the format of child {@link TraceContext}s, of which parent is the linked
+     * span. For each entry in the returned list, the linked span's {@code traceId} can be retrieved through
+     * {@link TraceContext#getTraceId()} and the {@code spanId} can be retrieved through {@link TraceContext#getParentId()}.
+     *
+     * @return a list of child {@link TraceContext}s of linked spans
+     */
+    public List<TraceContext> getSpanLinks() {
+        return spanLinks;
+    }
+
     @Override
     public void resetState() {
         finished = true;
@@ -388,10 +438,19 @@ public void resetState() {
         userOutcome = null;
         hasCapturedExceptions = false;
         bufferedSpan.set(null);
+        recycleSpanLinks();
         otelKind = null;
         otelAttributes.clear();
     }
 
+    private void recycleSpanLinks() {
+        //noinspection ForLoopReplaceableByForEach
+        for (int i = 0; i < spanLinks.size(); i++) {
+            tracer.recycle(spanLinks.get(i));
+        }
+        spanLinks.clear();
+    }
+
     public Span createSpan() {
         return createSpan(traceContext.getClock().getEpochMicros());
     }

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

@@ -265,13 +265,13 @@ public static TraceContext with128BitId(ElasticApmTracer tracer) {
     }
 
     @SuppressWarnings("unchecked")
-    public static <C> ChildContextCreatorTwoArg<C, TextHeaderGetter<C>> getFromTraceContextTextHeaders() {
-        return (ChildContextCreatorTwoArg<C, TextHeaderGetter<C>>) FROM_TRACE_CONTEXT_TEXT_HEADERS;
+    public static <C> ChildContextCreatorTwoArg<C, HeaderGetter<String, C>> getFromTraceContextTextHeaders() {
+        return (ChildContextCreatorTwoArg<C, HeaderGetter<String, C>>) FROM_TRACE_CONTEXT_TEXT_HEADERS;
     }
 
     @SuppressWarnings("unchecked")
-    public static <C> ChildContextCreatorTwoArg<C, BinaryHeaderGetter<C>> getFromTraceContextBinaryHeaders() {
-        return (ChildContextCreatorTwoArg<C, BinaryHeaderGetter<C>>) FROM_TRACE_CONTEXT_BINARY_HEADERS;
+    public static <C> ChildContextCreatorTwoArg<C, HeaderGetter<byte[], C>> getFromTraceContextBinaryHeaders() {
+        return (ChildContextCreatorTwoArg<C, HeaderGetter<byte[], C>>) FROM_TRACE_CONTEXT_BINARY_HEADERS;
     }
 
     public static ChildContextCreator<Tracer> fromActive() {

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

@@ -0,0 +1,96 @@
+/*
+ * 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.transaction;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * A {@link TraceContext} list that enforces uniqueness for the purpose of span links storage.
+ * We cannot naturally use {@link Set} implementations because {@link TraceContext#equals} is comparing own ID with other's ID. This is
+ * inappropriate for span links, which we consider equal if their <b>parent IDs</b> are equal.
+ * So instead, we use a limited subclass of {@link ArrayList} that maintains a parent ID cache for equality checks.
+ * As a side benefit, this gives us the ability to iterate over span links based on index and avoid the related iterator allocation, which
+ * we wouldn't have if using a {@link Set}.
+ */
+public class UniqueSpanLinkArrayList extends ArrayList<TraceContext> {
+
+    private final Set<Id> parentIdSet = new HashSet<>();
+
+    @Override
+    public boolean retainAll(Collection<?> c) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public TraceContext set(int index, TraceContext traceContext) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean add(TraceContext traceContext) {
+        if (parentIdSet.add(traceContext.getParentId())) {
+            return super.add(traceContext);
+        }
+        return false;
+    }
+
+    @Override
+    public void add(int index, TraceContext traceContext) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public TraceContext remove(int index) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean remove(Object o) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void clear() {
+        parentIdSet.clear();
+        super.clear();
+    }
+
+    @Override
+    public boolean addAll(Collection<? extends TraceContext> c) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean addAll(int index, Collection<? extends TraceContext> c) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    protected void removeRange(int fromIndex, int toIndex) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean removeAll(Collection<?> c) {
+        throw new UnsupportedOperationException();
+    }
+}

apm-agent-core/src/main/java/co/elastic/apm/agent/objectpool/ObjectPoolFactory.java

@@ -21,6 +21,7 @@
 import co.elastic.apm.agent.impl.ElasticApmTracer;
 import co.elastic.apm.agent.impl.error.ErrorCapture;
 import co.elastic.apm.agent.impl.transaction.Span;
+import co.elastic.apm.agent.impl.transaction.TraceContext;
 import co.elastic.apm.agent.impl.transaction.Transaction;
 import co.elastic.apm.agent.objectpool.impl.QueueBasedObjectPool;
 import org.jctools.queues.atomic.AtomicQueueFactory;
@@ -59,4 +60,13 @@ public ErrorCapture createInstance() {
                 }
             });
     }
+
+    public ObjectPool<TraceContext> createSpanLinkPool(int maxCapacity, final ElasticApmTracer tracer) {
+        return createRecyclableObjectPool(maxCapacity, new Allocator<TraceContext>() {
+                @Override
+                public TraceContext createInstance() {
+                    return TraceContext.with64BitId(tracer);
+                }
+            });
+    }
 }