#5 - Propagate the X-Ray tracing header by setting _X_AMZN_TRACE_ID (#6)

Author: mfriesen

.classpath

@@ -13,6 +13,19 @@
 			<attribute name="test" value="true"/>
 		</attributes>
 	</classpathentry>
+	<classpathentry kind="src" output="bin/main" path="src/main/java">
+		<attributes>
+			<attribute name="gradle_scope" value="main"/>
+			<attribute name="gradle_used_by_scope" value="main,test"/>
+		</attributes>
+	</classpathentry>
+	<classpathentry kind="src" output="bin/test" path="src/test/java">
+		<attributes>
+			<attribute name="gradle_scope" value="test"/>
+			<attribute name="gradle_used_by_scope" value="test"/>
+			<attribute name="test" value="true"/>
+		</attributes>
+	</classpathentry>
 	<classpathentry kind="con" path="org.eclipse.jdt.launching.JRE_CONTAINER/org.eclipse.jdt.internal.debug.ui.launcher.StandardVMType/JavaSE-11/"/>
 	<classpathentry kind="con" path="org.eclipse.buildship.core.gradleclasspathcontainer"/>
 	<classpathentry kind="output" path="bin/default"/>

.settings/org.eclipse.buildship.core.prefs

@@ -1,2 +1,2 @@
-connection.project.dir=samples/helloworld
+connection.project.dir=
 eclipse.preferences.version=1

build.gradle

@@ -22,7 +22,7 @@ sourceCompatibility = JavaVersion.VERSION_11
 targetCompatibility = JavaVersion.VERSION_11
 
 group 'com.formkiq'
-version '2.1'
+version '2.2.0'
 
 spotless {
   java {
@@ -44,7 +44,7 @@ checkstyle {
 	toolVersion '8.29'
 	configFile file("config/checkstyle/checkstyle.xml")
 	configProperties = [project_loc: "${projectDir}"]
-	}
+}
 
 tasks.withType(Checkstyle).each { checkstyleTask ->
     checkstyleTask.doLast {
@@ -57,6 +57,8 @@ tasks.withType(Checkstyle).each { checkstyleTask ->
     }
 }
 
+checkstyleMain.dependsOn spotlessApply
+
 repositories {
     jcenter()
 }
@@ -91,6 +93,7 @@ afterEvaluate {
 dependencies {
 	implementation group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.0'
     implementation group: 'com.google.code.gson', name: 'gson', version: '2.8.6'
+	testImplementation group: 'com.amazonaws', name: 'aws-lambda-java-events', version: '3.1.0'
     testImplementation group: 'junit', name: 'junit', version:'4.+'
     testImplementation group: 'org.mock-server', name: 'mockserver-netty', version: '5.10.0'
     testImplementation group: 'org.slf4j', name: 'slf4j-simple', version: '1.7.30'

src/main/java/com/formkiq/lambda/runtime/graalvm/LambdaRuntime.java

@@ -24,6 +24,8 @@
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Parameter;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Type;
 import java.nio.charset.StandardCharsets;
 import java.text.MessageFormat;
 import java.util.HashMap;
@@ -183,17 +185,23 @@ private static void invokeClass(
     while (true) {
 
       // Get next Lambda Event
+      Context context = null;
       String eventBody = null;
       String requestId = UUID.randomUUID().toString();
 
       if (runtimeUrl != null) {
         HttpResponse event = HttpClient.get(runtimeUrl);
         requestId = event.getHeaderValue("Lambda-Runtime-Aws-Request-Id");
+
+        String xamazTraceId = event.getHeaderValue("Lambda-Runtime-Trace-Id");
+        if (xamazTraceId != null) {
+          System.setProperty("_X_AMZN_TRACE_ID", xamazTraceId);
+        }
+
+        context = new LambdaContext(requestId);
         eventBody = event.getBody();
       }
 
-      Context context = new LambdaContext(requestId);
-
       try {
         // Invoke Handler Method
         Object handler = clazz.getConstructor().newInstance();
@@ -266,18 +274,20 @@ private static String invokeLambdaRequestHandler(
    * @throws InvocationTargetException InvocationTargetException
    * @throws IllegalArgumentException IllegalArgumentException
    * @throws IllegalAccessException IllegalAccessException
+   * @throws ClassNotFoundException ClassNotFoundException
    */
   @SuppressWarnings({"unchecked", "rawtypes"})
   private static String invokeMethod(
       final Object object, final String methodName, final String payload, final Context context)
-      throws IllegalAccessException, IllegalArgumentException, InvocationTargetException {
+      throws IllegalAccessException, IllegalArgumentException, InvocationTargetException,
+          ClassNotFoundException {
 
     String val = "";
     Method method = findRequestHandlerMethod(object.getClass(), methodName);
-    Parameter parameter = method.getParameters()[0];
-
+    Class<?> parameterType = getParameterType(object, method);
     Gson gson = new GsonBuilder().create();
-    Object input = gson.fromJson(payload, parameter.getType());
+
+    Object input = gson.fromJson(payload, parameterType);
 
     Object value = null;
     if (methodName == null && object instanceof RequestHandler) {
@@ -300,6 +310,33 @@ private static String invokeMethod(
     return val;
   }
 
+  /**
+   * Get the Parameter Type of the Object.
+   *
+   * @param object {@link Object}
+   * @param method {@link Method}
+   * @return {@link Class}
+   * @throws ClassNotFoundException ClassNotFoundException
+   */
+  private static Class<?> getParameterType(final Object object, final Method method)
+      throws ClassNotFoundException {
+    Parameter parameter = method.getParameters()[0];
+    Class<?> parameterType = parameter.getType();
+
+    if (Object.class.equals(parameterType)) {
+
+      Type[] types = object.getClass().getGenericInterfaces();
+      if (types.length > 0 && types[0] instanceof ParameterizedType) {
+        ParameterizedType p = (ParameterizedType) types[0];
+        if (p.getActualTypeArguments().length > 0) {
+          parameterType = Class.forName(p.getActualTypeArguments()[0].getTypeName());
+        }
+      }
+    }
+
+    return parameterType;
+  }
+
   /**
    * Invoke {@link RequestStreamHandler}.
    *

src/test/java/com/formkiq/lambda/runtime/graalvm/InvocationNextHandler.java

@@ -21,19 +21,12 @@ public class InvocationNextHandler implements ExpectationResponseCallback {
 
   /** Lambda Request Id. */
   private static final String REQUEST_ID = "testrequestid";
+  /** Lambda Request Id. */
+  private static final String TRACE_ID = "testtraceid";
 
   /** Response Content. */
   private String responseContent = "test";
 
-  /**
-   * Get Response Content.
-   *
-   * @return {@link String}
-   */
-  public String getResponseContent() {
-    return this.responseContent;
-  }
-
   /**
    * Set Response Content.
    *
@@ -48,6 +41,8 @@ public HttpResponse handle(final HttpRequest httpRequest) throws Exception {
     final int statusCode = 200;
     return HttpResponse.response()
         .withHeader("Lambda-Runtime-Aws-Request-Id", REQUEST_ID)
+        .withHeader("Lambda-Runtime-Trace-Id", TRACE_ID)
+        .withBody(responseContent)
         .withStatusCode(Integer.valueOf(statusCode));
   }
 }

src/test/java/com/formkiq/lambda/runtime/graalvm/LambdaRuntimeTest.java

@@ -13,6 +13,7 @@
 package com.formkiq.lambda.runtime.graalvm;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.mockserver.integration.ClientAndServer.startClientAndServer;
 import static org.mockserver.model.HttpRequest.request;
 
@@ -286,12 +287,32 @@ public void testInvoke10() throws Exception {
       System.setProperty(e.getKey(), e.getValue());
     }
 
-    System.out.println("MAP: " + env);
     // when
     LambdaRuntime.main(new String[] {});
 
     // then
     String expected = "this is a run string";
     assertEquals(expected, invocationResponseHandler.getResponse());
   }
+
+  /**
+   * Test invoke Lambda with {@link TestRequestApiGatewayProxyHandler} special method and without
+   * environment variables.
+   *
+   * @throws Exception Exception
+   */
+  @Test
+  public void testInvoke11() throws Exception {
+    // given
+    invocationNextHandler.setResponseContent("{\"body\":\"this is some data\"}");
+    Map<String, String> env = createEnv(TestRequestApiGatewayProxyHandler.class.getName());
+
+    // when
+    LambdaRuntime.invoke(env);
+
+    // then
+    String expected = "{\"body\":\"this is some data\"}";
+    assertEquals(expected, invocationResponseHandler.getResponse());
+    assertNotNull(System.getProperty("_X_AMZN_TRACE_ID"));
+  }
 }

src/test/java/com/formkiq/lambda/runtime/graalvm/TestRequestApiGatewayProxyHandler.java

@@ -0,0 +1,32 @@
+/**
+ * Copyright [2020] FormKiQ Inc. Licensed under the Apache License, Version 2.0 (the "License"); you
+ * may not use this file except in compliance with the License. You may obtain a copy of the License
+ * at
+ *
+ * <p>http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * <p>Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.formkiq.lambda.runtime.graalvm;
+
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestHandler;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
+import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
+
+/** Test {@link RequestStreamHandler} and returns a {@link String}. */
+public class TestRequestApiGatewayProxyHandler
+    implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {
+
+  @Override
+  public APIGatewayProxyResponseEvent handleRequest(
+      final APIGatewayProxyRequestEvent input, final Context context) {
+    APIGatewayProxyResponseEvent e = new APIGatewayProxyResponseEvent();
+    e.setBody(input.getBody());
+    return e;
+  }
+}