From f9618ddeb649a6562d96d3de932f4d97a825834b Mon Sep 17 00:00:00 2001
From: Jan Gassen <jan.gassen@vorwerk.de>
Date: Thu, 13 Aug 2020 21:42:03 +0200
Subject: [PATCH] Add timeout to cached instances

---
 pom.xml                                       |  5 ++++
 .../java/de/jangassen/lambda/LambdaProxy.java |  4 ++-
 .../loader/MethodInvocationContextCache.java  | 27 +++++++++++++------
 .../DefaultLambdaMethodInvokerTest.java       |  3 ++-
 4 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/pom.xml b/pom.xml
index 720ba64..fe98016 100644
--- a/pom.xml
+++ b/pom.xml
@@ -188,5 +188,10 @@
             <version>${junit-jupiter.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.google.guava</groupId>
+            <artifactId>guava</artifactId>
+            <version>29.0-jre</version>
+        </dependency>
     </dependencies>
 </project>
\ No newline at end of file
diff --git a/src/main/java/de/jangassen/lambda/LambdaProxy.java b/src/main/java/de/jangassen/lambda/LambdaProxy.java
index 456375a..319043b 100644
--- a/src/main/java/de/jangassen/lambda/LambdaProxy.java
+++ b/src/main/java/de/jangassen/lambda/LambdaProxy.java
@@ -17,10 +17,12 @@
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.concurrent.Executors;
 
 public class LambdaProxy {
     public static final int PORT = 3000;
+    public static final Duration INSTANCE_TIMEOUT = Duration.ofMinutes(5);
 
     public static final String LAMBDA_PROXY = "lambdaProxy";
     public static final String TEMPLATE_YAML = "template.yaml";
@@ -58,7 +60,7 @@ private LambdaProxyServlet getLambdaServlet(Path rootPath, Path templateFile) th
         SamApiDescription apiDescription = new SamApiDescription(samTemplate, projectPath);
         LambdaClassLoaderFactory classLoaderFactory = new LambdaClassLoaderFactory(new SamArtifactResolver(rootPath));
         DefaultLambdaMethodInvoker lambdaMethodInvoker = new DefaultLambdaMethodInvoker();
-        MethodInvocationContextCache methodInvocationContextCache = new MethodInvocationContextCache(classLoaderFactory);
+        MethodInvocationContextCache methodInvocationContextCache = new MethodInvocationContextCache(classLoaderFactory, INSTANCE_TIMEOUT);
 
         return new LambdaProxyServlet(lambdaMethodInvoker, methodInvocationContextCache, getCorsSettings(samTemplate), apiDescription.getApiMethods());
     }
diff --git a/src/main/java/de/jangassen/lambda/loader/MethodInvocationContextCache.java b/src/main/java/de/jangassen/lambda/loader/MethodInvocationContextCache.java
index be8c7f6..b50114a 100644
--- a/src/main/java/de/jangassen/lambda/loader/MethodInvocationContextCache.java
+++ b/src/main/java/de/jangassen/lambda/loader/MethodInvocationContextCache.java
@@ -1,21 +1,32 @@
 package de.jangassen.lambda.loader;
 
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
 import de.jangassen.lambda.api.ApiResource;
 
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
+import javax.annotation.Nonnull;
+import java.time.Duration;
 
 public class MethodInvocationContextCache implements MethodInvocationContextProvider {
-    private final ClassLoaderFactory classLoaderFactory;
+    private final LoadingCache<ApiResource, MethodInvocationContext> cache;
 
-    private final Map<ApiResource, MethodInvocationContext> cache = new ConcurrentHashMap<>();
-
-    public MethodInvocationContextCache(ClassLoaderFactory classLoaderFactory) {
-        this.classLoaderFactory = classLoaderFactory;
+    public MethodInvocationContextCache(ClassLoaderFactory classLoaderFactory, Duration timeout) {
+        cache = CacheBuilder.newBuilder()
+                .expireAfterAccess(timeout)
+                .build(new CacheLoader<ApiResource, MethodInvocationContext>() {
+                    @Override
+                    public MethodInvocationContext load(@Nonnull ApiResource apiResource) {
+                        return MethodInvocationContextBuilder
+                                .start(classLoaderFactory)
+                                .withApiResource(apiResource)
+                                .build();
+                    }
+                });
     }
 
     @Override
     public MethodInvocationContext getMethodInvocationContext(ApiResource apiResource) {
-        return cache.computeIfAbsent(apiResource, key -> MethodInvocationContextBuilder.start(classLoaderFactory).withApiResource(key).build());
+        return cache.getUnchecked(apiResource);
     }
 }
diff --git a/src/test/java/de/jangassen/lambda/loader/DefaultLambdaMethodInvokerTest.java b/src/test/java/de/jangassen/lambda/loader/DefaultLambdaMethodInvokerTest.java
index bb981eb..0f7acfd 100644
--- a/src/test/java/de/jangassen/lambda/loader/DefaultLambdaMethodInvokerTest.java
+++ b/src/test/java/de/jangassen/lambda/loader/DefaultLambdaMethodInvokerTest.java
@@ -7,13 +7,14 @@
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
+import java.time.Duration;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 public class DefaultLambdaMethodInvokerTest {
     @Test
     public void testSimpleInvocationInSameClassLoader() throws IOException, ReflectiveOperationException {
-        MethodInvocationContextCache methodInvocationContextCache = new MethodInvocationContextCache(h -> DefaultLambdaMethodInvokerTest.class.getClassLoader());
+        MethodInvocationContextCache methodInvocationContextCache = new MethodInvocationContextCache(h -> DefaultLambdaMethodInvokerTest.class.getClassLoader(), Duration.ofSeconds(10));
         DefaultLambdaMethodInvoker defaultLambdaMethodInvoker = new DefaultLambdaMethodInvoker();
 
         Request req = createTestRequest();