Convenience methods for running code in the UI or VaadinSession context (#11219)

Change-Id: If31a965f925ca2bedb25c712b83ccb070a9e71a0

4 files changed, 188 insertions, 2 deletions

diff --git a/server/src/com/vaadin/server/VaadinService.java b/server/src/com/vaadin/server/VaadinService.java
index 9be1c64880..e5d72969c6 100644
--- a/server/src/com/vaadin/server/VaadinService.java
+++ b/server/src/com/vaadin/server/VaadinService.java
@@ -579,12 +579,21 @@ public abstract class VaadinService implements Serializable {
      */
     public void setCurrentInstances(VaadinRequest request,
             VaadinResponse response) {
-        CurrentInstance.setInheritable(VaadinService.class, this);
+        setCurrent(this);
         CurrentInstance.set(VaadinRequest.class, request);
         CurrentInstance.set(VaadinResponse.class, response);
     }
 
     /**
+     * Sets the given Vaadin service as the current service.
+     * 
+     * @param service
+     */
+    public static void setCurrent(VaadinService service) {
+        CurrentInstance.setInheritable(VaadinService.class, service);
+    }
+
+    /**
      * Gets the currently processed Vaadin request. The current request is
      * automatically defined when the request is started. The current request
      * can not be used in e.g. background threads because of the way server
diff --git a/server/src/com/vaadin/server/VaadinSession.java b/server/src/com/vaadin/server/VaadinSession.java
index 92a3399bb5..322bfff924 100644
--- a/server/src/com/vaadin/server/VaadinSession.java
+++ b/server/src/com/vaadin/server/VaadinSession.java
@@ -927,4 +927,37 @@ public class VaadinSession implements HttpSessionBindingListener, Serializable {
     private static final Logger getLogger() {
         return Logger.getLogger(VaadinSession.class.getName());
     }
+
+    /**
+     * Performs a safe update of this VaadinSession.
+     * <p>
+     * This method runs the runnable code so that it is safe to update session
+     * variables. It also ensures that all thread locals are set correctly when
+     * executing the runnable.
+     * </p>
+     * <p>
+     * Note that using this method for a VaadinSession which has been detached
+     * from its underlying HTTP session is not necessarily safe. Exclusive
+     * access is provided through locking which is done using the underlying
+     * session.
+     * </p>
+     * 
+     * @param runnable
+     *            The runnable which updates the session
+     */
+    public void runSafely(Runnable runnable) {
+        Map<Class<?>, CurrentInstance> old = null;
+        lock();
+        try {
+            old = CurrentInstance.setThreadLocals(this);
+            runnable.run();
+        } finally {
+            unlock();
+            if (old != null) {
+                CurrentInstance.restoreThreadLocals(old);
+            }
+        }
+
+    }
+
 }
diff --git a/server/src/com/vaadin/ui/UI.java b/server/src/com/vaadin/ui/UI.java
index 796d1f08ea..6a21ba3357 100644
--- a/server/src/com/vaadin/ui/UI.java
+++ b/server/src/com/vaadin/ui/UI.java
@@ -1054,4 +1054,44 @@ public abstract class UI extends AbstractSingleComponentContainer implements
     public int getTabIndex() {
         return getState(false).tabIndex;
     }
+
+    /**
+     * Performs a safe update of this UI.
+     * <p>
+     * This method runs the runnable code so that it is safe to update UI and
+     * session variables. It also ensures that all thread locals are set
+     * correctly when executing the runnable.
+     * </p>
+     * <p>
+     * Note that exclusive access is only guaranteed as long as the UI is
+     * attached to a VaadinSession. If the UI is not attached to a session, this
+     * method makes no guarantees. If the UI is detached then the current
+     * session will also be null.
+     * </p>
+     * 
+     * @param runnable
+     *            The runnable which updates the UI
+     */
+    public void runSafely(Runnable runnable) {
+        Map<Class<?>, CurrentInstance> old = null;
+
+        VaadinSession session = getSession();
+
+        if (session != null) {
+            session.lock();
+        }
+        try {
+            old = CurrentInstance.setThreadLocals(this);
+            runnable.run();
+        } finally {
+            if (session != null) {
+                session.unlock();
+            }
+            if (old != null) {
+                CurrentInstance.restoreThreadLocals(old);
+            }
+        }
+
+    }
+
 }
diff --git a/server/src/com/vaadin/util/CurrentInstance.java b/server/src/com/vaadin/util/CurrentInstance.java
index adf6d963c3..a2cfd4fff7 100644
--- a/server/src/com/vaadin/util/CurrentInstance.java
+++ b/server/src/com/vaadin/util/CurrentInstance.java
@@ -21,8 +21,28 @@ import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
 
+import com.vaadin.server.VaadinPortlet;
+import com.vaadin.server.VaadinPortletService;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.server.VaadinResponse;
+import com.vaadin.server.VaadinService;
+import com.vaadin.server.VaadinServlet;
+import com.vaadin.server.VaadinServletService;
+import com.vaadin.server.VaadinSession;
+import com.vaadin.ui.UI;
+
 /**
  * Keeps track of various thread local instances used by the framework.
+ * <p>
+ * Currently the framework uses the following instances:
+ * </p>
+ * <p>
+ * Inheritable: {@link UI}, {@link VaadinPortlet}, {@link VaadinService},
+ * {@link VaadinServlet}, {@link VaadinSession}.
+ * </p>
+ * <p>
+ * Non-inheritable: {@link VaadinRequest}, {@link VaadinResponse}.
+ * </p>
  * 
  * @author Vaadin Ltd
  * @version @VERSION@
@@ -32,6 +52,18 @@ public class CurrentInstance implements Serializable {
     private final Object instance;
     private final boolean inheritable;
 
+    private static boolean portletAvailable = false;
+    {
+        try {
+            /*
+             * VaadinPortlet depends on portlet API which is available only if
+             * running in a portal.
+             */
+            portletAvailable = (VaadinPortlet.class.getName() != null);
+        } catch (Throwable t) {
+        }
+    }
+
     private static InheritableThreadLocal<Map<Class<?>, CurrentInstance>> instances = new InheritableThreadLocal<Map<Class<?>, CurrentInstance>>() {
         @Override
         protected Map<Class<?>, CurrentInstance> childValue(
@@ -118,7 +150,11 @@ public class CurrentInstance implements Serializable {
                     new CurrentInstance(instance, inheritable));
             if (previousInstance != null) {
                 assert previousInstance.inheritable == inheritable : "Inheritable status mismatch for "
-                        + type;
+                        + type
+                        + " (previous was "
+                        + previousInstance.inheritable
+                        + ", new is "
+                        + inheritable + ")";
             }
         }
     }
@@ -129,4 +165,72 @@ public class CurrentInstance implements Serializable {
     public static void clearAll() {
         instances.get().clear();
     }
+
+    /**
+     * Restores the given thread locals to the given values. Note that this
+     * should only be used internally to restore Vaadin classes.
+     * 
+     * @param old
+     *            A Class -> Object map to set as thread locals
+     */
+    public static void restoreThreadLocals(Map<Class<?>, CurrentInstance> old) {
+        for (Class c : old.keySet()) {
+            CurrentInstance ci = old.get(c);
+            set(c, ci.instance, ci.inheritable);
+        }
+    }
+
+    /**
+     * Sets thread locals for the UI and all related classes
+     * 
+     * @param ui
+     *            The UI
+     * @return A map containing the old values of the thread locals this method
+     *         updated.
+     */
+    public static Map<Class<?>, CurrentInstance> setThreadLocals(UI ui) {
+        Map<Class<?>, CurrentInstance> old = new HashMap<Class<?>, CurrentInstance>();
+        old.put(UI.class, new CurrentInstance(UI.getCurrent(), true));
+        UI.setCurrent(ui);
+        old.putAll(setThreadLocals(ui.getSession()));
+        return old;
+    }
+
+    /**
+     * Sets thread locals for the {@link VaadinSession} and all related classes
+     * 
+     * @param session
+     *            The VaadinSession
+     * @return A map containing the old values of the thread locals this method
+     *         updated.
+     */
+    public static Map<Class<?>, CurrentInstance> setThreadLocals(
+            VaadinSession session) {
+        Map<Class<?>, CurrentInstance> old = new HashMap<Class<?>, CurrentInstance>();
+        old.put(VaadinSession.class,
+                new CurrentInstance(VaadinSession.getCurrent(), true));
+        old.put(VaadinService.class,
+                new CurrentInstance(VaadinService.getCurrent(), true));
+        VaadinService service = null;
+        if (session != null) {
+            service = session.getService();
+        }
+
+        VaadinSession.setCurrent(session);
+        VaadinService.setCurrent(service);
+
+        if (service instanceof VaadinServletService) {
+            old.put(VaadinServlet.class,
+                    new CurrentInstance(VaadinServlet.getCurrent(), true));
+            VaadinServlet.setCurrent(((VaadinServletService) service)
+                    .getServlet());
+        } else if (portletAvailable && service instanceof VaadinPortletService) {
+            old.put(VaadinPortlet.class,
+                    new CurrentInstance(VaadinPortlet.getCurrent(), true));
+            VaadinPortlet.setCurrent(((VaadinPortletService) service)
+                    .getPortlet());
+        }
+
+        return old;
+    }
 }