feat(context): add w3c() helper for W3C trace fields

.github/workflows/dockerized-test.yml

@@ -0,0 +1,68 @@
+name: dockerized-test
+
+permissions:
+  contents: read
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: ['*']
+    paths:
+      - 'aws-lambda-java-core/**'
+      - 'aws-lambda-java-serialization/**'
+      - 'aws-lambda-java-runtime-interface-client/**'
+      - '.github/workflows/dockerized-test.yml'
+  workflow_dispatch:
+
+jobs:
+  dockerized-test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+
+      - name: Set up JDK
+        uses: actions/setup-java@v5
+        with:
+          # The core/serialization/RIC builds require a JDK 8 toolchain
+          # (maven-toolchains-plugin: [1.8,9)); 25 is the JVM Maven runs on.
+          java-version: |
+            8
+            25
+          distribution: corretto
+          cache: maven
+
+      - name: Build core, serialization and the RIC under test
+        run: |
+          mvn -B -q install -DskipTests --file aws-lambda-java-core/pom.xml
+          mvn -B -q install -DskipTests --file aws-lambda-java-serialization/pom.xml
+          mvn -B -q clean install -DskipTests -DmultiArch=false \
+            --file aws-lambda-java-runtime-interface-client/pom.xml
+          # Stage core + serialization next to the RIC jar so the image build
+          # can pick all three up from a single ./target folder.
+          cp aws-lambda-java-core/target/aws-lambda-java-core-*.jar \
+             aws-lambda-java-runtime-interface-client/target/
+          cp aws-lambda-java-serialization/target/aws-lambda-java-serialization-*.jar \
+             aws-lambda-java-runtime-interface-client/target/
+
+      - name: Compile the W3C task handler
+        working-directory: aws-lambda-java-runtime-interface-client
+        run: |
+          javac --release 8 -cp "$(ls target/aws-lambda-java-core-*.jar)" \
+            -d test/dockerized/tasks \
+            test/dockerized/handler/W3CHandler.java
+
+      - name: Build the test image
+        working-directory: aws-lambda-java-runtime-interface-client
+        run: |
+          docker build . \
+            -t local/test \
+            -f Dockerfile.test \
+            --build-arg BASE_IMAGE=public.ecr.aws/lambda/java:25
+
+      - name: Run dockerized suites
+        uses: aws/containerized-test-runner-for-aws-lambda@0863dd17b5fc19585250a2405c0f939a77b4f397 # main
+        with:
+          suiteFileArray: '["./aws-lambda-java-runtime-interface-client/test/dockerized/suites/*.json"]'
+          dockerImageName: 'local/test'
+          taskFolder: './aws-lambda-java-runtime-interface-client/test/dockerized/tasks'

aws-lambda-java-core/src/main/java/com/amazonaws/services/lambda/runtime/Context.java

@@ -2,6 +2,9 @@
 
 package com.amazonaws.services.lambda.runtime;
 
+import java.util.Collections;
+import java.util.Map;
+
 /**
  * 
  * The context object allows you to access useful information available within
@@ -119,4 +122,20 @@ default String getTenantId() {
 	default String getXrayTraceId() {
 		return null;
 	}
+
+	/**
+	 * Returns the W3C trace-context fields associated with the request at invoke
+	 * time.
+	 * <p>
+	 * The returned map only ever contains the allowlisted W3C trace-context
+	 * keys ({@code traceparent}, {@code tracestate}, {@code baggage}) that were
+	 * carried on the invoke's client context. Any other key, and any allowlisted
+	 * key whose value is not a string, is dropped.
+	 * </p>
+	 *
+	 * @return an unmodifiable map of W3C trace-context fields; empty by default
+	 */
+	default Map<String, String> w3c() {
+		return Collections.emptyMap();
+	}
 }

aws-lambda-java-runtime-interface-client/Dockerfile.test

@@ -0,0 +1,16 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+ARG BASE_IMAGE=public.ecr.aws/lambda/java:25
+FROM $BASE_IMAGE
+
+COPY target/aws-lambda-java-runtime-interface-client-*.jar /opt/ric/
+COPY target/aws-lambda-java-core-*.jar /opt/ric/
+COPY target/aws-lambda-java-serialization-*.jar /opt/ric/
+
+ENTRYPOINT ["/bin/sh", "-c", \
+"exec /usr/local/bin/aws-lambda-rie java \
+-cp \"/opt/ric/*:/var/task\" \
+com.amazonaws.services.lambda.runtime.api.client.AWSLambda \"$@\"", \
+"sh"]
+CMD ["W3CHandler::getW3c"]

aws-lambda-java-runtime-interface-client/pom.xml

@@ -65,7 +65,7 @@
         <dependency>
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-lambda-java-core</artifactId>
-            <version>1.4.0</version>
+            <version>1.4.1-SNAPSHOT</version>
         </dependency>
         <dependency>
             <groupId>com.amazonaws</groupId>

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaClientContext.java

@@ -14,6 +14,7 @@ public class LambdaClientContext implements ClientContext {
     private LambdaClientContextClient client;
     private Map<String, String> custom;
     private Map<String, String> env;
+    private Object w3c;
 
     public Client getClient() {
         return client;
@@ -26,4 +27,10 @@ public Map<String, String> getCustom() {
     public Map<String, String> getEnvironment() {
         return env;
     }
+
+    Object readAndStripW3c() {
+        Object raw = this.w3c;
+        this.w3c = null;
+        return raw;
+    }
 }

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContext.java

@@ -9,9 +9,22 @@
 import com.amazonaws.services.lambda.runtime.CognitoIdentity;
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.LambdaLogger;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
 
 public class LambdaContext implements Context {
 
+    /**
+     * Allowlist of W3C trace-context fields that may be surfaced through
+     * {@link #w3c()}. Any other key carried on {@code clientContext.w3c} is
+     * ignored, and any allowlisted key whose value is not a string is dropped.
+     */
+    public static final List<String> W3C_ALLOWED_FIELDS =
+            Collections.unmodifiableList(Arrays.asList("traceparent", "tracestate", "baggage"));
+
     private int memoryLimit;
     private final String awsRequestId;
     private final String logGroupName;
@@ -25,6 +38,7 @@ public class LambdaContext implements Context {
     private final String tenantId;
     private final String xrayTraceId;
     private final LambdaLogger logger;
+    private final Map<String, String> w3cFields;
 
     public LambdaContext(
             int memoryLimit,
@@ -52,6 +66,7 @@ public LambdaContext(
         this.invokedFunctionArn = invokedFunctionArn;
         this.tenantId = tenantId;
         this.xrayTraceId = xrayTraceId;
+        this.w3cFields = extractAndStripW3c(clientContext);
         this.logger = com.amazonaws.services.lambda.runtime.LambdaRuntime.getLogger();
     }
 
@@ -105,6 +120,32 @@ public String getXrayTraceId() {
         return xrayTraceId;
     }
 
+    @Override
+    public Map<String, String> w3c() {
+        return w3cFields;
+    }
+
+    private static Map<String, String> extractAndStripW3c(ClientContext clientContext) {
+        if (!(clientContext instanceof LambdaClientContext)) {
+            return Collections.emptyMap();
+        }
+
+        Object rawW3c = ((LambdaClientContext) clientContext).readAndStripW3c();
+        if (!(rawW3c instanceof Map)) {
+            return Collections.emptyMap();
+        }
+
+        Map<?, ?> source = (Map<?, ?>) rawW3c;
+        Map<String, String> fields = new LinkedHashMap<>();
+        for (String key : W3C_ALLOWED_FIELDS) {
+            Object value = source.get(key);
+            if (value instanceof String) {
+                fields.put(key, (String) value);
+            }
+        }
+        return Collections.unmodifiableMap(fields);
+    }
+
     public LambdaLogger getLogger() {
         return logger;
     }

aws-lambda-java-runtime-interface-client/src/test/java/com/amazonaws/services/lambda/runtime/api/client/api/LambdaContextTest.java

@@ -2,9 +2,17 @@
 
 package com.amazonaws.services.lambda.runtime.api.client.api;
 
+import com.amazonaws.services.lambda.runtime.ClientContext;
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.serialization.PojoSerializer;
+import com.amazonaws.services.lambda.runtime.serialization.factories.GsonFactory;
 import org.junit.jupiter.api.Test;
 
+import java.util.Map;
+
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class LambdaContextTest {
@@ -58,4 +66,97 @@ private LambdaContext createContextWithDeadline(long deadlineTimeInMs) {
         return new LambdaContext(MEMORY_LIMIT, deadlineTimeInMs, REQUEST_ID, LOG_GROUP_NAME, LOG_STREAM_NAME,
                 FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, TENANT_ID, X_RAY_TRACE_ID, CLIENT_CONTEXT);
     }
+
+    private LambdaContext w3cContext(ClientContext clientContext) {
+        return new LambdaContext(MEMORY_LIMIT, System.currentTimeMillis() + 1000, REQUEST_ID, LOG_GROUP_NAME,
+                LOG_STREAM_NAME, FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, TENANT_ID,
+                X_RAY_TRACE_ID, clientContext);
+    }
+
+    private LambdaClientContext clientContextFromJson(String json) {
+        PojoSerializer<LambdaClientContext> serializer =
+                GsonFactory.getInstance().getSerializer(LambdaClientContext.class);
+        return serializer.fromJson(json);
+    }
+
+    @Test
+    public void w3c_returnsEmptyWhenClientContextIsNull() {
+        assertTrue(w3cContext(null).w3c().isEmpty());
+    }
+
+    @Test
+    public void w3c_returnsEmptyWhenNoW3cKey() {
+        LambdaClientContext cc = clientContextFromJson("{\"custom\": {\"value\": \"test\"}}");
+        assertTrue(w3cContext(cc).w3c().isEmpty());
+    }
+
+    @Test
+    public void w3c_returnsAllAllowlistedFields() {
+        LambdaClientContext cc = clientContextFromJson(
+                "{\"w3c\": {" +
+                        "\"traceparent\": \"00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01\"," +
+                        "\"tracestate\": \"rojo=00f067aa0ba902b7\"," +
+                        "\"baggage\": \"userId=alice\"}}");
+
+        Map<String, String> w3c = w3cContext(cc).w3c();
+
+        assertEquals(3, w3c.size());
+        assertEquals("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", w3c.get("traceparent"));
+        assertEquals("rojo=00f067aa0ba902b7", w3c.get("tracestate"));
+        assertEquals("userId=alice", w3c.get("baggage"));
+    }
+
+    @Test
+    public void w3c_dropsNonAllowlistedKeysAndNonStringValues() {
+        LambdaClientContext cc = clientContextFromJson(
+                "{\"w3c\": {" +
+                        "\"baggage\": \"keep=me\"," +
+                        "\"unknownField\": \"drop-me\"," +
+                        "\"traceparent\": 42," +
+                        "\"tracestate\": null}}");
+
+        Map<String, String> w3c = w3cContext(cc).w3c();
+
+        assertEquals(1, w3c.size());
+        assertEquals("keep=me", w3c.get("baggage"));
+    }
+
+    @Test
+    public void w3c_treatsNonObjectPayloadAsEmpty() {
+        assertTrue(w3cContext(clientContextFromJson("{\"w3c\": \"not-an-object\"}")).w3c().isEmpty());
+        assertTrue(w3cContext(clientContextFromJson("{\"w3c\": [\"baggage=abc\"]}")).w3c().isEmpty());
+    }
+
+    @Test
+    public void w3c_isStrippedFromClientContextAndIsUnmodifiable() {
+        LambdaClientContext cc = clientContextFromJson(
+                "{\"custom\": {\"value\": \"test\"}, \"w3c\": {\"baggage\": \"userId=alice\"}}");
+
+        LambdaContext context = w3cContext(cc);
+
+        assertEquals("userId=alice", context.w3c().get("baggage"));
+        assertThrows(UnsupportedOperationException.class, () -> context.w3c().put("x", "y"));
+        assertEquals("test", cc.getCustom().get("value"));
+        assertTrue(w3cContext(cc).w3c().isEmpty());
+    }
+
+    @Test
+    public void w3c_defaultsToEmptyOnBareContextInterface() {
+        Context bare = new Context() {
+            public String getAwsRequestId() { return null; }
+            public String getLogGroupName() { return null; }
+            public String getLogStreamName() { return null; }
+            public String getFunctionName() { return null; }
+            public String getFunctionVersion() { return null; }
+            public String getInvokedFunctionArn() { return null; }
+            public com.amazonaws.services.lambda.runtime.CognitoIdentity getIdentity() { return null; }
+            public ClientContext getClientContext() { return null; }
+            public int getRemainingTimeInMillis() { return 0; }
+            public int getMemoryLimitInMB() { return 0; }
+            public com.amazonaws.services.lambda.runtime.LambdaLogger getLogger() { return null; }
+        };
+
+        assertTrue(bare.w3c().isEmpty());
+        assertFalse(bare.w3c() == null);
+    }
 }

aws-lambda-java-runtime-interface-client/test/dockerized/handler/W3CHandler.java

@@ -0,0 +1,26 @@
+/*
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+SPDX-License-Identifier: Apache-2.0
+*/
+
+import com.amazonaws.services.lambda.runtime.ClientContext;
+import com.amazonaws.services.lambda.runtime.Context;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+
+public class W3CHandler {
+
+    public Map<String, String> getW3c(Map<String, Object> event, Context context) {
+        return context.w3c();
+    }
+
+    public Map<String, Object> getW3cAndCustom(Map<String, Object> event, Context context) {
+        Map<String, Object> result = new LinkedHashMap<>();
+        result.put("w3c", context.w3c());
+        ClientContext clientContext = context.getClientContext();
+        result.put("custom", clientContext == null ? null : clientContext.getCustom());
+        return result;
+    }
+}