Set current instances when calling UI.push from VaadinSession.unlock (#12168)

Change-Id: I27795ab9ae3e3692f508e847936ccaa5a1ebadd4

server/src/com/vaadin/server/VaadinSession.java

@@ -885,9 +885,9 @@ public class VaadinSession implements HttpSessionBindingListener, Serializable {
      * Unlocks this session. This method should always be used in a finally
      * block after {@link #lock()} to ensure that the lock is always released.
      * <p>
-     * If {@link #getPushMode() the push mode} is {@link PushMode#AUTOMATIC
-     * automatic}, pushes the changes in all UIs in this session to their
-     * respective clients.
+     * For UIs in this session that have its push mode set to
+     * {@link PushMode#AUTOMATIC automatic}, pending changes will be pushed to
+     * their respective clients.
      * 
      * @see #lock()
      * @see UI#push()
@@ -904,7 +904,13 @@ public class VaadinSession implements HttpSessionBindingListener, Serializable {
 
                 for (UI ui : getUIs()) {
                     if (ui.getPushConfiguration().getPushMode() == PushMode.AUTOMATIC) {
-                        ui.push();
+                        Map<Class<?>, CurrentInstance> oldCurrent = CurrentInstance
+                                .setCurrent(ui);
+                        try {
+                            ui.push();
+                        } finally {
+                            CurrentInstance.restoreInstances(oldCurrent);
+                        }
                     }
                 }
             }

server/src/com/vaadin/ui/UI.java

@@ -1294,15 +1294,18 @@ public abstract class UI extends AbstractSingleComponentContainer implements
      * Pushes the pending changes and client RPC invocations of this UI to the
      * client-side.
      * <p>
-     * As with all UI methods, it is not safe to call push() without holding the
-     * {@link VaadinSession#lock() session lock}.
+     * As with all UI methods, the session must be locked when calling this
+     * method. It is also recommended that {@link UI#getCurrent()} is set up to
+     * return this UI since writing the response may invoke logic in any
+     * attached component or extension. The recommended way of fulfilling these
+     * conditions is to use {@link #access(Runnable)}.
      * 
      * @throws IllegalStateException
      *             if push is disabled.
      * @throws UIDetachedException
      *             if this UI is not attached to a session.
      * 
-     * @see #getPushMode()
+     * @see #getPushConfiguration()
      * 
      * @since 7.1
      */

uitest/src/com/vaadin/tests/components/ui/UiAccess.html

@@ -161,6 +161,32 @@
 	<td>vaadin=runcomvaadintestscomponentsuiUiAccess::PID_SLog_row_0</td>
 	<td>3. Test value after access: Set before run pending</td>
 </tr>
+<!-- Run last part with push enabled -->
+<tr>
+	<td>open</td>
+	<td>/run/com.vaadin.tests.components.ui.UiAccess?restartApplication&amp;transport=websocket</td>
+	<td></td>
+</tr>
+<tr>
+	<td>click</td>
+	<td>vaadin=runcomvaadintestscomponentsuiUiAccess::/VVerticalLayout[0]/Slot[2]/VVerticalLayout[0]/Slot[7]/VButton[0]/domChild[0]/domChild[0]</td>
+	<td></td>
+</tr>
+<tr>
+	<td>waitForNotText</td>
+	<td>vaadin=runcomvaadintestscomponentsuiUiAccess::PID_SLog_row_0</td>
+	<td></td>
+</tr>
+<tr>
+	<td>assertText</td>
+	<td>vaadin=runcomvaadintestscomponentsuiUiAccess::PID_SLog_row_0</td>
+	<td>exact:1. Current session matches in beforeResponse? true</td>
+</tr>
+<tr>
+	<td>assertText</td>
+	<td>vaadin=runcomvaadintestscomponentsuiUiAccess::PID_SLog_row_1</td>
+	<td>exact:0. Current UI matches in beforeResponse? true</td>
+</tr>
 </tbody></table>
 </body>
 </html>

uitest/src/com/vaadin/tests/components/ui/UiAccess.java

@@ -23,13 +23,18 @@ import java.util.concurrent.locks.ReentrantLock;
 
 import com.vaadin.server.VaadinRequest;
 import com.vaadin.server.VaadinService;
+import com.vaadin.server.VaadinSession;
+import com.vaadin.shared.communication.PushMode;
 import com.vaadin.tests.components.AbstractTestUIWithLog;
 import com.vaadin.ui.Button;
 import com.vaadin.ui.Button.ClickEvent;
+import com.vaadin.ui.UI;
 import com.vaadin.util.CurrentInstance;
 
 public class UiAccess extends AbstractTestUIWithLog {
 
+    private volatile boolean checkCurrentInstancesBeforeResponse = false;
+
     private Future<Void> checkFromBeforeClientResponse;
 
     private class CurrentInstanceTestType {
@@ -283,6 +288,46 @@ public class UiAccess extends AbstractTestUIWithLog {
                                         .get(CurrentInstanceTestType.class));
                     }
                 }));
+
+        addComponent(new Button("CurrentInstance when pushing",
+                new Button.ClickListener() {
+                    @Override
+                    public void buttonClick(ClickEvent event) {
+                        log.clear();
+                        if (getPushConfiguration().getPushMode() != PushMode.AUTOMATIC) {
+                            log("Can only test with automatic push enabled");
+                            return;
+                        }
+
+                        final VaadinSession session = getSession();
+                        new Thread() {
+                            @Override
+                            public void run() {
+                                // Pretend this isn't a Vaadin thread
+                                CurrentInstance.clearAll();
+
+                                /*
+                                 * Get explicit lock to ensure the (implicit)
+                                 * push does not happen during normal request
+                                 * handling.
+                                 */
+                                session.lock();
+                                try {
+                                    access(new Runnable() {
+                                        @Override
+                                        public void run() {
+                                            checkCurrentInstancesBeforeResponse = true;
+                                            // Trigger beforeClientResponse
+                                            markAsDirty();
+                                        }
+                                    });
+                                } finally {
+                                    session.unlock();
+                                }
+                            }
+                        }.start();
+                    }
+                }));
     }
 
     @Override
@@ -292,6 +337,15 @@ public class UiAccess extends AbstractTestUIWithLog {
                     + checkFromBeforeClientResponse.isDone());
             checkFromBeforeClientResponse = null;
         }
+        if (checkCurrentInstancesBeforeResponse) {
+            UI currentUI = UI.getCurrent();
+            VaadinSession currentSession = VaadinSession.getCurrent();
+
+            log("Current UI matches in beforeResponse? " + (currentUI == this));
+            log("Current session matches in beforeResponse? "
+                    + (currentSession == getSession()));
+            checkCurrentInstancesBeforeResponse = false;
+        }
     }
 
     @Override