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

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.