feat: add tenant id to lambda context and structured log message

Author: Ramisa Alam

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

@@ -581,7 +581,8 @@ public ByteArrayOutputStream call(InvocationRequest request) throws Error, Excep
                         cognitoIdentity,
                         LambdaEnvironment.FUNCTION_VERSION,
                         request.getInvokedFunctionArn(),
-                        clientContext
+                        clientContext,
+                        request.getTenantId()
                 );
 
                 safeAddContextToLambdaLogger(context);

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

@@ -22,6 +22,7 @@ public class LambdaContext implements Context {
     private final long deadlineTimeInMs;
     private final CognitoIdentity cognitoIdentity;
     private final ClientContext clientContext;
+    private final String tenantId;
     private final LambdaLogger logger;
 
     public LambdaContext(
@@ -34,7 +35,8 @@ public LambdaContext(
             CognitoIdentity identity,
             String functionVersion,
             String invokedFunctionArn,
-            ClientContext clientContext
+            ClientContext clientContext,
+            String tenantId
     ) {
         this.memoryLimit = memoryLimit;
         this.deadlineTimeInMs = deadlineTimeInMs;
@@ -46,6 +48,7 @@ public LambdaContext(
         this.clientContext = clientContext;
         this.functionVersion = functionVersion;
         this.invokedFunctionArn = invokedFunctionArn;
+        this.tenantId = tenantId;
         this.logger = com.amazonaws.services.lambda.runtime.LambdaRuntime.getLogger();
     }
 
@@ -91,6 +94,10 @@ public int getRemainingTimeInMillis() {
         return delta > 0 ? delta : 0;
     }
 
+    public String getTenantId() {
+        return tenantId;
+    }
+
     public LambdaLogger getLogger() {
         return logger;
     }

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

@@ -41,6 +41,7 @@ private StructuredLogMessage createLogMessage(String message, LogLevel logLevel)
 
         if (lambdaContext != null) {
             msg.AWSRequestId = lambdaContext.getAwsRequestId();
+            msg.tenantId = lambdaContext.getTenantId();
         }
         return msg;
     }

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

@@ -12,4 +12,5 @@ class StructuredLogMessage {
     public String message;
     public LogLevel level;
     public String AWSRequestId;
+    public String tenantId;
 }

aws-lambda-java-runtime-interface-client/src/main/java/com/amazonaws/services/lambda/runtime/api/client/runtimeapi/dto/InvocationRequest.java

@@ -40,6 +40,11 @@ public class InvocationRequest {
      */
     private String cognitoIdentity;
 
+    /**
+     * The tenant ID associated with the request.
+     */
+    private String tenantId;
+
     private byte[] content;
 
     public String getId() {
@@ -94,6 +99,14 @@ public void setCognitoIdentity(String cognitoIdentity) {
         this.cognitoIdentity = cognitoIdentity;
     }
 
+    public String getTenantId() {
+        return tenantId;
+    }
+
+    public void setTenantId(String tenantId) {
+        this.tenantId = tenantId;
+    }
+
     public byte[] getContent() {
         return content;
     }

aws-lambda-java-runtime-interface-client/src/main/jni/com_amazonaws_services_lambda_runtime_api_client_runtimeapi_NativeClient.cpp

@@ -20,6 +20,7 @@ static jfieldID contentField;
 static jfieldID clientContextField;
 static jfieldID cognitoIdentityField;
 static jfieldID xrayTraceIdField;
+static jfieldID tenantIdField;
 
 
 jint JNI_OnLoad(JavaVM* vm, void* reserved) {
@@ -41,6 +42,7 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) {
     xrayTraceIdField = env->GetFieldID(invocationRequestClass , "xrayTraceId", "Ljava/lang/String;");
     clientContextField = env->GetFieldID(invocationRequestClass , "clientContext", "Ljava/lang/String;");
     cognitoIdentityField = env->GetFieldID(invocationRequestClass , "cognitoIdentity", "Ljava/lang/String;");
+    tenantIdField = env->GetFieldID(invocationRequestClass, "tenantId", "Ljava/lang/String;");
 
     return JNI_VERSION;
 }
@@ -106,6 +108,10 @@ JNIEXPORT jobject JNICALL Java_com_amazonaws_services_lambda_runtime_api_client_
     CHECK_EXCEPTION(env, env->SetObjectField(invocationRequest, cognitoIdentityField, env->NewStringUTF(response.cognito_identity.c_str())));
   }
 
+  if(response.tenant_id != ""){
+    CHECK_EXCEPTION(env, env->SetObjectField(invocationRequest, tenantIdField, env->NewStringUTF(response.tenant_id.c_str())));
+  }
+
   bytes = reinterpret_cast<const jbyte*>(response.payload.c_str());
   CHECK_EXCEPTION(env, jArray = env->NewByteArray(response.payload.length()));
   CHECK_EXCEPTION(env, env->SetByteArrayRegion(jArray, 0, response.payload.length(), bytes));

aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/include/aws/lambda-runtime/runtime.h

@@ -61,6 +61,11 @@ struct invocation_request {
      */
     std::chrono::time_point<std::chrono::system_clock> deadline;
 
+    /**
+     * Tenant ID of the current invocation.
+     */
+    std::string tenant_id;
+
     /**
      * The number of milliseconds left before lambda terminates the current execution.
      */

aws-lambda-java-runtime-interface-client/src/main/jni/deps/aws-lambda-cpp-0.2.7/src/runtime.cpp

@@ -40,6 +40,7 @@ static constexpr auto CLIENT_CONTEXT_HEADER = "lambda-runtime-client-context";
 static constexpr auto COGNITO_IDENTITY_HEADER = "lambda-runtime-cognito-identity";
 static constexpr auto DEADLINE_MS_HEADER = "lambda-runtime-deadline-ms";
 static constexpr auto FUNCTION_ARN_HEADER = "lambda-runtime-invoked-function-arn";
+static constexpr auto TENANT_ID_HEADER = "lambda-runtime-aws-tenant-id";
 
 enum Endpoints {
     INIT,
@@ -301,6 +302,10 @@ runtime::next_outcome runtime::get_next()
             req.payload.c_str(),
             static_cast<int64_t>(req.get_time_remaining().count()));
     }
+
+    if (resp.has_header(TENANT_ID_HEADER)) {
+        req.tenant_id = resp.get_header(TENANT_ID_HEADER);
+    }
     return next_outcome(req);
 }
 

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

@@ -18,6 +18,7 @@ public class LambdaContextTest {
     private static final String INVOKED_FUNCTION_ARN = "invoked-function-arn";
     private static final LambdaClientContext CLIENT_CONTEXT = new LambdaClientContext();
     public static final int MEMORY_LIMIT = 128;
+    public static final String TENANT_ID = "tenant-id";
 
     @Test
     public void getRemainingTimeInMillis() {
@@ -54,6 +55,6 @@ public void getRemainingTimeInMillis_Deadline() throws InterruptedException {
 
     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, CLIENT_CONTEXT);
+                FUNCTION_NAME, IDENTITY, FUNCTION_VERSION, INVOKED_FUNCTION_ARN, CLIENT_CONTEXT, TENANT_ID);
     }
 }

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

@@ -29,7 +29,8 @@ void testFormattingWithLambdaContext() {
                 null,
                 null,
                 "function-arn",
-                null
+                null,
+                "tenant-id"
         );
         assertFormatsString("test log", LogLevel.WARN, context);
     }
@@ -52,6 +53,7 @@ void assert_expected_log_message(StructuredLogMessage result, String message, Lo
 
         if (context != null) {
             assertEquals(context.getAwsRequestId(), result.AWSRequestId);
+            assertEquals(context.getTenantId(), result.tenantId);
         }
     }
 }

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

@@ -15,6 +15,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.ErrorRequest;
+import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.InvocationRequest;
 import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.StackElement;
 import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.XRayErrorCause;
 import com.amazonaws.services.lambda.runtime.api.client.runtimeapi.dto.XRayException;
@@ -333,6 +334,31 @@ public void nextTest() {
         }
     }
 
+    @Test
+    public void nextWithTenantIdTest() {
+        try {
+            MockResponse mockResponse = new MockResponse();
+            mockResponse.setResponseCode(HTTP_ACCEPTED);
+            mockResponse.setHeader("lambda-runtime-aws-request-id", "1234567890");
+            mockResponse.setHeader("Content-Type", "application/json");
+            String expectedTenantId = "my-tenant-id";
+            mockResponse.setHeader("lambda-runtime-aws-tenant-id", expectedTenantId);
+            mockWebServer.enqueue(mockResponse);
+
+            InvocationRequest invocationRequest = lambdaRuntimeApiClientImpl.nextInvocation();
+            assertEquals(expectedTenantId, invocationRequest.getTenantId());
+
+            RecordedRequest recordedRequest = mockWebServer.takeRequest();
+            HttpUrl actualUrl = recordedRequest.getRequestUrl();
+            String expectedUrl = "http://" + getHostnamePort() + "/2018-06-01/runtime/invocation/next";
+            assertEquals(expectedUrl, actualUrl.toString());
+
+            String actualBody = recordedRequest.getBody().readUtf8();
+            assertEquals("", actualBody);
+        } catch(Exception e) {
+            fail();
+        }
+    }
 
     @Test
     public void createUrlMalformedTest() {