Protect CurrentInstance instances from garbage collection (#12509)

Change-Id: I9ec103a1a42d8888d6f680f477393807223740cf

2 files changed, 74 insertions, 236 deletions

diff --git a/server/src/com/vaadin/util/CurrentInstance.java b/server/src/com/vaadin/util/CurrentInstance.java
index a1c543117d..3320a04366 100644
--- a/server/src/com/vaadin/util/CurrentInstance.java
+++ b/server/src/com/vaadin/util/CurrentInstance.java
@@ -17,10 +17,14 @@
 package com.vaadin.util;
 
 import java.io.Serializable;
+import java.lang.ref.WeakReference;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 import com.vaadin.server.VaadinRequest;
 import com.vaadin.server.VaadinResponse;
@@ -36,6 +40,10 @@ import com.vaadin.ui.UI;
  * when {@link VaadinSession#access(Runnable)} or {@link UI#access(Runnable)} is
  * used.
  * <p>
+ * Please note that the instances are stored using {@link WeakReference}. This
+ * means that the a current instance value may suddenly disappear if there a no
+ * other references to the object.
+ * <p>
  * Currently the framework uses the following instances:
  * </p>
  * <p>
@@ -49,7 +57,7 @@ import com.vaadin.ui.UI;
  * @since 7.0.0
  */
 public class CurrentInstance implements Serializable {
-    private final Object instance;
+    private final WeakReference<Object> instance;
     private final boolean inheritable;
 
     private static InheritableThreadLocal<Map<Class<?>, CurrentInstance>> instances = new InheritableThreadLocal<Map<Class<?>, CurrentInstance>>() {
@@ -60,7 +68,7 @@ public class CurrentInstance implements Serializable {
                 return null;
             }
 
-            Map<Class<?>, CurrentInstance> value = new WeakValueMap<Class<?>, CurrentInstance>();
+            Map<Class<?>, CurrentInstance> value = new HashMap<Class<?>, CurrentInstance>();
 
             // Copy all inheritable values to child map
             for (Entry<Class<?>, CurrentInstance> e : parentValue.entrySet()) {
@@ -74,7 +82,7 @@ public class CurrentInstance implements Serializable {
     };
 
     private CurrentInstance(Object instance, boolean inheritable) {
-        this.instance = instance;
+        this.instance = new WeakReference<Object>(instance);
         this.inheritable = inheritable;
     }
 
@@ -93,12 +101,49 @@ public class CurrentInstance implements Serializable {
         }
         CurrentInstance currentInstance = map.get(type);
         if (currentInstance != null) {
-            return type.cast(currentInstance.instance);
+            Object value = currentInstance.instance.get();
+            if (value == null) {
+                /*
+                 * This is believed to never actually happen since the
+                 * ThreadLocal should only outlive the referenced object on
+                 * threads that are not doing anything related to Vaadin, which
+                 * should thus never invoke CurrentInstance.get().
+                 * 
+                 * At this point, there might also be other values that have
+                 * been collected, so we'll scan the entire map and remove stale
+                 * CurrentInstance objects. Using a ReferenceQueue could make
+                 * this assumingly rare case slightly more efficient, but would
+                 * significantly increase the complexity of the code for
+                 * maintaining a separate ReferenceQueue for each Thread.
+                 */
+                removeStaleInstances(map);
+
+                if (map.isEmpty()) {
+                    instances.remove();
+                }
+
+                return null;
+            }
+            return type.cast(value);
         } else {
             return null;
         }
     }
 
+    private static void removeStaleInstances(Map<Class<?>, CurrentInstance> map) {
+        for (Iterator<Entry<Class<?>, CurrentInstance>> iterator = map
+                .entrySet().iterator(); iterator.hasNext();) {
+            Entry<Class<?>, CurrentInstance> entry = iterator.next();
+            Object instance = entry.getValue().instance.get();
+            if (instance == null) {
+                iterator.remove();
+                getLogger().log(Level.WARNING,
+                        "CurrentInstance for {0} has been garbage collected.",
+                        entry.getKey());
+            }
+        }
+    }
+
     /**
      * Sets the current instance of the given type.
      * 
@@ -183,9 +228,19 @@ public class CurrentInstance implements Serializable {
      *            A Class -> CurrentInstance map to set as current instances
      */
     public static void restoreInstances(Map<Class<?>, CurrentInstance> old) {
+        boolean removeStale = false;
         for (Class c : old.keySet()) {
             CurrentInstance ci = old.get(c);
-            set(c, ci.instance, ci.inheritable);
+            Object v = ci.instance.get();
+            if (v == null) {
+                removeStale = true;
+            } else {
+                set(c, v, ci.inheritable);
+            }
+        }
+
+        if (removeStale) {
+            removeStaleInstances(old);
         }
     }
 
@@ -207,12 +262,21 @@ public class CurrentInstance implements Serializable {
             return Collections.emptyMap();
         } else {
             Map<Class<?>, CurrentInstance> copy = new HashMap<Class<?>, CurrentInstance>();
+            boolean removeStale = false;
             for (Class<?> c : map.keySet()) {
                 CurrentInstance ci = map.get(c);
-                if (ci.inheritable || !onlyInheritable) {
+                if (ci.instance.isEnqueued()) {
+                    removeStale = true;
+                } else if (ci.inheritable || !onlyInheritable) {
                     copy.put(c, ci);
                 }
             }
+            if (removeStale) {
+                removeStaleInstances(map);
+                if (map.isEmpty()) {
+                    instances.remove();
+                }
+            }
             return copy;
         }
     }
@@ -266,4 +330,8 @@ public class CurrentInstance implements Serializable {
 
         return old;
     }
+
+    private static Logger getLogger() {
+        return Logger.getLogger(CurrentInstance.class.getName());
+    }
 }
diff --git a/server/src/com/vaadin/util/WeakValueMap.java b/server/src/com/vaadin/util/WeakValueMap.java
deleted file mode 100644
index 1134594cba..0000000000
--- a/server/src/com/vaadin/util/WeakValueMap.java
+++ /dev/null
@@ -1,230 +0,0 @@
-/*
- * 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.ref.Reference;
-import java.lang.ref.ReferenceQueue;
-import java.lang.ref.WeakReference;
-import java.util.AbstractMap;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
-
-/**
- * A Map holding weak references to its values. It is internally backed by a
- * normal HashMap and all values are stored as WeakReferences. Garbage collected
- * entries are removed when touched.
- * <p>
- * <em>Note</em> this class is not serializable.
- * 
- * @author Vaadin Ltd
- * @since 7.1.4
- */
-public class WeakValueMap<K, V> implements Map<K, V> {
-
-    /**
-     * This class holds a weak reference to the value and a strong reference to
-     * the key for efficient removal of stale values.
-     */
-    private static class WeakValueReference<K, V> extends WeakReference<V> {
-        private final K key;
-
-        WeakValueReference(K key, V value, ReferenceQueue<V> refQueue) {
-            super(value, refQueue);
-            this.key = key;
-        }
-
-        K getKey() {
-            return key;
-        }
-    }
-
-    private final HashMap<K, WeakValueReference<K, V>> backingMap;
-    private final ReferenceQueue<V> refQueue;
-
-    /**
-     * Constructs a new WeakValueMap, where all values are stored as weak
-     * references.
-     */
-    public WeakValueMap() {
-        backingMap = new HashMap<K, WeakValueReference<K, V>>();
-        refQueue = new ReferenceQueue<V>();
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public V put(K key, V value) {
-        if (key == null) {
-            throw new NullPointerException("key cannot be null");
-        }
-        if (value == null) {
-            throw new NullPointerException("value cannot be null");
-        }
-        removeStaleEntries();
-        backingMap.put(key, new WeakValueReference<K, V>(key, value, refQueue));
-        return value;
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public V remove(Object o) {
-        removeStaleEntries();
-        WeakReference<V> value = backingMap.remove(o);
-        return value == null ? null : value.get();
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public void putAll(Map<? extends K, ? extends V> map) {
-        if (map != null) {
-            for (Map.Entry<? extends K, ? extends V> entry : map.entrySet()) {
-                put(entry.getKey(), entry.getValue());
-            }
-        }
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public void clear() {
-        backingMap.clear();
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public Set<K> keySet() {
-        removeStaleEntries();
-        return backingMap.keySet();
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public V get(Object o) {
-        removeStaleEntries();
-        WeakReference<V> weakValue = backingMap.get(o);
-        if (weakValue != null) {
-            return weakValue.get();
-        }
-        return null;
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public int size() {
-        removeStaleEntries();
-        return backingMap.size();
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public boolean isEmpty() {
-        removeStaleEntries();
-        return backingMap.isEmpty();
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public boolean containsKey(Object o) {
-        removeStaleEntries();
-        return backingMap.containsKey(o);
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public boolean containsValue(Object o) {
-        removeStaleEntries();
-        for (V value : values()) {
-            if (o.equals(value)) {
-                return true;
-            }
-        }
-        return false;
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public Collection<V> values() {
-        removeStaleEntries();
-        Collection<V> values = new HashSet<V>();
-        for (WeakReference<V> weakValue : backingMap.values()) {
-            V value = weakValue.get();
-            if (value != null) {
-                // null values have been GC'd, which may happen long before
-                // anything is enqueued in the ReferenceQueue.
-                values.add(value);
-            }
-        }
-        return values;
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public Set<Entry<K, V>> entrySet() {
-        removeStaleEntries();
-        Set<Entry<K, V>> entrySet = new HashSet<Entry<K, V>>();
-        for (Entry<K, WeakValueReference<K, V>> entry : backingMap.entrySet()) {
-            V value = entry.getValue().get();
-            if (value != null) {
-                // null values have been GC'd, which may happen long before
-                // anything is enqueued in the ReferenceQueue.
-                entrySet.add(new AbstractMap.SimpleEntry<K, V>(entry.getKey(),
-                        value));
-            }
-        }
-        return entrySet;
-    }
-
-    /**
-     * Cleans up stale entries by polling the ReferenceQueue.
-     * <p>
-     * Depending on the GC implementation and strategy, the ReferenceQueue is
-     * not necessarily notified immediately when a reference is garbage
-     * collected, but it will eventually be.
-     */
-    private void removeStaleEntries() {
-        Reference<? extends V> ref;
-        while ((ref = refQueue.poll()) != null) {
-            Object key = ((WeakValueReference<?, ?>) ref).getKey();
-            backingMap.remove(key);
-        }
-    }
-}