GH-1187 Fix AWS Context initialization for Custom Runtime

Updated sample to show that Context is not null

Resolves #1187

Author: olegz

spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/CustomRuntimeEventLoop.java

@@ -28,10 +28,14 @@
 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 @@ private void eventLoop(ConfigurableApplicationContext context) {
 			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 @@ private void eventLoop(ConfigurableApplicationContext context) {
 					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 @@ private void eventLoop(ConfigurableApplicationContext context) {
 					}
 				}
 				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();

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/FunctionTypeProcessor.java

@@ -31,6 +31,9 @@
 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 void applyTo(GenerationContext generationContext, BeanFactoryInitializati
 			// 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);
 		}
 
 	}

spring-cloud-function-samples/function-sample-aws-custom/pom.xml

@@ -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>

spring-cloud-function-samples/function-sample-aws-native/pom.xml

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

spring-cloud-function-samples/function-sample-aws-native/src/main/java/com/example/demo/NativeFunctionApplication.java

@@ -7,10 +7,14 @@
 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 static void main(String[] args) {
 //	}
 
 	@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);
 		};
 	}