fixes #90 Update router middleware to call getFinalizedRequest to complete the request chain

Author: stevehu

src/main/java/com/networknt/aws/lambda/handler/middleware/metrics/APMMetricsMiddleware.java

@@ -83,12 +83,12 @@ public Status execute(LightLambdaExchange exchange) {
 
         exchange.addResponseCompleteListener(finalExchange -> {
             Map<String, Object> auditInfo = (Map<String, Object>)finalExchange.getAttachment(AUDIT_ATTACHMENT_KEY);
-            if(logger.isTraceEnabled()) logger.trace("auditInfo = " + auditInfo);
+            if(logger.isTraceEnabled()) logger.trace("auditInfo = {}", auditInfo);
             if (auditInfo != null && !auditInfo.isEmpty()) {
                 Map<String, String> tags = new HashMap<>();
                 tags.put("endpoint", (String) auditInfo.get(Constants.ENDPOINT_STRING));
                 String clientId = auditInfo.get(Constants.CLIENT_ID_STRING) != null ? (String) auditInfo.get(Constants.CLIENT_ID_STRING) : "unknown";
-                if(logger.isTraceEnabled()) logger.trace("clientId = " + clientId);
+                if(logger.isTraceEnabled()) logger.trace("clientId = {}", clientId);
                 tags.put("clientId", clientId);
                 // scope client id will only be available if two token is used. For example, authorization code flow.
                 if (config.isSendScopeClientId()) {
@@ -120,7 +120,8 @@ public Status execute(LightLambdaExchange exchange) {
                 metricName = metricName.tagged(tags);
                 long time = Clock.defaultClock().getTick() - startTime;
                 registry.getOrAdd(metricName, MetricRegistry.MetricBuilder.TIMERS).update(time, TimeUnit.NANOSECONDS);
-                if(logger.isTraceEnabled()) logger.trace("metricName = " + metricName  + " commonTags = " + JsonMapper.toJson(commonTags) + " tags = " + JsonMapper.toJson(tags));
+                if(logger.isTraceEnabled())
+                    logger.trace("metricName = {} commonTags = {} tags = {}", metricName, JsonMapper.toJson(commonTags), JsonMapper.toJson(tags));
                 incCounterForStatusCode(finalExchange.getFinalizedResponse(true).getStatusCode(), commonTags, tags);
             } else {
                 // when we reach here, it will be in light-gateway so no specification is loaded on the server and also the security verification is failed.

src/main/java/com/networknt/aws/lambda/handler/middleware/router/LambdaRouterMiddleware.java

@@ -1,5 +1,6 @@
 package com.networknt.aws.lambda.handler.middleware.router;
 
+import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
 import com.networknt.aws.lambda.LightLambdaExchange;
 import com.networknt.aws.lambda.handler.Handler;
@@ -63,11 +64,13 @@ public Status execute(LightLambdaExchange exchange) {
             return this.successMiddlewareStatus();
         } else {
             if (!exchange.hasFailedState()) {
+                // get the finalized request to trigger the state change for the request complete.
+                APIGatewayProxyRequestEvent requestEvent = exchange.getFinalizedRequest(false);
                 /* invoke http service */
                 String serviceId = serviceIdOptional.get();
-                var originalPath = exchange.getRequest().getPath();
+                var originalPath = requestEvent.getPath();
                 var targetPath = originalPath;
-                var method = exchange.getRequest().getHttpMethod().toLowerCase();
+                var method = requestEvent.getHttpMethod().toLowerCase();
                 LOG.debug("Request path: {} -- Request method: {} -- Start time: {}", originalPath, method, System.currentTimeMillis());
                 // lookup the host from the serviceId
                 String host = cluster.serviceToUrl(protocol, serviceId, null, null);
@@ -89,8 +92,8 @@ public Status execute(LightLambdaExchange exchange) {
                 if("get".equalsIgnoreCase(method) || "delete".equalsIgnoreCase(method)) {
                     HttpClientRequest request = new HttpClientRequest();
                     try {
-                        HttpRequest.Builder builder = request.initBuilder(host + targetPath, HttpMethod.valueOf(exchange.getRequest().getHttpMethod()));
-                        exchange.getRequest().getHeaders().forEach(builder::header);
+                        HttpRequest.Builder builder = request.initBuilder(host + targetPath, HttpMethod.valueOf(requestEvent.getHttpMethod()));
+                        requestEvent.getHeaders().forEach(builder::header);
                         builder.timeout(Duration.ofMillis(CONFIG.getMaxRequestTime()));
                         HttpResponse<String> response = (HttpResponse<String>) request.send(builder, HttpResponse.BodyHandlers.ofString());
                         APIGatewayProxyResponseEvent res = new APIGatewayProxyResponseEvent()
@@ -113,8 +116,8 @@ public Status execute(LightLambdaExchange exchange) {
                 } else if("post".equalsIgnoreCase(method) || "put".equalsIgnoreCase(method) || "patch".equalsIgnoreCase(method)) {
                     HttpClientRequest request = new HttpClientRequest();
                     try {
-                        HttpRequest.Builder builder = request.initBuilder(host + targetPath, HttpMethod.valueOf(exchange.getRequest().getHttpMethod()), Optional.of(exchange.getRequest().getBody()));
-                        exchange.getRequest().getHeaders().forEach(builder::header);
+                        HttpRequest.Builder builder = request.initBuilder(host + targetPath, HttpMethod.valueOf(requestEvent.getHttpMethod()), Optional.of(requestEvent.getBody()));
+                        requestEvent.getHeaders().forEach(builder::header);
                         builder.timeout(Duration.ofMillis(CONFIG.getMaxRequestTime()));
                         HttpResponse<String> response = (HttpResponse<String>) request.send(builder, HttpResponse.BodyHandlers.ofString());
                         APIGatewayProxyResponseEvent res = new APIGatewayProxyResponseEvent()

src/test/java/com/networknt/aws/lambda/LightLambdaExchangeTest.java

@@ -126,4 +126,52 @@ void exchangeCompleteListener() {
         Assertions.assertNotNull(exchange.getAttachment(TEST_ATTACHMENT));
     }
 
+    @Test
+    public void testExchangeState() {
+        var apiGatewayProxyRequestEvent = TestUtils.createTestRequestEvent();
+        apiGatewayProxyRequestEvent.setPath("/v1/pets");
+        Map<String, String> headerMap = new HashMap<>();
+        headerMap.put("x-Traceability-id", "abc");
+        apiGatewayProxyRequestEvent.setHeaders(headerMap);
+        // set request body
+        apiGatewayProxyRequestEvent.setBody("{\"id\": 1, \"name\": \"dog\"}");
+        InvocationResponse invocation = InvocationResponse.builder()
+                .requestId("12345")
+                .event(apiGatewayProxyRequestEvent)
+                .build();
+        APIGatewayProxyRequestEvent requestEvent = invocation.getEvent();
+        Context lambdaContext = new LambdaContext(invocation.getRequestId());
+
+        var testAsynchronousMiddleware = new TestAsynchronousMiddleware();
+        var testInvocationHandler = new TestInvocationHandler();
+
+        Chain requestChain = new Chain(false);
+
+        requestChain.addChainable(testAsynchronousMiddleware);
+        requestChain.addChainable(testAsynchronousMiddleware);
+        requestChain.addChainable(testAsynchronousMiddleware);
+        requestChain.addChainable(testInvocationHandler);
+        requestChain.addChainable(testAsynchronousMiddleware);
+        requestChain.addChainable(testAsynchronousMiddleware);
+        requestChain.addChainable(testAsynchronousMiddleware);
+        requestChain.setupGroupedChain();
+
+        LightLambdaExchange exchange = new LightLambdaExchange(lambdaContext, requestChain);
+
+        /* Request should NOT be in progress without setting the initial request data. */
+        Assertions.assertFalse(exchange.isRequestInProgress());
+
+
+        /* Request should be in progress with the initial request data set. */
+        exchange.setInitialRequest(requestEvent);
+        Assertions.assertTrue(exchange.isRequestInProgress());
+        exchange.getFinalizedRequest(false);
+        // set the response
+        APIGatewayProxyResponseEvent response = new APIGatewayProxyResponseEvent();
+        response.setStatusCode(200);
+        exchange.setInitialResponse(response);
+        Assertions.assertEquals(22, exchange.getState());
+        Assertions.assertTrue(exchange.isRequestComplete());
+
+    }
 }