Preserve AWS X-Ray Context In Spring Cloud Function Custom Runtime

[Richardmbs12]

Issue Proposal: Preserve AWS X-Ray Context In Spring Cloud Function Custom Runtime

Summary

When running Spring Cloud Function on AWS Lambda’s Custom Runtime (provided.al2023) with GraalVM native images, downstream Micrometer instrumentation receives a new trace instead of the AWS-generated root because the framework’s CustomRuntimeEventLoop neither forwards nor surfaces the X-Ray headers that Lambda injects (Lambda-Runtime-Trace-Id, _X_AMZN_TRACE_ID). As a result, Micrometer spans lose parental lineage, causing Application Signals / X-Ray views to fragment.

Environment

Component	Value
Spring Cloud Function	4.3.0
Micrometer	1.12.5 (via Spring Boot 3.5.3)
Runtime	AWS Lambda Custom Runtime (provided.al2023), GraalVM 21 native image
Deployment	Native ZIP with ADOT Application Signals collector

Issue is still relevant to latest code on main branch here on spring-cloud-function

Reproduction

1. Deploy any Spring Cloud Function application as a GraalVM native binary using the Custom Runtime end-to-end example.
2. Invoke through API Gateway (HTTP or REST). Lambda delivers Lambda-Runtime-Trace-Id and _X_AMZN_TRACE_ID, while API Gateway’s X-Amzn-Trace-Id often lacks Parent=.
3. Observe Micrometer spans (debug exporter / logs / X-Ray). The spans mint a brand-new trace ID; parentSpanContext.remote is false.

Actual outcome

• System.getProperty("com.amazonaws.xray.traceHeader") remains empty.
• Message delivered to user function lacks Lambda-Runtime-Trace-Id / _X_AMZN_TRACE_ID.
• Micrometer propagation therefore starts a new trace and breaks lineage with AWS-managed segments.

Expected outcome

• Spring Cloud Function should either forward or synthesise the complete AWS trace header so Micrometer (or any tracer) can extract the remote parent context.
• com.amazonaws.xray.traceHeader should be populated (maintaining parity with the Java managed runtime).

Workaround & Evidence

We forked the event loop to:

1. Prefer the Lambda-Runtime-Trace-Id header (always contains Parent= for custom runtimes).
2. Fallback to _X_AMZN_TRACE_ID or the API-supplied X-Amzn-Trace-Id.
3. Update System.setProperty("com.amazonaws.xray.traceHeader", …) and attach headers to the Spring Message.

After patching we confirmed:

• Direct invoke: trace 1-68e77c02-abc9d4df78bf4c77bdb7a030.
• HTTP API: trace 1-68e77c14-3c721fbc6c8ddd025f8ef5d6 shows Micrometer spans under AWS root.
• REST API: trace 1-68e77c31-2f825ff737168a5a229da70a exhibits the same continuity.

Proposed Fix (for Spring Cloud Function)

Apply a small enrichment step inside CustomRuntimeEventLoop so downstream tracers see the AWS context. Suggested patch (abridged for clarity):

diff --git a/org/springframework/cloud/function/adapter/aws/CustomRuntimeEventLoop.java b/.../CustomRuntimeEventLoop.java
@@
-import org.springframework.messaging.Message;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.support.MessageBuilder;
@@
-					Message<?> requestMessage = AWSLambdaUtils.generateMessage(is, function.getInputType(), function.isSupplier(), mapper, clientContext);
+					Message<?> requestMessage = AWSLambdaUtils.generateMessage(is, function.getInputType(), function.isSupplier(), mapper, clientContext);
+					requestMessage = enrichTraceHeaders(response.getHeaders(), requestMessage);

-					Object functionResponse = function.apply(requestMessage);
+					Object functionResponse = function.apply(requestMessage);
@@
 	}
+
+	private Message<?> enrichTraceHeaders(HttpHeaders headers, Message<?> message) {
+		String runtimeTrace = trim(headers.getFirst("Lambda-Runtime-Trace-Id"));
+		String envTrace = trim(System.getenv("_X_AMZN_TRACE_ID"));
+		String headerTrace = trim(headers.getFirst("X-Amzn-Trace-Id"));
+
+		// prefer Lambda runtime header, then environment, then inbound header
+		String resolved = runtimeTrace != null ? runtimeTrace
+				: envTrace != null ? envTrace
+				: headerTrace;
+
+		if (resolved != null) {
+			System.setProperty("com.amazonaws.xray.traceHeader", resolved);
+		}
+		else {
+			System.clearProperty("com.amazonaws.xray.traceHeader");
+			return message;
+		}
+
+		return MessageBuilder.fromMessage(message)
+				.setHeader("Lambda-Runtime-Trace-Id", runtimeTrace != null ? runtimeTrace : resolved)
+				.setHeader("X-Amzn-Trace-Id", resolved)
+				.setHeader("_X_AMZN_TRACE_ID", envTrace != null ? envTrace : resolved)
+				.build();
+	}
+
+	private String trim(String value) {
+		return (value == null || value.isBlank()) ? null : value.trim();
+	}

Notes

• Lambda-Runtime-Trace-Id is the only header that consistently includes a Parent segment for custom runtimes; API Gateway’s header often omits it. Preferring the runtime header restores the link to the client span.
• Clearing the system property when the trace is absent avoids leaking previous values across invocations.
• If desired, the helper can be conditioned on the presence of Micrometer or made a dedicated extension point; the core requirement is to surface the runtime trace data before user code executes.

Request

Please integrate similar enrichment into CustomRuntimeEventLoop (and/or expose hooks for developers) so Spring Cloud Function preserves AWS trace lineage on custom runtimes without requiring custom event loops. This keeps Micrometer instrumentation aligned with Lambda-managed traces and avoids duplicating framework code.