Add API methods to enable RUM trace correlation (#278)

Author: felixbarny

apm-agent-api/src/main/java/co/elastic/apm/api/NoopSpan.java

@@ -55,6 +55,12 @@ public String getId() {
         return "";
     }
 
+    @Nonnull
+    @Override
+    public String getTraceId() {
+        return "";
+    }
+
     @Override
     public Scope activate() {
         return NoopScope.INSTANCE;

apm-agent-api/src/main/java/co/elastic/apm/api/NoopTransaction.java

@@ -61,6 +61,18 @@ public String getId() {
         return "";
     }
 
+    @Nonnull
+    @Override
+    public String ensureParentId() {
+        return "";
+    }
+
+    @Nonnull
+    @Override
+    public String getTraceId() {
+        return "";
+    }
+
     @Override
     public Scope activate() {
         return NoopScope.INSTANCE;

apm-agent-api/src/main/java/co/elastic/apm/api/Span.java

@@ -132,6 +132,22 @@ public interface Span {
     @Nonnull
     String getId();
 
+    /**
+     * Returns the id of this trace (never {@code null})
+     * <p>
+     * The trace-ID is consistent across all transactions and spans which belong to the same logical trace,
+     * even for spans which happened in another service (given this service is also monitored by Elastic APM).
+     * </p>
+     * <p>
+     * If this span represents a noop,
+     * this method returns an empty string.
+     * </p>
+     *
+     * @return the id of this span (never {@code null})
+     */
+    @Nonnull
+    String getTraceId();
+
     /**
      * Makes this span the active span on the current thread until {@link Scope#close()} has been called.
      * <p>

apm-agent-api/src/main/java/co/elastic/apm/api/SpanImpl.java

@@ -82,6 +82,13 @@ public String getId() {
         return "";
     }
 
+    @Nonnull
+    @Override
+    public String getTraceId() {
+        // co.elastic.apm.plugin.api.SpanInstrumentation.GetTraceIdInstrumentation
+        return "";
+    }
+
     @Override
     public Scope activate() {
         // co.elastic.apm.plugin.api.SpanInstrumentation.ActivateInstrumentation

apm-agent-api/src/main/java/co/elastic/apm/api/Transaction.java

@@ -117,6 +117,50 @@ public interface Transaction extends Span {
     @Nonnull
     String getId();
 
+    /**
+     * <p>
+     * If the transaction does not have a parent-ID yet,
+     * calling this method generates a new ID,
+     * sets it as the parent-ID of this transaction,
+     * and returns it as a `String`.
+     * </p>
+     * <p>
+     * This enables the correlation of the spans the JavaScript Real User Monitoring (RUM) agent creates for the initial page load
+     * with the transaction of the backend service.
+     * If your backend service generates the HTML page dynamically,
+     * initializing the JavaScript RUM agent with the value of this method allows analyzing the time spent in the browser vs in the backend services.
+     * </p>
+     * <p>
+     * To enable the JavaScript RUM agent when using an HTML templating language like Freemarker,
+     * add {@code ElasticApm.currentTransaction()} with the key {@code "transaction"} to the model.
+     * </p>
+     * <p>
+     * Also, add a snippet similar to this to the body of your HTML pages,
+     * preferably before other JS libraries:
+     * </p>
+     *
+     * <pre>{@code
+     * <script src="elastic-apm-js-base/dist/bundles/elastic-apm-js-base.umd.min.js"></script>
+     * <script>
+     *   var elasticApm = initApm({
+     *     serviceName: '',
+     *     serverUrl: 'http://localhost:8200',
+     *     pageLoadTraceId: "${transaction.traceId}",
+     *     pageLoadSpanId: "${transaction.ensureParentId()}",
+     *     pageLoadSampled: ${transaction.sampled}
+     *   })
+     * </script>
+     * }</pre>
+     *
+     * <p>
+     * See the JavaScript RUM agent documentation for more information.
+     * </p>
+     *
+     * @return the parent-ID for this transaction. Updates the transaction to use a new parent-ID if it has previously been unset.
+     */
+    @Nonnull
+    String ensureParentId();
+
     /**
      * Makes this transaction the active transaction on the current thread until {@link Scope#close()} has been called.
      * <p>

apm-agent-api/src/main/java/co/elastic/apm/api/TransactionImpl.java

@@ -40,4 +40,11 @@ public void setUser(String id, String email, String username) {
         // co.elastic.apm.plugin.api.TransactionInstrumentation$SetUserInstrumentation.setUser
     }
 
+    @Nonnull
+    @Override
+    public String ensureParentId() {
+        // co.elastic.apm.plugin.api.TransactionInstrumentation.EnsureParentIdInstrumentation
+        return "";
+    }
+
 }

apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/plugin/api/SpanInstrumentation.java

@@ -146,6 +146,21 @@ public static void getId(@Advice.FieldValue(value = "span", typing = Assigner.Ty
         }
     }
 
+    public static class GetTraceIdInstrumentation extends SpanInstrumentation {
+        public GetTraceIdInstrumentation() {
+            super(named("getTraceId").and(takesArguments(0)));
+        }
+
+        @VisibleForAdvice
+        @Advice.OnMethodExit
+        public static void getTraceId(@Advice.FieldValue(value = "span", typing = Assigner.Typing.DYNAMIC) AbstractSpan<?> span,
+                                        @Advice.Return(readOnly = false) String traceId) {
+            if (tracer != null) {
+                traceId = span.getTraceContext().getTraceId().toString();
+            }
+        }
+    }
+
     public static class AddTagInstrumentation extends SpanInstrumentation {
         public AddTagInstrumentation() {
             super(named("addTag"));

apm-agent-plugins/apm-api-plugin/src/main/java/co/elastic/apm/plugin/api/TransactionInstrumentation.java

@@ -21,6 +21,7 @@
 
 import co.elastic.apm.bci.ElasticApmInstrumentation;
 import co.elastic.apm.bci.VisibleForAdvice;
+import co.elastic.apm.impl.transaction.TraceContext;
 import co.elastic.apm.impl.transaction.Transaction;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
@@ -77,4 +78,23 @@ public static void setUser(@Advice.FieldValue(value = "span", typing = Assigner.
             transaction.setUser(id, email, username);
         }
     }
+
+    public static class EnsureParentIdInstrumentation extends TransactionInstrumentation {
+        public EnsureParentIdInstrumentation() {
+            super(named("ensureParentId"));
+        }
+
+        @VisibleForAdvice
+        @Advice.OnMethodExit
+        public static void ensureParentId(@Advice.FieldValue(value = "span", typing = Assigner.Typing.DYNAMIC) Transaction transaction,
+                                          @Advice.Return(readOnly = false) String spanId) {
+            if (tracer != null) {
+                final TraceContext traceContext = transaction.getTraceContext();
+                if (traceContext.getParentId().isEmpty()) {
+                    traceContext.getParentId().setToRandomValue();
+                }
+                spanId = traceContext.getParentId().toString();
+            }
+        }
+    }
 }

apm-agent-plugins/apm-api-plugin/src/main/resources/META-INF/services/co.elastic.apm.bci.ElasticApmInstrumentation

@@ -3,12 +3,14 @@ co.elastic.apm.plugin.api.ElasticApmApiInstrumentation$CurrentTransactionInstrum
 co.elastic.apm.plugin.api.ElasticApmApiInstrumentation$CurrentSpanInstrumentation
 co.elastic.apm.plugin.api.ElasticApmApiInstrumentation$CaptureExceptionInstrumentation
 co.elastic.apm.plugin.api.TransactionInstrumentation$SetUserInstrumentation
+co.elastic.apm.plugin.api.TransactionInstrumentation$EnsureParentIdInstrumentation
 co.elastic.apm.plugin.api.SpanInstrumentation$SetNameInstrumentation
 co.elastic.apm.plugin.api.SpanInstrumentation$SetTypeInstrumentation
 co.elastic.apm.plugin.api.SpanInstrumentation$DoCreateSpanInstrumentation
 co.elastic.apm.plugin.api.SpanInstrumentation$EndInstrumentation
 co.elastic.apm.plugin.api.SpanInstrumentation$CaptureExceptionInstrumentation
 co.elastic.apm.plugin.api.SpanInstrumentation$GetIdInstrumentation
+co.elastic.apm.plugin.api.SpanInstrumentation$GetTraceIdInstrumentation
 co.elastic.apm.plugin.api.SpanInstrumentation$AddTagInstrumentation
 co.elastic.apm.plugin.api.SpanInstrumentation$ActivateInstrumentation
 co.elastic.apm.plugin.api.CaptureExceptionInstrumentation

apm-agent-plugins/apm-api-plugin/src/test/java/co/elastic/apm/api/ElasticApmApiInstrumentationTest.java

@@ -108,10 +108,13 @@ void testGetId_distributedTracingEnabled() {
         co.elastic.apm.impl.transaction.Transaction transaction = tracer.startTransaction().withType(Transaction.TYPE_REQUEST);
         try (Scope scope = transaction.activateInScope()) {
             assertThat(ElasticApm.currentTransaction().getId()).isEqualTo(transaction.getTraceContext().getId().toString());
+            assertThat(ElasticApm.currentTransaction().getTraceId()).isEqualTo(transaction.getTraceContext().getTraceId().toString());
             assertThat(ElasticApm.currentSpan().getId()).isEqualTo(transaction.getTraceContext().getId().toString());
+            assertThat(ElasticApm.currentSpan().getTraceId()).isEqualTo(transaction.getTraceContext().getTraceId().toString());
             co.elastic.apm.impl.transaction.Span span = transaction.createSpan().withType("db").withName("SELECT");
             try (Scope spanScope = span.activateInScope()) {
                 assertThat(ElasticApm.currentSpan().getId()).isEqualTo(span.getTraceContext().getId().toString());
+                assertThat(ElasticApm.currentSpan().getTraceId()).isEqualTo(span.getTraceContext().getTraceId().toString());
             } finally {
                 span.end();
             }
@@ -162,4 +165,16 @@ void testScopes() {
         assertThat(ElasticApm.currentTransaction()).isSameAs(NoopTransaction.INSTANCE);
 
     }
+
+    @Test
+    void testEnsureParentId() {
+        final Transaction transaction = ElasticApm.startTransaction();
+        try (co.elastic.apm.api.Scope scope = transaction.activate()) {
+            assertThat(tracer.currentTransaction()).isNotNull();
+            assertThat(tracer.currentTransaction().getTraceContext().getParentId().isEmpty()).isTrue();
+            String rumTransactionId = transaction.ensureParentId();
+            assertThat(tracer.currentTransaction().getTraceContext().getParentId().toString()).isEqualTo(rumTransactionId);
+            assertThat(transaction.ensureParentId()).isEqualTo(rumTransactionId);
+        }
+    }
 }