fix: propagate tracing header via system property (#1087)

close: #234
micronaut-projects · Jun 1, 2021 · 853f203 · 853f203
1 parent 510a796
commit 853f203
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 4 deletions.
diff --git a/...ntime/src/main/java/io/micronaut/function/aws/runtime/AbstractMicronautLambdaRuntime.java b/...ntime/src/main/java/io/micronaut/function/aws/runtime/AbstractMicronautLambdaRuntime.java
@@ -38,6 +38,7 @@
 import io.micronaut.core.util.ArrayUtils;
 import io.micronaut.core.util.StringUtils;
 import io.micronaut.function.aws.MicronautLambdaContext;
+import io.micronaut.function.aws.MicronautRequestHandler;
 import io.micronaut.http.HttpHeaders;
 import io.micronaut.http.HttpResponse;
 import io.micronaut.http.HttpStatus;
@@ -392,7 +393,7 @@ protected void propagateTraceId(HttpHeaders headers) {
         String traceId = headers.get(LambdaRuntimeInvocationResponseHeaders.LAMBDA_RUNTIME_TRACE_ID);
         logn(LogLevel.DEBUG, "Trace id: ", traceId, '\n');
         if (StringUtils.isNotEmpty(traceId)) {
-            //TODO Set Env.variable _X_AMZN_TRACE_ID with value traceId
+            System.setProperty(MicronautRequestHandler.LAMBDA_TRACE_HEADER_PROP, traceId);
         }
     }
 

diff --git a/...ntime/src/test/groovy/io/micronaut/function/aws/runtime/MicronautLambdaRuntimeSpec.groovy b/...ntime/src/test/groovy/io/micronaut/function/aws/runtime/MicronautLambdaRuntimeSpec.groovy
@@ -21,6 +21,8 @@ import com.amazonaws.serverless.proxy.model.AwsProxyRequestContext
 import com.amazonaws.serverless.proxy.model.AwsProxyResponse
 import io.micronaut.context.ApplicationContext
 import io.micronaut.context.annotation.Requires
+import io.micronaut.function.aws.MicronautRequestHandler
+import io.micronaut.http.HttpHeaders
 import io.micronaut.http.HttpResponse
 import io.micronaut.http.annotation.Body
 import io.micronaut.http.annotation.Controller
@@ -30,9 +32,35 @@ import io.micronaut.http.annotation.Post
 import io.micronaut.runtime.server.EmbeddedServer
 import spock.lang.Specification
 import spock.util.concurrent.PollingConditions
+import spock.util.environment.RestoreSystemProperties
 
 class MicronautLambdaRuntimeSpec extends Specification {
 
+    @RestoreSystemProperties
+    void "Tracing header propagated as system property"() {
+        given:
+        String traceHeader = 'Root=1-5759e988-bd862e3fe1be46a994272793;Sampled=1'
+        EmbeddedServer embeddedServer = ApplicationContext.run(EmbeddedServer, ['spec.name': 'MicronautLambdaRuntimeSpec'])
+        String serverUrl = "localhost:$embeddedServer.port"
+        CustomMicronautLambdaRuntime customMicronautLambdaRuntime = new CustomMicronautLambdaRuntime(serverUrl)
+        Thread t = new Thread({ ->
+            customMicronautLambdaRuntime.run([] as String[])
+        })
+        t.start()
+
+        when:
+        def httpHeaders = Stub(HttpHeaders) {
+            get(LambdaRuntimeInvocationResponseHeaders.LAMBDA_RUNTIME_TRACE_ID) >> traceHeader
+        }
+        customMicronautLambdaRuntime.propagateTraceId(httpHeaders)
+
+        then:
+        System.getProperty(MicronautRequestHandler.LAMBDA_TRACE_HEADER_PROP) == traceHeader
+
+        cleanup:
+        embeddedServer.close()
+    }
+
     void "test runtime API loop"() {
         given:
         EmbeddedServer embeddedServer = ApplicationContext.run(EmbeddedServer, ['spec.name': 'MicronautLambdaRuntimeSpec'])

diff --git a/function-aws/src/main/java/io/micronaut/function/aws/MicronautRequestHandler.java b/function-aws/src/main/java/io/micronaut/function/aws/MicronautRequestHandler.java
@@ -24,6 +24,7 @@
 import io.micronaut.core.convert.ConversionError;
 import io.micronaut.core.reflect.GenericTypeUtils;
 import io.micronaut.core.util.ArrayUtils;
+import io.micronaut.core.util.StringUtils;
 import io.micronaut.function.executor.AbstractFunctionExecutor;
 import org.slf4j.MDC;
 import java.util.Optional;
@@ -39,6 +40,10 @@
 public abstract class MicronautRequestHandler<I, O> extends AbstractFunctionExecutor<I, O, Context> implements RequestHandler<I, O>, MicronautLambdaContext {
 
     public static final String ENV_X_AMZN_TRACE_ID = "_X_AMZN_TRACE_ID";
+
+    // See: https://github.com/aws/aws-xray-sdk-java/issues/251
+    public static final String LAMBDA_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader";
+
     public static final String MDC_DEFAULT_AWS_REQUEST_ID = "AWSRequestId";
     public static final String MDC_DEFAULT_FUNCTION_NAME = "AWSFunctionName";
     public static final String MDC_DEFAULT_FUNCTION_VERSION = "AWSFunctionVersion";
@@ -130,9 +135,14 @@ protected void populateMappingDiagnosticContextWithXrayTraceId() {
      */
     @NonNull
     protected static Optional<String> parseXrayTraceId() {
-        final String X_AMZN_TRACE_ID = System.getenv(ENV_X_AMZN_TRACE_ID);
-        if (X_AMZN_TRACE_ID != null) {
-            return Optional.of(X_AMZN_TRACE_ID.split(";")[0].replace("Root=", ""));
+        String lambdaTraceHeaderKey = System.getenv(ENV_X_AMZN_TRACE_ID);
+        lambdaTraceHeaderKey = StringUtils.isNotEmpty(lambdaTraceHeaderKey) ? lambdaTraceHeaderKey
+                : System.getProperty(LAMBDA_TRACE_HEADER_PROP);
+        if (lambdaTraceHeaderKey != null) {
+            String[] arr = lambdaTraceHeaderKey.split(";");
+            if (arr.length >= 1) {
+                return Optional.of(arr[0].replace("Root=", ""));
+            }
         }
         return Optional.empty();
     }