2 files changed, 171 insertions, 9 deletions
diff --git a/server/src/com/vaadin/util/CurrentInstance.java b/server/src/com/vaadin/util/CurrentInstance.java
index a2cfd4fff7..60489d596e 100644
--- a/server/src/com/vaadin/util/CurrentInstance.java
+++ b/server/src/com/vaadin/util/CurrentInstance.java
@@ -68,6 +68,10 @@ public class CurrentInstance implements Serializable {
         @Override
         protected Map<Class<?>, CurrentInstance> childValue(
                 Map<Class<?>, CurrentInstance> parentValue) {
+            if (parentValue == null) {
+                return null;
+            }
+
             Map<Class<?>, CurrentInstance> value = new HashMap<Class<?>, CurrentInstance>();
 
             // Copy all inheritable values to child map
@@ -79,11 +83,6 @@ public class CurrentInstance implements Serializable {
 
             return value;
         }
-
-        @Override
-        protected Map<java.lang.Class<?>, CurrentInstance> initialValue() {
-            return new HashMap<Class<?>, CurrentInstance>();
-        }
     };
 
     private CurrentInstance(Object instance, boolean inheritable) {
@@ -100,7 +99,11 @@ public class CurrentInstance implements Serializable {
      *         if there is no current instance.
      */
     public static <T> T get(Class<T> type) {
-        CurrentInstance currentInstance = instances.get().get(type);
+        Map<Class<?>, CurrentInstance> map = instances.get();
+        if (map == null) {
+            return null;
+        }
+        CurrentInstance currentInstance = map.get(type);
         if (currentInstance != null) {
             return type.cast(currentInstance.instance);
         } else {
@@ -142,11 +145,25 @@ public class CurrentInstance implements Serializable {
     }
 
     private static <T> void set(Class<T> type, T instance, boolean inheritable) {
+        Map<Class<?>, CurrentInstance> map = instances.get();
         if (instance == null) {
-            instances.get().remove(type);
+            // remove the instance
+            if (map == null) {
+                return;
+            }
+            map.remove(type);
+            if (map.isEmpty()) {
+                instances.remove();
+                map = null;
+            }
         } else {
             assert type.isInstance(instance) : "Invald instance type";
-            CurrentInstance previousInstance = instances.get().put(type,
+            if (map == null) {
+                map = new HashMap<Class<?>, CurrentInstance>();
+                instances.set(map);
+            }
+
+            CurrentInstance previousInstance = map.put(type,
                     new CurrentInstance(instance, inheritable));
             if (previousInstance != null) {
                 assert previousInstance.inheritable == inheritable : "Inheritable status mismatch for "
@@ -163,7 +180,7 @@ public class CurrentInstance implements Serializable {
      * Clears all current instances.
      */
     public static void clearAll() {
-        instances.get().clear();
+        instances.remove();
     }
 
     /**
diff --git a/server/tests/src/com/vaadin/util/TestCurrentInstance.java b/server/tests/src/com/vaadin/util/TestCurrentInstance.java
new file mode 100644
index 0000000000..da986abe31
--- /dev/null
+++ b/server/tests/src/com/vaadin/util/TestCurrentInstance.java
@@ -0,0 +1,145 @@
+/*
+ * Copyright 2000-2013 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.lang.reflect.Field;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import junit.framework.Assert;
+
+import org.junit.Test;
+
+public class TestCurrentInstance {
+
+    @Test
+    public void testInitiallyCleared() throws Exception {
+        assertCleared();
+    }
+
+    @Test
+    public void testClearedAfterRemove() throws Exception {
+        CurrentInstance.set(TestCurrentInstance.class, this);
+        Assert.assertEquals(this,
+                CurrentInstance.get(TestCurrentInstance.class));
+        CurrentInstance.set(TestCurrentInstance.class, null);
+
+        assertCleared();
+    }
+
+    @Test
+    public void testClearedAfterRemoveInheritable() throws Exception {
+        CurrentInstance.setInheritable(TestCurrentInstance.class, this);
+        Assert.assertEquals(this,
+                CurrentInstance.get(TestCurrentInstance.class));
+        CurrentInstance.setInheritable(TestCurrentInstance.class, null);
+
+        assertCleared();
+    }
+
+    @Test
+    public void testInheritableThreadLocal() throws Exception {
+        final AtomicBoolean threadFailed = new AtomicBoolean(true);
+
+        CurrentInstance.setInheritable(TestCurrentInstance.class, this);
+        Assert.assertEquals(this,
+                CurrentInstance.get(TestCurrentInstance.class));
+        Thread t = new Thread() {
+            @Override
+            public void run() {
+                Assert.assertEquals(TestCurrentInstance.this,
+                        CurrentInstance.get(TestCurrentInstance.class));
+                threadFailed.set(false);
+            };
+        };
+        t.start();
+        CurrentInstance.set(TestCurrentInstance.class, null);
+
+        assertCleared();
+        while (t.isAlive()) {
+            Thread.sleep(1000);
+        }
+        Assert.assertFalse("Thread failed", threadFailed.get());
+
+    }
+
+    @Test
+    public void testClearedAfterRemoveInSeparateThread() throws Exception {
+        final AtomicBoolean threadFailed = new AtomicBoolean(true);
+
+        CurrentInstance.setInheritable(TestCurrentInstance.class, this);
+        Assert.assertEquals(this,
+                CurrentInstance.get(TestCurrentInstance.class));
+        Thread t = new Thread() {
+            @Override
+            public void run() {
+                try {
+                    Assert.assertEquals(TestCurrentInstance.this,
+                            CurrentInstance.get(TestCurrentInstance.class));
+                    CurrentInstance.set(TestCurrentInstance.class, null);
+                    assertCleared();
+
+                    threadFailed.set(false);
+                } catch (Exception e) {
+                    e.printStackTrace();
+                }
+            };
+        };
+        t.start();
+
+        while (t.isAlive()) {
+            Thread.sleep(1000);
+        }
+        Assert.assertFalse("Thread failed", threadFailed.get());
+
+        // Clearing the threadlocal in the thread should not have cleared it
+        // here
+        Assert.assertEquals(this,
+                CurrentInstance.get(TestCurrentInstance.class));
+
+        // Clearing the only remaining threadlocal should free all memory
+        CurrentInstance.set(TestCurrentInstance.class, null);
+        assertCleared();
+    }
+
+    @Test
+    public void testClearedWithClearAll() throws Exception {
+        CurrentInstance.set(TestCurrentInstance.class, this);
+        Assert.assertEquals(this,
+                CurrentInstance.get(TestCurrentInstance.class));
+        CurrentInstance.clearAll();
+
+        assertCleared();
+    }
+
+    private void assertCleared() throws SecurityException,
+            NoSuchFieldException, IllegalAccessException {
+        Assert.assertNull(getInternalCurrentInstanceVariable().get());
+    }
+
+    private InheritableThreadLocal<Map<Class<?>, CurrentInstance>> getInternalCurrentInstanceVariable()
+            throws SecurityException, NoSuchFieldException,
+            IllegalAccessException {
+        Field f = CurrentInstance.class.getDeclaredField("instances");
+        f.setAccessible(true);
+        return (InheritableThreadLocal<Map<Class<?>, CurrentInstance>>) f
+                .get(null);
+    }
+
+    public void testInheritedClearedAfterRemove() {
+
+    }
+}