Protect CurrentInstance instances from garbage collection (#12509)

Change-Id: I9ec103a1a42d8888d6f680f477393807223740cf

4 files changed, 219 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);
-        }
-    }
-}
diff --git a/uitest/src/com/vaadin/tests/components/ui/CurrentUiRetained.html b/uitest/src/com/vaadin/tests/components/ui/CurrentUiRetained.html
new file mode 100644
index 0000000000..fe030f2ea7
--- /dev/null
+++ b/uitest/src/com/vaadin/tests/components/ui/CurrentUiRetained.html
@@ -0,0 +1,52 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-strict.dtd">
+<html xmlns="http://www.w3.org/1999/xhtml" xml:lang="en" lang="en">
+<head profile="http://selenium-ide.openqa.org/profiles/test-case">
+<meta http-equiv="Content-Type" content="text/html; charset=UTF-8" />
+<link rel="selenium.base" href="" />
+<title>New Test</title>
+</head>
+<body>
+<table cellpadding="1" cellspacing="1" border="1">
+<thead>
+<tr><td rowspan="1" colspan="3">New Test</td></tr>
+</thead><tbody>
+<tr>
+	<td>open</td>
+	<td>/run/com.vaadin.tests.components.ui.CurrentUiRetained?restartApplication</td>
+	<td></td>
+</tr>
+<tr>
+	<td>pause</td>
+	<td>2000</td>
+	<td></td>
+</tr>
+<tr>
+	<td>click</td>
+	<td>vaadin=runcomvaadintestscomponentsuiCurrentUiRetained::/VVerticalLayout[0]/Slot[2]/VVerticalLayout[0]/Slot[0]/VButton[0]/domChild[0]/domChild[0]</td>
+	<td></td>
+</tr>
+<tr>
+	<td>assertText</td>
+	<td>vaadin=runcomvaadintestscomponentsuiCurrentUiRetained::PID_SLog_row_3</td>
+	<td>1. Correct UI.getCurrent before GC: true</td>
+</tr>
+<tr>
+	<td>assertText</td>
+	<td>vaadin=runcomvaadintestscomponentsuiCurrentUiRetained::PID_SLog_row_2</td>
+	<td>2. Correct UI.getCurrent after GC: true</td>
+</tr>
+<tr>
+	<td>assertText</td>
+	<td>vaadin=runcomvaadintestscomponentsuiCurrentUiRetained::PID_SLog_row_1</td>
+	<td>3. GC probe available before GC: true</td>
+</tr>
+<tr>
+	<td>assertText</td>
+	<td>vaadin=runcomvaadintestscomponentsuiCurrentUiRetained::PID_SLog_row_0</td>
+	<td>4. GC probe available after GC: false</td>
+</tr>
+
+</tbody></table>
+</body>
+</html>
diff --git a/uitest/src/com/vaadin/tests/components/ui/CurrentUiRetained.java b/uitest/src/com/vaadin/tests/components/ui/CurrentUiRetained.java
new file mode 100644
index 0000000000..b0127e3a75
--- /dev/null
+++ b/uitest/src/com/vaadin/tests/components/ui/CurrentUiRetained.java
@@ -0,0 +1,93 @@
+/*
+ * 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.tests.components.ui;
+
+import java.util.ArrayList;
+
+import com.vaadin.server.VaadinRequest;
+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 CurrentUiRetained extends AbstractTestUIWithLog {
+    public static class GcProbe {
+
+    }
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        final ArrayList<UI> uiLog = new ArrayList<UI>();
+        final ArrayList<Boolean> probeLog = new ArrayList<Boolean>();
+
+        final Thread thread = new Thread(new Runnable() {
+            @Override
+            public void run() {
+                try {
+                    uiLog.add(UI.getCurrent());
+
+                    GcProbe gcProbe = new GcProbe();
+                    CurrentInstance.set(GcProbe.class, gcProbe);
+                    probeLog.add(CurrentInstance.get(GcProbe.class) != null);
+                    gcProbe = null;
+
+                    Thread.sleep(500l);
+                    System.gc();
+                    Thread.sleep(500l);
+
+                    probeLog.add(CurrentInstance.get(GcProbe.class) != null);
+                    uiLog.add(UI.getCurrent());
+                } catch (InterruptedException e) {
+                    e.printStackTrace();
+                }
+            }
+        });
+        thread.start();
+
+        addComponent(new Button("Show result", new Button.ClickListener() {
+            @Override
+            public void buttonClick(ClickEvent event) {
+                try {
+                    thread.join();
+
+                    log("Correct UI.getCurrent before GC: "
+                            + (uiLog.get(0) == CurrentUiRetained.this));
+                    log("Correct UI.getCurrent after GC: "
+                            + (uiLog.get(1) == CurrentUiRetained.this));
+
+                    log("GC probe available before GC: " + probeLog.get(0));
+                    log("GC probe available after GC: " + probeLog.get(1));
+
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        }));
+    }
+
+    @Override
+    protected String getTestDescription() {
+        return "Tests that garbage collection removes stale CurrentInstance values while retaining values not collected.";
+    }
+
+    @Override
+    protected Integer getTicketNumber() {
+        return Integer.valueOf(12509);
+    }
+
+}