Minor fixes for APM tracing (#138587)

* If running a recent version of APM server, it's not necessary to track an unsampled root span (aka transaction) to report a correct duration. APM agent only reports these for servers < v8.

* Cleanup parent trace headers if not tracing (recording), to avoid any further unnecessary attempts to create spans that are immediately discarded. This might be due to sampling or exceeding max transaction spans.

Relates to ES-13389

Author: mosche

build-tools/src/main/java/org/elasticsearch/gradle/testclusters/MockApmServer.java

@@ -114,8 +114,26 @@ public void stop() {
     }
 
     class RootHandler implements HttpHandler {
+        // checked by APM agent to identify the APM server version to adjust its behavior accordingly
+        private static final String FAKE_VERSION = """
+            {
+              "build_date": "2021-12-18T19:59:06Z",
+              "build_sha": "24fe620eeff5a19e2133c940c7e5ce1ceddb1445",
+              "publish_ready": true,
+              "version": "9.0.0"
+            }
+            """;
+
         public void handle(HttpExchange t) {
             try {
+                if ("GET".equals(t.getRequestMethod()) && "/".equals(t.getRequestURI().getPath())) {
+                    t.sendResponseHeaders(200, FAKE_VERSION.length());
+                    try (OutputStream os = t.getResponseBody()) {
+                        os.write(FAKE_VERSION.getBytes());
+                    }
+                    return;
+                }
+
                 InputStream body = t.getRequestBody();
                 if (metricFilter == null && transactionFilter == null) {
                     logRequestBody(body);

modules/apm/src/main/java/org/elasticsearch/telemetry/apm/internal/tracing/APMTracer.java

@@ -192,10 +192,14 @@ public void startTrace(TraceContext traceContext, Traceable traceable, String sp
             }
 
             final Span span = spanBuilder.startSpan();
-            // If not a root span (meaning a local parent exists) and the agent decided not to record the span, discard it immediately.
-            // Root spans (transactions), however, have to be kept to correctly report their duration.
-            if (localParentContext != null && span.isRecording() == false) {
-                logger.trace("Span [{}] [{}] will not be recorded due to transaction_max_spans reached", spanId, spanName);
+            if (span.isRecording() == false) {
+                if (localParentContext == null) {
+                    // this root span (transactions) is dropped due to sampling; the agent might report these when connected to
+                    // very old versions of apm server, however (with an incorrect duration)
+                    logger.trace("Root span [{}] [{}] will not be recorded due to sampling", spanId, spanName);
+                } else {
+                    logger.trace("Span [{}] [{}] will not be recorded due to transaction_max_spans reached", spanId, spanName);
+                }
                 span.end(); // end span immediately to release any resources.
                 return null; // return null to discard and not record in map of spans
             }

modules/apm/src/test/java/org/elasticsearch/telemetry/apm/internal/tracing/APMTracerTests.java

@@ -109,9 +109,7 @@ public void test_onTraceStarted_ifNotRecorded_doesNotStartTracing() {
         apmTracer.startTrace(traceContext, TRACEABLE1, "name1_discard", null);
 
         assertThat(traceContext.getTransient(Task.APM_TRACE_CONTEXT), nullValue());
-        // the root span (transaction) is tracked
-        assertThat(apmTracer.getSpans(), aMapWithSize(1));
-        assertThat(apmTracer.getSpans(), hasKey(TRACEABLE1.getSpanId()));
+        assertThat(apmTracer.getSpans(), anEmptyMap());
     }
 
     /**

server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java

@@ -205,16 +205,41 @@ public StoredContext newTraceContext() {
         // this is the context when this method returns
         final ThreadContextStruct newContext;
 
-        final boolean hasTraceHeaders = originalContext.requestHeaders.containsKey(Task.TRACE_PARENT_HTTP_HEADER)
-            || originalContext.requestHeaders.containsKey(Task.TRACE_STATE)
-            || originalContext.transientHeaders.containsKey(Task.APM_TRACE_CONTEXT);
+        final var requestHeaders = originalContext.requestHeaders;
+        final var transientHeaders = originalContext.transientHeaders;
+
+        final boolean hasTraceHeaders = requestHeaders.containsKey(Task.TRACE_PARENT_HTTP_HEADER)
+            || requestHeaders.containsKey(Task.TRACE_STATE)
+            || transientHeaders.containsKey(Task.APM_TRACE_CONTEXT);
 
         if (hasTraceHeaders == false) {
-            // no need to copy if no trace headers are present
-            newContext = originalContext;
+            final boolean hasParentTraceHeaders = (transientHeaders.containsKey(Task.PARENT_TRACE_PARENT_HEADER)
+                || transientHeaders.containsKey(Task.PARENT_TRACE_STATE)
+                || transientHeaders.containsKey(Task.PARENT_APM_TRACE_CONTEXT));
+
+            if (hasParentTraceHeaders == false) {
+                // no need to copy if no trace headers are present
+                newContext = originalContext;
+            } else {
+                // tracing was stopped (e.g. after reaching max spans);
+                // remove parent trace headers to not attempt creating any further, immediately discarded spans
+                final Map<String, Object> newTransientHeaders = new HashMap<>(transientHeaders);
+                newTransientHeaders.remove(Task.PARENT_TRACE_PARENT_HEADER);
+                newTransientHeaders.remove(Task.PARENT_TRACE_STATE);
+                newTransientHeaders.remove(Task.PARENT_APM_TRACE_CONTEXT);
+
+                newContext = new ThreadContextStruct(
+                    requestHeaders,
+                    originalContext.responseHeaders,
+                    newTransientHeaders,
+                    originalContext.isSystemContext,
+                    originalContext.warningHeadersSize
+                );
+                threadLocal.set(newContext);
+            }
         } else {
-            final Map<String, String> newRequestHeaders = new HashMap<>(originalContext.requestHeaders);
-            final Map<String, Object> newTransientHeaders = new HashMap<>(originalContext.transientHeaders);
+            final Map<String, String> newRequestHeaders = new HashMap<>(requestHeaders);
+            final Map<String, Object> newTransientHeaders = new HashMap<>(transientHeaders);
 
             final String previousTraceParent = newRequestHeaders.remove(Task.TRACE_PARENT_HTTP_HEADER);
             if (previousTraceParent != null) {

server/src/test/java/org/elasticsearch/tasks/TaskManagerTests.java

@@ -287,8 +287,8 @@ public void testRegisterTaskStartsTracingIfTraceParentExists() {
         final Tracer mockTracer = mock(Tracer.class);
         final TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Set.of(), mockTracer);
 
-        // fake a parent APM trace context
-        threadPool.getThreadContext().putTransient(Task.PARENT_APM_TRACE_CONTEXT, null);
+        // fake an APM trace context
+        threadPool.getThreadContext().putTransient(Task.APM_TRACE_CONTEXT, new Object());
         final boolean hasParentTask = randomBoolean();
         final TaskId parentTask = hasParentTask ? new TaskId("parentNode", 1) : TaskId.EMPTY_TASK_ID;
 
@@ -408,8 +408,8 @@ public void testRegisterAndExecuteStartsTracingIfTraceParentExists() {
         final Tracer mockTracer = mock(Tracer.class);
         final TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Set.of(), mockTracer);
 
-        // fake a parent APM trace context
-        threadPool.getThreadContext().putTransient(Task.PARENT_APM_TRACE_CONTEXT, null);
+        // fake an APM trace context
+        threadPool.getThreadContext().putTransient(Task.APM_TRACE_CONTEXT, new Object());
 
         final Task task = taskManager.registerAndExecute(
             "testType",