Fallback to finding disconnected UI based on AtmosphereResource (#14251)

Change-Id: Icdac51322a90c32c122a182bc692c4eff3d8285b

server/src/com/vaadin/server/communication/PushHandler.java

@@ -18,6 +18,7 @@ package com.vaadin.server.communication;
 
 import java.io.IOException;
 import java.io.Reader;
+import java.util.Collection;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -48,7 +49,7 @@ import com.vaadin.ui.UI;
 
 /**
  * Establishes bidirectional ("push") communication channels
- *
+ * 
  * @author Vaadin Ltd
  * @since 7.1
  */
@@ -194,7 +195,7 @@ public class PushHandler extends AtmosphereResourceEventListenerAdapter {
 
     /**
      * Find the UI for the atmosphere resource, lock it and invoke the callback.
-     *
+     * 
      * @param resource
      *            the atmosphere resource for the current request
      * @param callback
@@ -357,9 +358,31 @@ public class PushHandler extends AtmosphereResourceEventListenerAdapter {
             // Sets UI.currentInstance
             ui = service.findUI(vaadinRequest);
             if (ui == null) {
-                getLogger().log(Level.SEVERE,
-                        "Could not get UI. This should never happen");
-                return;
+                /*
+                 * UI not found, could be because FF has asynchronously closed
+                 * the websocket connection and Atmosphere has already done
+                 * cleanup of the request attributes.
+                 * 
+                 * In that case, we still have a chance of finding the right UI
+                 * by iterating through the UIs in the session looking for one
+                 * using the same AtmosphereResource.
+                 */
+                ui = findUiUsingResource(resource, session.getUIs());
+
+                if (ui == null) {
+                    getLogger()
+                            .log(Level.SEVERE,
+                                    "Could not get UI. This should never happen,"
+                                            + " except when reloading in Firefox -"
+                                            + " see http://dev.vaadin.com/ticket/14251.");
+                    return;
+                } 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.");
+                }
             }
 
             PushMode pushMode = ui.getPushConfiguration().getPushMode();
@@ -411,6 +434,19 @@ public class PushHandler extends AtmosphereResourceEventListenerAdapter {
         }
     }
 
+    private static UI findUiUsingResource(AtmosphereResource resource,
+            Collection<UI> uIs) {
+        for (UI ui : uIs) {
+            PushConnection pushConnection = ui.getPushConnection();
+            if (pushConnection instanceof AtmospherePushConnection) {
+                if (((AtmospherePushConnection) pushConnection).getResource() == resource) {
+                    return ui;
+                }
+            }
+        }
+        return null;
+    }
+
     /**
      * Sends a refresh message to the given atmosphere resource. Uses an
      * AtmosphereResource instead of an AtmospherePushConnection even though it
@@ -419,10 +455,10 @@ public class PushHandler extends AtmosphereResourceEventListenerAdapter {
      * two push connections which try to use the same UI. Using the
      * AtmosphereResource directly guarantees the message goes to the correct
      * recipient.
-     *
+     * 
      * @param resource
      *            The atmosphere resource to send refresh to
-     *
+     * 
      */
     private static void sendRefreshAndDisconnect(AtmosphereResource resource)
             throws IOException {

uitest/src/com/vaadin/tests/push/RefreshCloseConnection.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2000-2014 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.tests.push;
+
+import com.vaadin.annotations.PreserveOnRefresh;
+import com.vaadin.annotations.Push;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUIWithLog;
+
+@Push
+@PreserveOnRefresh
+public class RefreshCloseConnection extends AbstractTestUIWithLog {
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        log("Init");
+    }
+
+    @Override
+    protected void refresh(VaadinRequest request) {
+        if (getPushConnection().isConnected()) {
+            log("Still connected");
+        }
+        log("Refresh");
+        new Thread() {
+            @Override
+            public void run() {
+                accessSynchronously(new Runnable() {
+                    @Override
+                    public void run() {
+                        log("Push");
+                    }
+                });
+            }
+        }.start();
+    }
+
+    @Override
+    protected String getTestDescription() {
+        return "A log row should get pushed after reloading the page";
+    }
+
+    @Override
+    protected Integer getTicketNumber() {
+        return Integer.valueOf(14251);
+    }
+
+}

uitest/src/com/vaadin/tests/push/RefreshCloseConnectionTest.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2000-2014 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.tests.push;
+
+import java.util.List;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.openqa.selenium.remote.DesiredCapabilities;
+
+import com.vaadin.tests.tb3.MultiBrowserTest;
+import com.vaadin.tests.tb3.WebsocketTest;
+
+public class RefreshCloseConnectionTest extends MultiBrowserTest {
+    @Test
+    public void testSessionRefresh() {
+        openTestURL();
+
+        Assert.assertEquals("1. Init", getLogRow(0));
+
+        openTestURL();
+
+        Assert.assertEquals("2. Refresh", getLogRow(1));
+        Assert.assertEquals("3. Push", getLogRow(0));
+    }
+
+    @Override
+    public List<DesiredCapabilities> getBrowsersToTest() {
+        return WebsocketTest.getWebsocketBrowsers();
+    }
+}