From 38b41dcd3edfc2a1c6a37bb57efa972bc2455c0f Mon Sep 17 00:00:00 2001
From: amond <amond@amond.net>
Date: Tue, 17 Jan 2023 03:01:45 +0900
Subject: [PATCH] GH-987 feat: propagate aws x-ray tracing header

Resolves #987
Resolves #988
---
 .../adapter/aws/CustomRuntimeEventLoop.java       | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/CustomRuntimeEventLoop.java b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/CustomRuntimeEventLoop.java
index d67a2edea..33c43d1d3 100644
--- a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/CustomRuntimeEventLoop.java
+++ b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/CustomRuntimeEventLoop.java
@@ -146,6 +146,21 @@ public final class CustomRuntimeEventLoop implements SmartLifecycle {
 				String invocationUrl = MessageFormat
 						.format(LAMBDA_INVOCATION_URL_TEMPLATE, runtimeApi, LAMBDA_VERSION_DATE, requestId);
 
+				String traceId = response.getHeaders().getFirst("Lambda-Runtime-Trace-Id");
+				if (traceId != null) {
+					if (logger.isDebugEnabled()) {
+						logger.debug("Lambda-Runtime-Trace-Id: " + traceId);
+					}
+					try {
+						// The X-Ray SDK uses this value to connect trace data between services.
+						System.setProperty("com.amazonaws.xray.traceHeader", traceId);
+					} catch (Exception e) {
+						if (logger.isDebugEnabled()) {
+							logger.debug("Failed to set amazon x-ray trace id", e);
+						}
+					}
+				}
+
 				try {
 					Message<byte[]> responseMessage = (Message<byte[]>) function.apply(eventMessage);