Unify ThreadLocal handling (#9469)

7 files changed, 156 insertions, 41 deletions

diff --git a/server/src/com/vaadin/Application.java b/server/src/com/vaadin/Application.java
index d2313c0566..088934c5f9 100644
--- a/server/src/com/vaadin/Application.java
+++ b/server/src/com/vaadin/Application.java
@@ -70,6 +70,7 @@ import com.vaadin.ui.AbstractField;
 import com.vaadin.ui.Table;
 import com.vaadin.ui.UI;
 import com.vaadin.ui.Window;
+import com.vaadin.util.CurrentInstance;
 import com.vaadin.util.ReflectTools;
 
 /**
@@ -1799,13 +1800,6 @@ public class Application implements Terminal.ErrorListener, Serializable {
     }
 
     /**
-     * Thread local for keeping track of currently used application instance
-     * 
-     * @since 7.0
-     */
-    private static final ThreadLocal<Application> currentApplication = new ThreadLocal<Application>();
-
-    /**
      * Gets the currently used application. The current application is
      * automatically defined when processing requests to the server. In other
      * cases, (e.g. from background threads), the current application is not
@@ -1819,7 +1813,7 @@ public class Application implements Terminal.ErrorListener, Serializable {
      * @since 7.0
      */
     public static Application getCurrent() {
-        return currentApplication.get();
+        return CurrentInstance.get(Application.class);
     }
 
     /**
@@ -1840,7 +1834,7 @@ public class Application implements Terminal.ErrorListener, Serializable {
      * @since 7.0
      */
     public static void setCurrent(Application application) {
-        currentApplication.set(application);
+        CurrentInstance.setInheritable(Application.class, application);
     }
 
     /**
diff --git a/server/src/com/vaadin/server/PortletApplicationContext2.java b/server/src/com/vaadin/server/PortletApplicationContext2.java
index 3efcc91c08..6a12eafc47 100644
--- a/server/src/com/vaadin/server/PortletApplicationContext2.java
+++ b/server/src/com/vaadin/server/PortletApplicationContext2.java
@@ -46,6 +46,7 @@ import javax.xml.namespace.QName;
 
 import com.vaadin.Application;
 import com.vaadin.ui.UI;
+import com.vaadin.util.CurrentInstance;
 
 /**
  * TODO Write documentation, fix JavaDoc tags.
@@ -153,8 +154,9 @@ public class PortletApplicationContext2 extends ApplicationContext {
     }
 
     private PortletResponse getCurrentResponse() {
-        WrappedPortletResponse currentResponse = VaadinPortlet
-                .getCurrentResponse();
+        WrappedPortletResponse currentResponse = (WrappedPortletResponse) CurrentInstance
+                .get(WrappedResponse.class);
+
         if (currentResponse != null) {
             return currentResponse.getPortletResponse();
         } else {
@@ -163,8 +165,10 @@ public class PortletApplicationContext2 extends ApplicationContext {
     }
 
     public PortletConfig getPortletConfig() {
-        return VaadinPortlet.getCurrentResponse().getDeploymentConfiguration()
-                .getPortlet().getPortletConfig();
+        WrappedPortletResponse response = (WrappedPortletResponse) CurrentInstance
+                .get(WrappedResponse.class);
+        return response.getDeploymentConfiguration().getPortlet()
+                .getPortletConfig();
     }
 
     public void addPortletListener(Application app, PortletListener listener) {
diff --git a/server/src/com/vaadin/server/ServletApplicationContext.java b/server/src/com/vaadin/server/ServletApplicationContext.java
index ecf6202917..4184b68de4 100644
--- a/server/src/com/vaadin/server/ServletApplicationContext.java
+++ b/server/src/com/vaadin/server/ServletApplicationContext.java
@@ -25,6 +25,7 @@ import javax.servlet.http.HttpSessionBindingEvent;
 import javax.servlet.http.HttpSessionBindingListener;
 
 import com.vaadin.Application;
+import com.vaadin.util.CurrentInstance;
 
 /**
  * Web application context for Vaadin applications.
@@ -84,7 +85,8 @@ public class ServletApplicationContext extends ApplicationContext {
         reinitializingSession = false;
 
         // Create a new session
-        HttpSession newSession = VaadinServlet.getCurrentRequest().getSession();
+        HttpSession newSession = WrappedHttpServletRequest.cast(
+                CurrentInstance.get(WrappedRequest.class)).getSession();
 
         // Restores all attributes (security key, reference to this context
         // instance)
diff --git a/server/src/com/vaadin/server/VaadinPortlet.java b/server/src/com/vaadin/server/VaadinPortlet.java
index 199b8e1fc1..0a99644735 100644
--- a/server/src/com/vaadin/server/VaadinPortlet.java
+++ b/server/src/com/vaadin/server/VaadinPortlet.java
@@ -57,6 +57,7 @@ import com.vaadin.Application.ApplicationStartEvent;
 import com.vaadin.server.AbstractCommunicationManager.Callback;
 import com.vaadin.server.ServletPortletHelper.ApplicationClassException;
 import com.vaadin.ui.UI;
+import com.vaadin.util.CurrentInstance;
 
 /**
  * Portlet 2.0 base class. This replaces the servlet in servlet/portlet 1.0
@@ -296,8 +297,6 @@ public class VaadinPortlet extends GenericPortlet implements Constants {
     private PortletDeploymentConfiguration deploymentConfiguration;
     private AddonContext addonContext;
 
-    private static ThreadLocal<WrappedPortletResponse> currentResponse = new ThreadLocal<WrappedPortletResponse>();
-
     @Override
     public void init(PortletConfig config) throws PortletException {
         super.init(config);
@@ -395,10 +394,6 @@ public class VaadinPortlet extends GenericPortlet implements Constants {
         return deploymentConfiguration.isProductionMode();
     }
 
-    public static WrappedPortletResponse getCurrentResponse() {
-        return currentResponse.get();
-    }
-
     protected void handleRequest(PortletRequest request,
             PortletResponse response) throws PortletException, IOException {
         RequestTimer requestTimer = new RequestTimer();
@@ -412,7 +407,8 @@ public class VaadinPortlet extends GenericPortlet implements Constants {
         WrappedPortletResponse wrappedResponse = new WrappedPortletResponse(
                 response, getDeploymentConfiguration());
 
-        currentResponse.set(wrappedResponse);
+        CurrentInstance.set(WrappedRequest.class, wrappedRequest);
+        CurrentInstance.set(WrappedResponse.class, wrappedResponse);
 
         RequestType requestType = getRequestType(wrappedRequest);
 
@@ -608,9 +604,7 @@ public class VaadinPortlet extends GenericPortlet implements Constants {
 
                         }
                     } finally {
-                        UI.setCurrent(null);
-                        Application.setCurrent(null);
-                        currentResponse.set(null);
+                        CurrentInstance.clearAll();
 
                         PortletSession session = request
                                 .getPortletSession(false);
diff --git a/server/src/com/vaadin/server/VaadinServlet.java b/server/src/com/vaadin/server/VaadinServlet.java
index 078263c623..81569030e7 100644
--- a/server/src/com/vaadin/server/VaadinServlet.java
+++ b/server/src/com/vaadin/server/VaadinServlet.java
@@ -50,6 +50,7 @@ import com.vaadin.server.AbstractCommunicationManager.Callback;
 import com.vaadin.server.ServletPortletHelper.ApplicationClassException;
 import com.vaadin.shared.ApplicationConstants;
 import com.vaadin.ui.UI;
+import com.vaadin.util.CurrentInstance;
 
 @SuppressWarnings("serial")
 public class VaadinServlet extends HttpServlet implements Constants {
@@ -157,8 +158,6 @@ public class VaadinServlet extends HttpServlet implements Constants {
         }
     }
 
-    private static ThreadLocal<WrappedHttpServletRequest> currentRequest = new ThreadLocal<WrappedHttpServletRequest>();
-
     // TODO Move some (all?) of the constants to a separate interface (shared
     // with portlet)
 
@@ -263,17 +262,14 @@ public class VaadinServlet extends HttpServlet implements Constants {
         service(createWrappedRequest(request), createWrappedResponse(response));
     }
 
-    public static WrappedHttpServletRequest getCurrentRequest() {
-        return currentRequest.get();
-    }
-
     private void service(WrappedHttpServletRequest request,
             WrappedHttpServletResponse response) throws ServletException,
             IOException {
         RequestTimer requestTimer = new RequestTimer();
         requestTimer.start();
 
-        currentRequest.set(request);
+        CurrentInstance.set(WrappedResponse.class, response);
+        CurrentInstance.set(WrappedRequest.class, request);
 
         AbstractApplicationServletWrapper servletWrapper = new AbstractApplicationServletWrapper(
                 this);
@@ -423,9 +419,7 @@ public class VaadinServlet extends HttpServlet implements Constants {
                                 .onRequestEnd(request, response);
                     }
                 } finally {
-                    UI.setCurrent(null);
-                    Application.setCurrent(null);
-                    currentRequest.set(null);
+                    CurrentInstance.clearAll();
 
                     HttpSession session = request.getSession(false);
                     if (session != null) {
diff --git a/server/src/com/vaadin/ui/UI.java b/server/src/com/vaadin/ui/UI.java
index 7ae4e6bda3..e016e9432c 100644
--- a/server/src/com/vaadin/ui/UI.java
+++ b/server/src/com/vaadin/ui/UI.java
@@ -49,6 +49,7 @@ import com.vaadin.shared.ui.BorderStyle;
 import com.vaadin.shared.ui.ui.UIConstants;
 import com.vaadin.shared.ui.ui.UIServerRpc;
 import com.vaadin.shared.ui.ui.UIState;
+import com.vaadin.util.CurrentInstance;
 import com.vaadin.util.ReflectTools;
 
 /**
@@ -450,11 +451,6 @@ public abstract class UI extends AbstractComponentContainer implements
      */
     protected ActionManager actionManager;
 
-    /**
-     * Thread local for keeping track of the current UI.
-     */
-    private static final ThreadLocal<UI> currentUI = new ThreadLocal<UI>();
-
     /** Identifies the click event */
     private ConnectorTracker connectorTracker = new ConnectorTracker(this);
 
@@ -982,7 +978,7 @@ public abstract class UI extends AbstractComponentContainer implements
      * @see ThreadLocal
      */
     public static void setCurrent(UI ui) {
-        currentUI.set(ui);
+        CurrentInstance.setInheritable(UI.class, ui);
     }
 
     /**
@@ -995,7 +991,7 @@ public abstract class UI extends AbstractComponentContainer implements
      * @see #setCurrent(UI)
      */
     public static UI getCurrent() {
-        return currentUI.get();
+        return CurrentInstance.get(UI.class);
     }
 
     public void setScrollTop(int scrollTop) {
diff --git a/server/src/com/vaadin/util/CurrentInstance.java b/server/src/com/vaadin/util/CurrentInstance.java
new file mode 100644
index 0000000000..8478ba9486
--- /dev/null
+++ b/server/src/com/vaadin/util/CurrentInstance.java
@@ -0,0 +1,131 @@
+/*
+ * Copyright 2011 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.util;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Map.Entry;
+
+/**
+ * Keeps track of various thread local instances used by the framework.
+ * 
+ * @author Vaadin Ltd
+ * @version @VERSION@
+ * @since 7.0.0
+ */
+public class CurrentInstance {
+    private final Object instance;
+    private final boolean inheritable;
+
+    private static InheritableThreadLocal<Map<Class<?>, CurrentInstance>> instances = new InheritableThreadLocal<Map<Class<?>, CurrentInstance>>() {
+        @Override
+        protected Map<Class<?>, CurrentInstance> childValue(
+                Map<Class<?>, CurrentInstance> parentValue) {
+            Map<Class<?>, CurrentInstance> value = new HashMap<Class<?>, CurrentInstance>();
+
+            // Copy all inheritable values to child map
+            for (Entry<Class<?>, CurrentInstance> e : parentValue.entrySet()) {
+                if (e.getValue().inheritable) {
+                    value.put(e.getKey(), e.getValue());
+                }
+            }
+
+            return value;
+        }
+
+        @Override
+        protected Map<java.lang.Class<?>, CurrentInstance> initialValue() {
+            return new HashMap<Class<?>, CurrentInstance>();
+        }
+    };
+
+    private CurrentInstance(Object instance, boolean inheritable) {
+        this.instance = instance;
+        this.inheritable = inheritable;
+    }
+
+    /**
+     * Gets the current instance of a specific type if available.
+     * 
+     * @param type
+     *            the class to get an instance of
+     * @return the current instance or the provided type, or <code>null</code>
+     *         if there is no current instance.
+     */
+    public static <T> T get(Class<T> type) {
+        CurrentInstance currentInstance = instances.get().get(type);
+        if (currentInstance != null) {
+            return type.cast(currentInstance.instance);
+        } else {
+            return null;
+        }
+    }
+
+    /**
+     * Sets the current instance of the given type.
+     * 
+     * @see #setInheritable(Class, Object)
+     * @see ThreadLocal
+     * 
+     * @param type
+     *            the class that should be used when getting the current
+     *            instance back
+     * @param instance
+     *            the actual instance
+     */
+    public static <T> void set(Class<T> type, T instance) {
+        set(type, instance, false);
+    }
+
+    /**
+     * Sets the current inheritable instance of the given type. A current
+     * instance that is inheritable will be available for child threads.
+     * 
+     * @see #set(Class, Object)
+     * @see InheritableThreadLocal
+     * 
+     * @param type
+     *            the class that should be used when getting the current
+     *            instance back
+     * @param instance
+     *            the actual instance
+     */
+    public static <T> void setInheritable(Class<T> type, T instance) {
+        set(type, instance, true);
+    }
+
+    private static <T> void set(Class<T> type, T instance, boolean inheritable) {
+        if (instance == null) {
+            instances.get().remove(type);
+        } else {
+            assert type.isInstance(instance) : "Invald instance type";
+            CurrentInstance previousInstance = instances.get().put(type,
+                    new CurrentInstance(instance, inheritable));
+            if (previousInstance != null) {
+                assert previousInstance.inheritable == inheritable : "Inheritable status mismatch for "
+                        + type;
+            }
+        }
+    }
+
+    /**
+     * Clears all current instances.
+     */
+    public static void clearAll() {
+        instances.get().clear();
+    }
+}