Don't assign PushConnection to UI before it has been connected (#11506)

Change-Id: I728c830d6740f77a200ea69925772924e58f45a4

4 files changed, 43 insertions, 47 deletions

diff --git a/server/src/com/vaadin/server/communication/AtmospherePushConnection.java b/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
index 5ff1d14232..9b98153a23 100644
--- a/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
+++ b/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
@@ -36,7 +36,6 @@ import com.vaadin.ui.UI;
 public class AtmospherePushConnection implements Serializable, PushConnection {
 
     private UI ui;
-    private boolean pending = true;
     private transient AtmosphereResource resource;
 
     public AtmospherePushConnection(UI ui) {
@@ -45,16 +44,12 @@ public class AtmospherePushConnection implements Serializable, PushConnection {
 
     @Override
     public void push() {
-        if (!isConnected()) {
-            // Not currently connected; defer until connection established
-            setPending(true);
-        } else {
-            try {
-                push(true);
-            } catch (IOException e) {
-                // TODO Error handling
-                throw new RuntimeException("Push failed", e);
-            }
+        assert isConnected();
+        try {
+            push(true);
+        } catch (IOException e) {
+            // TODO Error handling
+            throw new RuntimeException("Push failed", e);
         }
     }
 
@@ -97,10 +92,6 @@ public class AtmospherePushConnection implements Serializable, PushConnection {
      */
     protected void connect(AtmosphereResource resource) throws IOException {
         this.resource = resource;
-        if (isPending()) {
-            push(true);
-            setPending(false);
-        }
     }
 
     /**
@@ -113,21 +104,6 @@ public class AtmospherePushConnection implements Serializable, PushConnection {
     }
 
     /**
-     * Marks that changes in the UI should be pushed as soon as a connection is
-     * established.
-     */
-    protected void setPending(boolean pending) {
-        this.pending = pending;
-    }
-
-    /**
-     * @return Whether the UI should be pushed as soon as a connection opens.
-     */
-    protected boolean isPending() {
-        return pending;
-    }
-
-    /**
      * @return the UI associated with this connection.
      */
     protected UI getUI() {
diff --git a/server/src/com/vaadin/server/communication/PushHandler.java b/server/src/com/vaadin/server/communication/PushHandler.java
index 6fb1ab7a09..5e61dd5c98 100644
--- a/server/src/com/vaadin/server/communication/PushHandler.java
+++ b/server/src/com/vaadin/server/communication/PushHandler.java
@@ -36,6 +36,7 @@ import com.vaadin.server.VaadinService;
 import com.vaadin.server.VaadinServletRequest;
 import com.vaadin.server.VaadinServletService;
 import com.vaadin.server.VaadinSession;
+import com.vaadin.shared.communication.PushMode;
 import com.vaadin.ui.UI;
 
 /**
@@ -88,9 +89,7 @@ public class PushHandler implements AtmosphereHandler {
                         "Could not find the requested UI in session");
                 return;
             }
-            assert ui.getPushConnection() instanceof AtmospherePushConnection;
-            AtmospherePushConnection connection = (AtmospherePushConnection) ui
-                    .getPushConnection();
+            assert session.getPushMode() != PushMode.DISABLED;
 
             if (req.getMethod().equalsIgnoreCase("GET")) {
                 /*
@@ -111,13 +110,20 @@ public class PushHandler implements AtmosphereHandler {
                 }
                 resource.suspend();
 
+                AtmospherePushConnection connection = new AtmospherePushConnection(
+                        ui);
                 connection.connect(resource);
+
+                ui.setPushConnection(connection);
             } else if (req.getMethod().equalsIgnoreCase("POST")) {
-                assert connection.isConnected() : "Got push from the client "
+                AtmospherePushConnection connection = getConnectionForUI(ui);
+
+                assert connection != null : "Got push from the client "
                         + "even though the connection does not seem to be "
-                        + "open. This might happen if a HttpSession is "
+                        + "valid. This might happen if a HttpSession is "
                         + "serialized and deserialized while the push "
-                        + "connection is kept open.";
+                        + "connection is kept open or if the UI has a "
+                        + "connection of unexpected type.";
 
                 /*
                  * We received a UIDL request through Atmosphere. If the push
@@ -156,6 +162,18 @@ public class PushHandler implements AtmosphereHandler {
         }
     }
 
+    private static AtmospherePushConnection getConnectionForUI(UI ui) {
+        PushConnection pushConnection = ui.getPushConnection();
+        if (pushConnection instanceof AtmospherePushConnection) {
+            AtmospherePushConnection apc = (AtmospherePushConnection) pushConnection;
+            if (apc.isConnected()) {
+                return apc;
+            }
+        }
+
+        return null;
+    }
+
     @Override
     public void onStateChange(AtmosphereResourceEvent event) throws IOException {
         AtmosphereResource resource = event.getResource();
diff --git a/server/src/com/vaadin/server/communication/UIInitHandler.java b/server/src/com/vaadin/server/communication/UIInitHandler.java
index e3d7264de9..efe11e02e9 100644
--- a/server/src/com/vaadin/server/communication/UIInitHandler.java
+++ b/server/src/com/vaadin/server/communication/UIInitHandler.java
@@ -37,7 +37,6 @@ import com.vaadin.server.VaadinRequest;
 import com.vaadin.server.VaadinResponse;
 import com.vaadin.server.VaadinService;
 import com.vaadin.server.VaadinSession;
-import com.vaadin.shared.communication.PushMode;
 import com.vaadin.shared.ui.ui.UIConstants;
 import com.vaadin.ui.UI;
 
@@ -206,14 +205,6 @@ public abstract class UIInitHandler extends SynchronizedRequestHandler {
         // Set thread local here so it is available in init
         UI.setCurrent(ui);
 
-        if (session.getPushMode() != PushMode.DISABLED) {
-            /*
-             * TODO This currently makes it very difficult to use any other push
-             * implementation than the bundled one.
-             */
-            ui.setPushConnection(new AtmospherePushConnection(ui));
-        }
-
         ui.doInit(request, uiId.intValue());
 
         session.addUI(ui);
diff --git a/server/src/com/vaadin/ui/UI.java b/server/src/com/vaadin/ui/UI.java
index 303dfc234a..1fb7ace9df 100644
--- a/server/src/com/vaadin/ui/UI.java
+++ b/server/src/com/vaadin/ui/UI.java
@@ -482,6 +482,8 @@ public abstract class UI extends AbstractSingleComponentContainer implements
 
     private PushConnection pushConnection = null;
 
+    private boolean hasPendingPush = false;
+
     /**
      * This method is used by Component.Focusable objects to request focus to
      * themselves. Focus renders must be handled at window level (instead of
@@ -1149,6 +1151,7 @@ public abstract class UI extends AbstractSingleComponentContainer implements
     public void push() {
         VaadinSession session = getSession();
         if (session != null) {
+            assert session.hasLock();
             if (!getConnectorTracker().hasDirtyConnectors()) {
                 // Do not push if there is nothing to push
                 return;
@@ -1157,8 +1160,12 @@ public abstract class UI extends AbstractSingleComponentContainer implements
             if (session.getPushMode() == PushMode.DISABLED) {
                 throw new IllegalStateException("Push not enabled");
             }
-            assert pushConnection != null;
-            pushConnection.push();
+
+            if (pushConnection == null) {
+                hasPendingPush = true;
+            } else {
+                pushConnection.push();
+            }
         } else {
             throw new UIDetachedException("Trying to push a detached UI");
         }
@@ -1180,6 +1187,10 @@ public abstract class UI extends AbstractSingleComponentContainer implements
         assert pushConnection == null;
         assert connection != null;
         pushConnection = connection;
+        if (hasPendingPush) {
+            hasPendingPush = false;
+            pushConnection.push();
+        }
     }
 
     /**