Close push connection when UI is detached (#11596)

Change-Id: Ibfc5923406b386786ae399b7f53cea47ac885f48

4 files changed, 39 insertions, 3 deletions

diff --git a/server/src/com/vaadin/server/VaadinSession.java b/server/src/com/vaadin/server/VaadinSession.java
index e9b3eb6ada..9c803924e0 100644
--- a/server/src/com/vaadin/server/VaadinSession.java
+++ b/server/src/com/vaadin/server/VaadinSession.java
@@ -321,7 +321,7 @@ public class VaadinSession implements HttpSessionBindingListener, Serializable {
     @Deprecated
     public void removeFromSession(VaadinService service) {
         assert hasLock();
-        session.setAttribute(getSessionAttributeName(service), null);
+        session.removeAttribute(getSessionAttributeName(service));
     }
 
     /**
diff --git a/server/src/com/vaadin/server/communication/AtmospherePushConnection.java b/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
index a5025e2356..770b0b6a59 100644
--- a/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
+++ b/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
@@ -20,6 +20,7 @@ import java.io.IOException;
 import java.io.Serializable;
 import java.io.StringWriter;
 import java.io.Writer;
+import java.util.concurrent.Future;
 
 import org.atmosphere.cpr.AtmosphereResource;
 import org.json.JSONException;
@@ -37,6 +38,7 @@ public class AtmospherePushConnection implements Serializable, PushConnection {
 
     private UI ui;
     private transient AtmosphereResource resource;
+    private Future<String> lastMessage;
 
     public AtmospherePushConnection(UI ui) {
         this.ui = ui;
@@ -79,7 +81,8 @@ public class AtmospherePushConnection implements Serializable, PushConnection {
      *            The message to send
      */
     void sendMessage(String message) {
-        getResource().getBroadcaster().broadcast(message, getResource());
+        lastMessage = getResource().getBroadcaster().broadcast(message,
+                getResource());
     }
 
     /**
@@ -97,7 +100,8 @@ public class AtmospherePushConnection implements Serializable, PushConnection {
     /**
      * Returns whether this connection is currently open.
      */
-    protected boolean isConnected() {
+    @Override
+    public boolean isConnected() {
         return resource != null
                 && resource.getBroadcaster().getAtmosphereResources()
                         .contains(resource);
@@ -120,6 +124,17 @@ public class AtmospherePushConnection implements Serializable, PushConnection {
 
     @Override
     public void disconnect() {
+        if (lastMessage != null) {
+            try {
+                // Wait for the last message to be sent before closing the
+                // connection (assumes that futures are completed in order)
+                lastMessage.get();
+            } catch (Exception e) {
+                e.printStackTrace();
+            }
+            lastMessage = null;
+        }
+
         resource.resume();
         assert !resource.getBroadcaster().getAtmosphereResources()
                 .contains(resource);
diff --git a/server/src/com/vaadin/server/communication/PushConnection.java b/server/src/com/vaadin/server/communication/PushConnection.java
index eecf4d93a4..4e043f565f 100644
--- a/server/src/com/vaadin/server/communication/PushConnection.java
+++ b/server/src/com/vaadin/server/communication/PushConnection.java
@@ -40,4 +40,9 @@ public interface PushConnection {
      */
     public void disconnect();
 
+    /**
+     * Returns whether this connection is currently open.
+     */
+    public boolean isConnected();
+
 }
\ No newline at end of file
diff --git a/server/src/com/vaadin/ui/UI.java b/server/src/com/vaadin/ui/UI.java
index 961ed289f3..442edfebb1 100644
--- a/server/src/com/vaadin/ui/UI.java
+++ b/server/src/com/vaadin/ui/UI.java
@@ -42,6 +42,7 @@ import com.vaadin.shared.EventId;
 import com.vaadin.shared.MouseEventDetails;
 import com.vaadin.shared.communication.PushMode;
 import com.vaadin.shared.ui.ui.ScrollClientRpc;
+import com.vaadin.shared.ui.ui.UIClientRpc;
 import com.vaadin.shared.ui.ui.UIConstants;
 import com.vaadin.shared.ui.ui.UIServerRpc;
 import com.vaadin.shared.ui.ui.UIState;
@@ -363,6 +364,12 @@ public abstract class UI extends AbstractSingleComponentContainer implements
         } else {
             if (session == null) {
                 detach();
+                if (pushConnection != null && pushConnection.isConnected()) {
+                    // Close the push connection when UI is detached. Otherwise
+                    // the push connection and possibly VaadinSession will live
+                    // on.
+                    pushConnection.disconnect();
+                }
             }
             this.session = session;
         }
@@ -1006,6 +1013,15 @@ public abstract class UI extends AbstractSingleComponentContainer implements
      */
     public void close() {
         closing = true;
+
+        boolean sessionExpired = (session == null || session.isClosing());
+        getRpcProxy(UIClientRpc.class).uiClosed(sessionExpired);
+        if (getPushConnection() != null && getPushConnection().isConnected()) {
+            // Push the Rpc to the client. The connection will be closed when
+            // the UI is detached and cleaned up.
+            getPushConnection().push();
+        }
+
     }
 
     /**