Fix AbstractClientConnector stateTypeCache memory leak (#9896)

Fixes #9883

2 files changed, 75 insertions, 7 deletions

diff --git a/server/src/main/java/com/vaadin/server/AbstractClientConnector.java b/server/src/main/java/com/vaadin/server/AbstractClientConnector.java
index fc05f67860..71ec5bfbb4 100644
--- a/server/src/main/java/com/vaadin/server/AbstractClientConnector.java
+++ b/server/src/main/java/com/vaadin/server/AbstractClientConnector.java
@@ -30,7 +30,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.WeakHashMap;
 import java.util.logging.Logger;
 
 import com.vaadin.event.EventRouter;
@@ -95,7 +95,13 @@ public abstract class AbstractClientConnector
 
     private ErrorHandler errorHandler = null;
 
-    private static final ConcurrentHashMap<Class<? extends AbstractClientConnector>, Class<? extends SharedState>> stateTypeCache = new ConcurrentHashMap<>();
+    /**
+     * Static cache mapping AbstractClientConnector classes to their respective
+     * ShareState classes. Using WeakHashMap since entries are recalculated on
+     * demand.
+     */
+    private static final Map<Class<? extends AbstractClientConnector>, Class<? extends SharedState>> stateTypeCache = Collections
+            .synchronizedMap(new WeakHashMap<>());
 
     @Override
     public Registration addAttachListener(AttachListener listener) {
@@ -310,11 +316,8 @@ public abstract class AbstractClientConnector
         // exceptions flying around
         if (stateType == null) {
             // Cache because we don't need to do this once per instance
-            stateType = stateTypeCache.get(this.getClass());
-            if (stateType == null) {
-                stateType = findStateType();
-                stateTypeCache.put(this.getClass(), stateType);
-            }
+            stateType = stateTypeCache.computeIfAbsent(this.getClass(),
+                    key -> findStateType());
         }
 
         return stateType;
diff --git a/server/src/test/java/com/vaadin/server/AbstractClientConnectorTest.java b/server/src/test/java/com/vaadin/server/AbstractClientConnectorTest.java
index d8c4be4c9e..dbf36377c2 100644
--- a/server/src/test/java/com/vaadin/server/AbstractClientConnectorTest.java
+++ b/server/src/test/java/com/vaadin/server/AbstractClientConnectorTest.java
@@ -19,6 +19,15 @@ import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.lang.ref.WeakReference;
+import java.lang.reflect.Field;
+import java.net.URL;
+import java.util.Map;
+
+import org.apache.commons.io.IOUtils;
 import org.junit.Assert;
 import org.junit.Test;
 import org.mockito.Mockito;
@@ -70,6 +79,42 @@ public class AbstractClientConnectorTest {
         verify(mock, times(1)).registerRpc(implementation, ClickRpc.class);
     }
 
+    @Test
+    public void stateTypeCacheDoesNotLeakMemory()
+            throws IllegalArgumentException, IllegalAccessException,
+            NoSuchFieldException, SecurityException, InterruptedException,
+            ClassNotFoundException {
+        Field stateTypeCacheField = AbstractClientConnector.class
+                .getDeclaredField("stateTypeCache");
+        stateTypeCacheField.setAccessible(true);
+        Map<Class<?>, ?> stateTypeCache = (Map<Class<?>, ?>) stateTypeCacheField
+                .get(null);
+
+        WeakReference<Class<?>> classRef = loadClass(
+                "com.vaadin.server.AbstractClientConnector");
+        stateTypeCache.put(classRef.get(), null);
+        int size = stateTypeCache.size();
+        Assert.assertNotNull("Class should not yet be garbage collected",
+                classRef.get());
+
+        for (int i = 0; i < 100; ++i) {
+            System.gc();
+            if (stateTypeCache.size() < size) {
+                break;
+            }
+            Thread.sleep(100);
+        }
+        Assert.assertTrue(stateTypeCache.size() < size);
+        Assert.assertNull("Class should be garbage collected", classRef.get());
+    }
+
+    private WeakReference<Class<?>> loadClass(String name)
+            throws ClassNotFoundException {
+        ClassLoader loader = new TestClassLoader();
+        Class<?> loaded = loader.loadClass(name);
+        return new WeakReference<>(loaded);
+    }
+
     private class ServerRpcLastMock
             implements Comparable<ServerRpcLastMock>, ClickRpc {
         private static final long serialVersionUID = -2822356895755286180L;
@@ -110,4 +155,24 @@ public class AbstractClientConnectorTest {
 
     }
 
+    private static class TestClassLoader extends ClassLoader {
+
+        @Override
+        public Class<?> loadClass(String name) throws ClassNotFoundException {
+            if (!name.startsWith("com.vaadin.")) {
+                return super.loadClass(name);
+            }
+            String path = name.replaceAll("\\.", File.separator)
+                    .concat(".class");
+            URL resource = Thread.currentThread().getContextClassLoader()
+                    .getResource(path);
+            try (InputStream stream = resource.openStream()) {
+                byte[] bytes = IOUtils.toByteArray(stream);
+                return defineClass(name, bytes, 0, bytes.length);
+            } catch (IOException e) {
+                throw new ClassNotFoundException();
+            }
+        }
+    }
+
 }