Propagate otel context through custom aws client context for lambda direct calls

instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/AwsLambdaFunctionInstrumenter.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.internal.ContextPropagationDebug;
 import io.opentelemetry.instrumentation.awslambdacore.v1_0.AwsLambdaRequest;
+import java.util.HashMap;
 import java.util.Locale;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -48,11 +49,20 @@ public void end(
 
   public Context extract(AwsLambdaRequest input) {
     ContextPropagationDebug.debugContextLeakIfEnabled();
+    // Look in both the http headers and the custom client context
+    Map<String, String> headers = input.getHeaders();
+    if (input.getAwsContext() != null && input.getAwsContext().getClientContext() != null) {
+      Map<String, String> customContext = input.getAwsContext().getClientContext().getCustom();
+      if (customContext != null) {
+        headers = new HashMap<>(headers);
+        headers.putAll(customContext);
+      }
+    }
 
     return openTelemetry
         .getPropagators()
         .getTextMapPropagator()
-        .extract(Context.root(), input.getHeaders(), MapGetter.INSTANCE);
+        .extract(Context.root(), headers, MapGetter.INSTANCE);
   }
 
   private enum MapGetter implements TextMapGetter<Map<String, String>> {

instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/InstrumenterExtractionTest.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awslambdacore.v1_0.internal;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.amazonaws.services.lambda.runtime.ClientContext;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.instrumentation.awslambdacore.v1_0.AwsLambdaRequest;
+import java.util.HashMap;
+import org.junit.jupiter.api.Test;
+
+class InstrumenterExtractionTest {
+  @Test
+  public void useCustomContext() {
+    AwsLambdaFunctionInstrumenter instr =
+        AwsLambdaFunctionInstrumenterFactory.createInstrumenter(
+            OpenTelemetry.propagating(
+                ContextPropagators.create(W3CTraceContextPropagator.getInstance())));
+    com.amazonaws.services.lambda.runtime.Context awsContext =
+        mock(com.amazonaws.services.lambda.runtime.Context.class);
+    ClientContext clientContext = mock(ClientContext.class);
+    when(awsContext.getClientContext()).thenReturn(clientContext);
+    HashMap<String, String> customMap = new HashMap<>();
+    customMap.put("traceparent", "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01");
+    when(clientContext.getCustom()).thenReturn(customMap);
+
+    AwsLambdaRequest input = AwsLambdaRequest.create(awsContext, new HashMap<>(), new HashMap<>());
+
+    Context extracted = instr.extract(input);
+    assertThat(extracted.toString().contains("3577")).isTrue();
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/library/build.gradle.kts

@@ -7,6 +7,7 @@ dependencies {
 
   library("software.amazon.awssdk:aws-core:2.2.0")
   library("software.amazon.awssdk:sqs:2.2.0")
+  library("software.amazon.awssdk:lambda:2.2.0")
   library("software.amazon.awssdk:sns:2.2.0")
   library("software.amazon.awssdk:aws-json-protocol:2.2.0")
   compileOnly(project(":muzzle")) // For @NoMuzzle

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/DirectLambdaAccess.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
+import software.amazon.awssdk.core.SdkRequest;
+
+final class DirectLambdaAccess {
+  private DirectLambdaAccess() {}
+
+  private static final boolean enabled = PluginImplUtil.isImplPresent("DirectLambdaImpl");
+
+  @NoMuzzle
+  public static SdkRequest modifyRequest(SdkRequest request, Context otelContext) {
+    return enabled ? DirectLambdaImpl.modifyRequest(request, otelContext) : null;
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/DirectLambdaImpl.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+import software.amazon.awssdk.core.SdkRequest;
+import software.amazon.awssdk.services.lambda.model.InvokeRequest;
+
+// this class is only used from DirectLambdaAccess from method with @NoMuzzle annotation
+
+// Direct lambda invocations (e.g., not through an api gateway) currently strip
+// away the otel propagation headers (but leave x-ray ones intact).  Use the
+// custom client context header as an additional propagation mechanism for this
+// very specific scenario.  For reference, the header is named "X-Amz-Client-Context" but the api to
+// manipulate
+// it abstracts that away.  The client context field is documented in
+// https://docs.aws.amazon.com/lambda/latest/api/API_Invoke.html#API_Invoke_RequestParameters
+
+final class DirectLambdaImpl {
+  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+  public static final String CLIENT_CONTEXT_CUSTOM_FIELDS_KEY = "custom";
+  public static final int MAX_CLIENT_CONTEXT_LENGTH = 3583;
+
+  private DirectLambdaImpl() {}
+
+  @Nullable
+  static SdkRequest modifyRequest(
+      SdkRequest request, io.opentelemetry.context.Context otelContext) {
+    if (isDirectLambdaInvocation(request)) {
+      try {
+        return modifyOrAddCustomContextHeader((InvokeRequest) request, otelContext);
+      } catch (Exception e) {
+        return null;
+      }
+    }
+    return null;
+  }
+
+  static boolean isDirectLambdaInvocation(SdkRequest request) {
+    return request instanceof InvokeRequest;
+  }
+
+  @SuppressWarnings("unchecked")
+  static SdkRequest modifyOrAddCustomContextHeader(
+      InvokeRequest request, io.opentelemetry.context.Context otelContext) throws Exception {
+    InvokeRequest.Builder builder = request.toBuilder();
+    // Unfortunately the value of this thing is a base64-encoded json with a character limit; also
+    // therefore not comma-composable like many http headers
+    String clientContextString = request.clientContext();
+    String clientContextJsonString = "{}";
+    if (clientContextString != null && !clientContextString.isEmpty()) {
+      clientContextJsonString =
+          new String(Base64.getDecoder().decode(clientContextString), StandardCharsets.UTF_8);
+    }
+    Map<String, Object> parsedJson =
+        (Map<String, Object>)
+            OBJECT_MAPPER.readValue(clientContextJsonString, new TypeReference<Object>() {});
+    Map<String, Object> customFields =
+        (Map<String, Object>)
+            parsedJson.getOrDefault(
+                CLIENT_CONTEXT_CUSTOM_FIELDS_KEY, new HashMap<String, Object>());
+
+    int numCustomFields = customFields.size();
+    GlobalOpenTelemetry.getPropagators()
+        .getTextMapPropagator()
+        .inject(otelContext, customFields, Map::put);
+    if (numCustomFields == customFields.size()) {
+      return null; // no modifications made
+    }
+
+    parsedJson.put(CLIENT_CONTEXT_CUSTOM_FIELDS_KEY, customFields);
+
+    // turn it back into a string (json encode)
+    String newJson = OBJECT_MAPPER.writeValueAsString(parsedJson);
+    // turn it back into a base64 string
+    String newJson64 = Base64.getEncoder().encodeToString(newJson.getBytes(StandardCharsets.UTF_8));
+    // check it for length (err on the safe side with >=)
+    if (newJson64.length() >= MAX_CLIENT_CONTEXT_LENGTH) {
+      return null;
+    }
+    builder.clientContext(newJson64);
+    return builder.build();
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/TracingExecutionInterceptor.java

@@ -212,6 +212,10 @@ public SdkRequest modifyRequest(
     if (modifiedRequest != null) {
       return modifiedRequest;
     }
+    modifiedRequest = DirectLambdaAccess.modifyRequest(request, otelContext);
+    if (modifiedRequest != null) {
+      return modifiedRequest;
+    }
 
     // Insert other special handling here, following the same pattern as SQS and SNS.
 

instrumentation/aws-sdk/aws-sdk-2.2/library/src/test/java/io/opentelemetry/instrumentation/awssdk/v2_2/DirectLambdaTest.java

@@ -0,0 +1,110 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import software.amazon.awssdk.services.lambda.model.InvokeRequest;
+
+public class DirectLambdaTest {
+  private Context context;
+
+  @Before
+  public void setup() {
+    GlobalOpenTelemetry.resetForTest();
+    ;
+    OpenTelemetrySdk.builder()
+        .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
+        .buildAndRegisterGlobal();
+
+    Span parent =
+        GlobalOpenTelemetry.getTracer("test")
+            .spanBuilder("parentSpan")
+            .setSpanKind(SpanKind.SERVER)
+            .startSpan();
+    context = parent.storeInContext(Context.current());
+    assertThat(context.toString().equals("{}")).isFalse();
+    parent.end();
+  }
+
+  @After
+  public void cleanup() {
+    GlobalOpenTelemetry.resetForTest();
+  }
+
+  private static String base64ify(String json) {
+    return Base64.getEncoder().encodeToString(json.getBytes(StandardCharsets.UTF_8));
+  }
+
+  @Test
+  public void noExistingClientContext() throws Exception {
+    InvokeRequest r = InvokeRequest.builder().build();
+
+    InvokeRequest newR =
+        (InvokeRequest) DirectLambdaImpl.modifyOrAddCustomContextHeader(r, context);
+
+    String newClientContext = newR.clientContext();
+    newClientContext =
+        new String(Base64.getDecoder().decode(newClientContext), StandardCharsets.UTF_8);
+    assertThat(newClientContext.contains("traceparent")).isTrue();
+  }
+
+  @Test
+  public void withExistingClientContext() throws Exception {
+    String clientContext =
+        base64ify(
+            "{\"otherStuff\": \"otherValue\", \"custom\": {\"preExisting\": \"somevalue\"} }");
+    InvokeRequest r = InvokeRequest.builder().clientContext(clientContext).build();
+
+    InvokeRequest newR =
+        (InvokeRequest) DirectLambdaImpl.modifyOrAddCustomContextHeader(r, context);
+
+    String newClientContext = newR.clientContext();
+    newClientContext =
+        new String(Base64.getDecoder().decode(newClientContext), StandardCharsets.UTF_8);
+    assertThat(newClientContext.contains("traceparent")).isTrue();
+    assertThat(newClientContext.contains("preExisting")).isTrue();
+    assertThat(newClientContext.contains("otherStuff")).isTrue();
+  }
+
+  @Test
+  public void exceedingMaximumLengthDoesNotModify() throws Exception {
+    // awkward way to build a valid json that is almost but not quite too long
+    boolean continueLengthingInput = true;
+    StringBuffer x = new StringBuffer("x");
+    String long64edClientContext = "";
+    while (continueLengthingInput) {
+      x.append("x");
+      String newClientContext = base64ify("{\"" + x + "\": \"" + x + "\"}");
+      if (newClientContext.length() >= DirectLambdaImpl.MAX_CLIENT_CONTEXT_LENGTH) {
+        continueLengthingInput = false;
+        break;
+      }
+      long64edClientContext = newClientContext;
+      continueLengthingInput =
+          long64edClientContext.length() < DirectLambdaImpl.MAX_CLIENT_CONTEXT_LENGTH;
+    }
+
+    InvokeRequest r = InvokeRequest.builder().clientContext(long64edClientContext).build();
+    assertThat(r.clientContext().equals(long64edClientContext)).isTrue();
+
+    InvokeRequest newR =
+        (InvokeRequest) DirectLambdaImpl.modifyOrAddCustomContextHeader(r, context);
+    assertThat(newR == null).isTrue(); // null return means no modification performed
+  }
+}