From 17c22b4e877b92e3e992fa2bac53edb7ee84e321 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 ca452db4c..c23c260c9 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
@@ -147,6 +147,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);