From 89d7bc60bb17a6984b7fe0ec1e62a60cc3b0eee0 Mon Sep 17 00:00:00 2001
From: Tyler Benson <tylerbenson@gmail.com>
Date: Fri, 3 Mar 2023 13:34:28 -0500
Subject: [PATCH] Add span link attributes and fix var names

---
 .../AwsXRayEnvSpanLinksExtractor.java         | 13 ++++++++----
 .../v1_0/AbstractAwsLambdaTest.java           | 20 +++++++++++++++++--
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/AwsXRayEnvSpanLinksExtractor.java b/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/AwsXRayEnvSpanLinksExtractor.java
index 2ed8527891b0..2935f985156a 100644
--- a/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/AwsXRayEnvSpanLinksExtractor.java
+++ b/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/AwsXRayEnvSpanLinksExtractor.java
@@ -5,6 +5,8 @@
 
 package io.opentelemetry.instrumentation.awslambdacore.v1_0.internal;
 
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.context.Context;
@@ -28,6 +30,9 @@ public final class AwsXRayEnvSpanLinksExtractor implements SpanLinksExtractor<Aw
   // lower-case map getter used for extraction
   private static final String AWS_TRACE_HEADER_PROPAGATOR_KEY = "x-amzn-trace-id";
 
+  private static final Attributes LINK_ATTRIBUTES =
+      Attributes.of(AttributeKey.stringKey("source"), "x-ray-env");
+
   @Override
   public void extract(
       SpanLinksBuilder spanLinks,
@@ -41,16 +46,16 @@ public static void extract(SpanLinksBuilder spanLinks) {
     if (parentTraceHeader == null || parentTraceHeader.isEmpty()) {
       return;
     }
-    Context parentCtx =
+    Context xrayContext =
         AwsXrayPropagator.getInstance()
             .extract(
                 // see BaseTracer#extract() on why we're using root() here
                 Context.root(),
                 Collections.singletonMap(AWS_TRACE_HEADER_PROPAGATOR_KEY, parentTraceHeader),
                 MapGetter.INSTANCE);
-    SpanContext parent = Span.fromContext(parentCtx).getSpanContext();
-    if (parent.isValid()) {
-      spanLinks.addLink(parent);
+    SpanContext envVarSpanCtx = Span.fromContext(xrayContext).getSpanContext();
+    if (envVarSpanCtx.isValid()) {
+      spanLinks.addLink(envVarSpanCtx, LINK_ATTRIBUTES);
     }
   }
 
diff --git a/instrumentation/aws-lambda/aws-lambda-core-1.0/testing/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/AbstractAwsLambdaTest.java b/instrumentation/aws-lambda/aws-lambda-core-1.0/testing/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/AbstractAwsLambdaTest.java
index 4b994e737df0..ccc8c6cbcf26 100644
--- a/instrumentation/aws-lambda/aws-lambda-core-1.0/testing/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/AbstractAwsLambdaTest.java
+++ b/instrumentation/aws-lambda/aws-lambda-core-1.0/testing/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/AbstractAwsLambdaTest.java
@@ -12,6 +12,8 @@
 
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestHandler;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.sdk.trace.data.StatusData;
@@ -108,8 +110,22 @@ void handlerLinksToInfrastructureTrace() {
                     span ->
                         span.hasName("my_function")
                             .hasKind(SpanKind.SERVER)
-                            .hasTraceId("8a3c60f7d188f8fa79d48a391a778fa6")
-                            .hasParentSpanId("0000000000000456")
+                            .hasLinksSatisfying(
+                                links ->
+                                    assertThat(links)
+                                        .singleElement()
+                                        .satisfies(
+                                            link -> {
+                                              assertThat(link.getSpanContext().getTraceId())
+                                                  .isEqualTo("8a3c60f7d188f8fa79d48a391a778fa6");
+                                              assertThat(link.getSpanContext().getSpanId())
+                                                  .isEqualTo("0000000000000456");
+                                              assertThat(link.getAttributes())
+                                                  .isEqualTo(
+                                                      Attributes.of(
+                                                          AttributeKey.stringKey("source"),
+                                                          "x-ray-env"));
+                                            }))
                             .hasAttributesSatisfying(
                                 attrs ->
                                     assertThat(attrs)