GH-1187 Fix AWS Context initialization for Custom Runtime

Updated sample to show that Context is not null

Resolves #1187
This commit is contained in:
Oleg Zhurakousky
2025-01-13 13:23:15 +01:00
parent a0405d306c
commit 26ef1fa76e
5 changed files with 114 additions and 9 deletions

View File

@@ -28,10 +28,14 @@ import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import com.amazonaws.services.lambda.runtime.ClientContext;
import com.amazonaws.services.lambda.runtime.CognitoIdentity;
import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.LambdaLogger;
import com.amazonaws.services.lambda.runtime.LambdaRuntime;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.cloud.function.context.FunctionCatalog;
import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
import org.springframework.cloud.function.context.config.RoutingFunction;
@@ -130,6 +134,8 @@ public final class CustomRuntimeEventLoop implements SmartLifecycle {
logger.debug("Attempting to get new event");
ResponseEntity<String> response = this.pollForData(rest, requestEntity);
Context clientContext = generateClientContext(response.getHeaders());
if (logger.isDebugEnabled()) {
logger.debug("New Event received: " + response);
}
@@ -140,9 +146,9 @@ public final class CustomRuntimeEventLoop implements SmartLifecycle {
FunctionInvocationWrapper function = locateFunction(environment, functionCatalog, response.getHeaders());
ByteArrayInputStream is = new ByteArrayInputStream(response.getBody().getBytes(StandardCharsets.UTF_8));
Message<?> requestMessage = AWSLambdaUtils.generateMessage(is, function.getInputType(), function.isSupplier(), mapper, null);
Message<?> requestMessage = AWSLambdaUtils.generateMessage(is, function.getInputType(), function.isSupplier(), mapper, clientContext);
Object functionResponse = function.apply(requestMessage);
byte[] responseBytes = AWSLambdaUtils.generateOutputFromObject(requestMessage, functionResponse, mapper, function.getOutputType());
String invocationUrl = MessageFormat
@@ -157,12 +163,91 @@ public final class CustomRuntimeEventLoop implements SmartLifecycle {
}
}
catch (Exception e) {
e.printStackTrace();
this.propagateAwsError(requestId, e, mapper, runtimeApi, rest);
}
}
}
}
private Context generateClientContext(HttpHeaders headers) {
Map<String, String> environment = System.getenv();
Context context = new Context() {
@Override
public int getRemainingTimeInMillis() {
long now = System.currentTimeMillis();
if (!headers.containsKey("Lambda-Runtime-Deadline-Ms")) {
return 0;
}
int delta = (int) (Long.parseLong(headers.getFirst("Lambda-Runtime-Deadline-Ms")) - now);
return delta > 0 ? delta : 0;
}
@Override
public int getMemoryLimitInMB() {
if (!environment.containsKey("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")) {
return 128;
}
return Integer.parseInt(environment.getOrDefault("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "128"));
}
@Override
public LambdaLogger getLogger() {
return LambdaRuntime.getLogger();
}
@Override
public String getLogStreamName() {
return environment.get("LOG_STREAM_NAME");
}
@Override
public String getLogGroupName() {
return environment.get("LOG_GROUP_NAME");
}
@Override
public String getInvokedFunctionArn() {
return headers.getFirst("Lambda-Runtime-Invoked-Function-Arn");
}
@Override
public CognitoIdentity getIdentity() {
return null;
}
@Override
public String getFunctionVersion() {
return environment.get("FUNCTION_VERSION");
}
@Override
public String getFunctionName() {
return environment.get("FUNCTION_NAME");
}
@Override
public ClientContext getClientContext() {
return null;
}
@Override
public String getAwsRequestId() {
return headers.getFirst("Lambda-Runtime-Aws-Request-Id");
}
public String toString() {
return "FUNCTION NAME: " + getFunctionName() + ", FUNCTION VERSION: " + getFunctionVersion()
+ ", FUNCTION ARN: " + getInvokedFunctionArn() + ", FUNCTION MEM LIMIT: " + getMemoryLimitInMB()
+ ", FUNCTION DEADLINE: " + getRemainingTimeInMillis();
}
};
return context;
}
private void propagateAwsError(String requestId, Exception e, JsonMapper mapper, String runtimeApi, RestTemplate rest) {
String errorMessage = e.getMessage();
String errorType = e.getClass().getSimpleName();

View File

@@ -31,6 +31,9 @@ import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContrib
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.boot.http.client.ClientHttpRequestFactorySettings;
import org.springframework.boot.web.server.Ssl;
import org.springframework.boot.web.server.Ssl.ServerNameSslBundle;
import org.springframework.cloud.function.context.catalog.FunctionTypeUtils;
import org.springframework.cloud.function.context.config.FunctionContextUtils;
import org.springframework.cloud.function.context.message.MessageUtils;
@@ -105,6 +108,15 @@ public class FunctionTypeProcessor implements BeanFactoryInitializationAotProces
// known static types
runtimeHints.reflection().registerType(MessageUtils.MessageStructureWithCaseInsensitiveHeaderKeys.class,
MemberCategory.INVOKE_PUBLIC_METHODS);
// temporary due to bug in boot
runtimeHints.reflection().registerType(ClientHttpRequestFactorySettings.class,
MemberCategory.INVOKE_PUBLIC_METHODS);
runtimeHints.reflection().registerType(Ssl.class,
MemberCategory.INVOKE_PUBLIC_METHODS);
runtimeHints.reflection().registerType(ServerNameSslBundle.class,
MemberCategory.INVOKE_PUBLIC_METHODS);
}
}

View File

@@ -42,7 +42,11 @@
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter</artifactId>
</dependency>
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-lambda-java-core</artifactId>
<version>1.1.0</version>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>

View File

@@ -47,7 +47,6 @@
<groupId>com.amazonaws</groupId>
<artifactId>aws-lambda-java-core</artifactId>
<version>1.1.0</version>
<scope>provided</scope>
</dependency>
<dependency>

View File

@@ -7,10 +7,14 @@ import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.cloud.function.adapter.aws.AWSLambdaUtils;
import org.springframework.context.annotation.Bean;
// import org.springframework.cloud.function.context.DefaultMessageRoutingHandler;
// import org.springframework.cloud.function.context.MessageRoutingCallback;
// import org.springframework.messaging.Message;
import org.springframework.messaging.Message;
import com.amazonaws.services.lambda.runtime.Context;
@SpringBootApplication
public class NativeFunctionApplication {
@@ -33,10 +37,11 @@ public class NativeFunctionApplication {
// }
@Bean
public Function<String, String> uppercase() {
return v -> {
System.out.println("Uppercasing " + v);
return v.toUpperCase(Locale.ROOT);
public Function<Message<String>, String> uppercase() {
return message -> {
System.out.println("AWS Context: " + message.getHeaders().get(AWSLambdaUtils.AWS_CONTEXT));
System.out.println("Uppercasing " + message.getPayload());
return message.getPayload().toUpperCase(Locale.ROOT);
};
}