Clear thread local instances on connection lost in push handler (#12042)

Adopted from https://github.com/vaadin/flow/pull/8567

5 files changed, 175 insertions, 7 deletions

diff --git a/server/src/main/java/com/vaadin/server/communication/PushHandler.java b/server/src/main/java/com/vaadin/server/communication/PushHandler.java
index 030d4323ef..77dd1b9b57 100644
--- a/server/src/main/java/com/vaadin/server/communication/PushHandler.java
+++ b/server/src/main/java/com/vaadin/server/communication/PushHandler.java
@@ -44,6 +44,7 @@ import com.vaadin.shared.ApplicationConstants;
 import com.vaadin.shared.JsonConstants;
 import com.vaadin.shared.communication.PushMode;
 import com.vaadin.ui.UI;
+import com.vaadin.util.CurrentInstance;
 
 import elemental.json.JsonException;
 
@@ -316,13 +317,29 @@ public class PushHandler {
     }
 
     void connectionLost(AtmosphereResourceEvent event) {
+        VaadinSession session = null;
+        try {
+            session = handleConnectionLost(event);
+        } finally {
+            if (session != null) {
+                session.access(new Runnable() {
+                    @Override
+                    public void run() {
+                        CurrentInstance.clearAll();
+                    }
+                });
+            }
+        }
+    }
+
+    private VaadinSession handleConnectionLost(AtmosphereResourceEvent event) {    
         // We don't want to use callWithUi here, as it assumes there's a client
         // request active and does requestStart and requestEnd among other
         // things.
         if (event == null) {
             getLogger().log(Level.SEVERE,
                     "Could not get event. This should never happen.");
-            return;
+            return null;
         }
 
         AtmosphereResource resource = event.getResource();
@@ -330,7 +347,7 @@ public class PushHandler {
         if (resource == null) {
             getLogger().log(Level.SEVERE,
                     "Could not get resource. This should never happen.");
-            return;
+            return null;
         }
 
         VaadinServletRequest vaadinRequest = new VaadinServletRequest(
@@ -342,7 +359,7 @@ public class PushHandler {
         } catch (ServiceException e) {
             getLogger().log(Level.SEVERE,
                     "Could not get session. This should never happen", e);
-            return;
+            return null;
         } catch (SessionExpiredException e) {
             // This happens at least if the server is restarted without
             // preserving the session. After restart the client reconnects, gets
@@ -351,7 +368,7 @@ public class PushHandler {
             getLogger().log(Level.FINER,
                     "Session expired before push disconnect event was received",
                     e);
-            return;
+            return session;
         }
 
         UI ui = null;
@@ -375,13 +392,13 @@ public class PushHandler {
                     getLogger().log(Level.FINE,
                             "Could not get UI. This should never happen,"
                                     + " except when reloading in Firefox and Chrome -"
-                                    + " see http://dev.vaadin.com/ticket/14251.");
-                    return;
+                                    + " see https://github.com/vaadin/framework/issues/5449.");
+                    return session;
                 } else {
                     getLogger().log(Level.INFO,
                             "No UI was found based on data in the request,"
                                     + " but a slower lookup based on the AtmosphereResource succeeded."
-                                    + " See http://dev.vaadin.com/ticket/14251 for more details.");
+                                    + " See https://github.com/vaadin/framework/issues/5449 for more details.");
                 }
             }
 
@@ -426,6 +443,7 @@ public class PushHandler {
                 // can't call ErrorHandler, we (hopefully) don't have a lock
             }
         }
+        return session; 
     }
 
     private static UI findUiUsingResource(AtmosphereResource resource,
diff --git a/server/src/test/java/com/vaadin/server/MockVaadinServletService.java b/server/src/test/java/com/vaadin/server/MockVaadinServletService.java
new file mode 100644
index 0000000000..dee45e99e6
--- /dev/null
+++ b/server/src/test/java/com/vaadin/server/MockVaadinServletService.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright 2000-2020 Vaadin Ltd.
+ *
+ * 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
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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.vaadin.server;
+
+import java.util.Collections;
+import java.util.List;
+
+import javax.servlet.ServletException;
+
+import com.vaadin.tests.util.MockDeploymentConfiguration;
+
+/**
+ *
+ * @author Vaadin Ltd
+ */
+public class MockVaadinServletService extends VaadinServletService {
+
+    public MockVaadinServletService() throws ServiceException {
+        this(new MockDeploymentConfiguration());
+    }
+
+    public MockVaadinServletService(
+            DeploymentConfiguration deploymentConfiguration) throws ServiceException {
+        this(new VaadinServlet(), deploymentConfiguration);
+    }
+
+    public MockVaadinServletService(VaadinServlet servlet,
+            DeploymentConfiguration deploymentConfiguration) throws ServiceException {
+        super(servlet, deploymentConfiguration);
+
+        try {
+            servlet.init(new MockServletConfig());
+        } catch (ServletException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    protected List<RequestHandler> createRequestHandlers()
+            throws ServiceException {
+        return Collections.emptyList();
+    }
+
+    @Override
+    public void init() {
+        try {
+            super.init();
+        } catch (ServiceException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+}
diff --git a/server/src/test/java/com/vaadin/server/MockVaadinSession.java b/server/src/test/java/com/vaadin/server/MockVaadinSession.java
index 51a29adcbf..a95bf8b55f 100644
--- a/server/src/test/java/com/vaadin/server/MockVaadinSession.java
+++ b/server/src/test/java/com/vaadin/server/MockVaadinSession.java
@@ -18,6 +18,10 @@ public class MockVaadinSession extends VaadinSession {
         super(service);
     }
 
+    public MockVaadinSession() throws ServiceException {
+        super(new MockVaadinServletService());
+    }
+
     @Override
     public void close() {
         super.close();
diff --git a/server/src/test/java/com/vaadin/server/communication/PushHandlerTest b/server/src/test/java/com/vaadin/server/communication/PushHandlerTest
new file mode 100644
index 0000000000..7f2c0ffb98
--- /dev/null
+++ b/server/src/test/java/com/vaadin/server/communication/PushHandlerTest
@@ -0,0 +1,79 @@
+/*
+ * Copyright 2000-2020 Vaadin Ltd.
+ *
+ * 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
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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.vaadin.server.communication;
+
+import java.util.concurrent.Future;
+
+import org.atmosphere.cpr.AtmosphereRequest;
+import org.atmosphere.cpr.AtmosphereResource;
+import org.atmosphere.cpr.AtmosphereResourceEvent;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import com.vaadin.server.MockVaadinServletService;
+import com.vaadin.server.MockVaadinSession;
+import com.vaadin.server.ServiceException;
+import com.vaadin.server.SessionExpiredException;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.server.VaadinSession;
+import com.vaadin.ui.UI;
+
+public class PushHandlerTest {
+
+    MockVaadinSession session = null;
+
+    @Test
+    public void connectionLost_currentInstancesAreCleared()
+            throws SessionExpiredException, ServiceException {
+        session = new MockVaadinSession() {
+            @Override
+            public Future<Void> access(Runnable runnable) {
+                runnable.run();
+                return Mockito.mock(Future.class);
+            }
+        };
+        VaadinSession.setCurrent(session);
+        Assert.assertNotNull(VaadinSession.getCurrent());
+        MockVaadinServletService service = null;
+        service = new MockVaadinServletService() {
+            @Override
+            public com.vaadin.server.VaadinSession findVaadinSession(
+                    VaadinRequest request) throws SessionExpiredException {
+                 return session;
+            }
+
+            @Override
+            public UI findUI(VaadinRequest request) {
+                return null;
+            }
+        };
+
+        service.init();
+        PushHandler handler = new PushHandler(service);
+
+        AtmosphereResource resource = Mockito.mock(AtmosphereResource.class);
+        AtmosphereRequest request = Mockito.mock(AtmosphereRequest.class);
+        Mockito.when(resource.getRequest()).thenReturn(request);
+
+        AtmosphereResourceEvent event = Mockito
+                .mock(AtmosphereResourceEvent.class);
+        Mockito.when(event.getResource()).thenReturn(resource);
+        handler.connectionLost(event);
+
+        Assert.assertNull(VaadinSession.getCurrent());
+    }
+}
diff --git a/server/src/test/java/com/vaadin/tests/server/ClassesSerializableTest.java b/server/src/test/java/com/vaadin/tests/server/ClassesSerializableTest.java
index 9d921e1a52..6626aa6046 100644
--- a/server/src/test/java/com/vaadin/tests/server/ClassesSerializableTest.java
+++ b/server/src/test/java/com/vaadin/tests/server/ClassesSerializableTest.java
@@ -61,6 +61,7 @@ public class ClassesSerializableTest {
             "com\\.vaadin\\.server\\.VaadinPortlet", //
             "com\\.vaadin\\.server\\.MockServletConfig", //
             "com\\.vaadin\\.server\\.MockServletContext", //
+            "com\\.vaadin\\.server\\.MockVaadinServletService", //
             "com\\.vaadin\\.server\\.Constants", //
             "com\\.vaadin\\.server\\.VaadinServiceClassLoaderUtil", //
             "com\\.vaadin\\.server\\.VaadinServiceClassLoaderUtil\\$GetClassLoaderPrivilegedAction", //