Enhance WeakKeySet to auto evict keys and avoid calling toString on K…

…eys.

This should fix https://code.google.com/p/google-guice/issues/detail?id=756.
-------------
Created by MOE: http://code.google.com/p/moe-java
MOE_MIGRATED_REVID=64083354

build.xml

@@ -105,6 +105,7 @@
         <pathelement location="lib/javax.inject.jar"/>
         <pathelement location="lib/aopalliance.jar"/>
         <pathelement location="lib/guava-16.0.1.jar"/>
+        <pathelement location="lib/build/guava-testlib-16.0.1.jar"/>
         <pathelement location="lib/build/junit.jar"/>
         <pathelement location="lib/build/servlet-api-2.5.jar"/>
         <pathelement location="lib/build/easymock.jar"/>

core/pom.xml

@@ -29,6 +29,12 @@
       <artifactId>guava</artifactId>
       <version>16.0.1</version>
     </dependency>
+    <dependency>
+      <groupId>com.google.guava</groupId>
+      <artifactId>guava-testlib</artifactId>
+      <version>16.0.1</version>
+      <scope>test</scope>
+    </dependency>
     <!--
      | CGLIB is embedded by default by the JarJar build profile
     -->

core/src/com/google/inject/internal/AbstractBindingProcessor.java

@@ -98,7 +98,7 @@ protected void putBinding(BindingImpl<?> binding) {
     }
 
     // prevent the parent from creating a JIT binding for this key
-    injector.state.parent().blacklist(key, binding.getSource());
+    injector.state.parent().blacklist(key, injector.state, binding.getSource());
     injector.state.putBinding(key, binding);
   }
 

core/src/com/google/inject/internal/InheritingState.java

@@ -56,12 +56,13 @@ final class InheritingState implements State {
   /*end[AOP]*/
   private final List<TypeListenerBinding> typeListenerBindings = Lists.newArrayList();
   private final List<ProvisionListenerBinding> provisionListenerBindings = Lists.newArrayList(); 
-  private final WeakKeySet blacklistedKeys = new WeakKeySet();
+  private final WeakKeySet blacklistedKeys;
   private final Object lock;
 
   InheritingState(State parent) {
     this.parent = checkNotNull(parent, "parent");
     this.lock = (parent == State.NONE) ? this : parent.lock();
+    this.blacklistedKeys = new WeakKeySet(lock);
   }
 
   public State parent() {
@@ -154,9 +155,9 @@ public List<ProvisionListenerBinding> getProvisionListenerBindings() {
     return result;
   }
 
-  public void blacklist(Key<?> key, Object source) {
-    parent.blacklist(key, source);
-    blacklistedKeys.add(key, source);
+  public void blacklist(Key<?> key, State state, Object source) {
+    parent.blacklist(key, state, source);
+    blacklistedKeys.add(key, state, source);
   }
 
   public boolean isBlacklisted(Key<?> key) {

core/src/com/google/inject/internal/InjectorImpl.java

@@ -791,7 +791,7 @@ private <T> BindingImpl<T> createJustInTimeBindingRecursive(Key<T> key, Errors e
     }
 
     BindingImpl<T> binding = createJustInTimeBinding(key, errors, jitDisabled, jitType);
-    state.parent().blacklist(key, binding.getSource());
+    state.parent().blacklist(key, state, binding.getSource());
     jitBindings.put(key, binding);
     return binding;
   }

core/src/com/google/inject/internal/State.java

@@ -105,7 +105,7 @@ public List<ProvisionListenerBinding> getProvisionListenerBindings() {
       return ImmutableList.of();
     }
 
-    public void blacklist(Key<?> key, Object source) {
+    public void blacklist(Key<?> key, State state, Object source) {
     }
 
     public boolean isBlacklisted(Key<?> key) {
@@ -165,9 +165,9 @@ TypeConverterBinding getConverter(
   /**
    * Forbids the corresponding injector from creating a binding to {@code key}. Child injectors
    * blacklist their bound keys on their parent injectors to prevent just-in-time bindings on the
-   * parent injector that would conflict.
+   * parent injector that would conflict and pass along their state to control the lifetimes.
    */
-  void blacklist(Key<?> key, Object source);
+  void blacklist(Key<?> key, State state, Object source);
 
   /**
    * Returns true if {@code key} is forbidden from being bound in this injector. This indicates that

core/src/com/google/inject/internal/WeakKeySet.java

@@ -16,37 +16,86 @@
 
 package com.google.inject.internal;
 
+import com.google.common.base.Objects;
+import com.google.common.base.Preconditions;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.RemovalCause;
+import com.google.common.cache.RemovalListener;
+import com.google.common.cache.RemovalNotification;
+import com.google.common.collect.LinkedHashMultiset;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Multiset;
 import com.google.common.collect.Sets;
 import com.google.inject.Key;
 import com.google.inject.internal.util.SourceProvider;
 
+import java.lang.annotation.Annotation;
 import java.util.Map;
 import java.util.Set;
 
 /**
  * Minimal set that doesn't hold strong references to the contained keys.
  *
- * @author jessewilson@google.com (Jesse Wilson)
+ * @author dweis@google.com (Daniel Weis)
  */
 final class WeakKeySet {
 
-  // TODO(user): Instead of relying on Key#toString(), consider maintaining a set of custom weak
-  // references for child injectors. On some actions, we can poll the ReferenceQueue and clear the
-  // relevant blacklisted values. This would allow us to avoid forcing the toString memoization on
-  // Key and fix the memory leak referenced in
-  // https://code.google.com/p/google-guice/issues/detail?id=756.
+  /**
+   * The key for the Map is {@link Key} for most bindings, String for multibindings.
+   * <p>
+   * Reason being that multibinding Key's annotations hold a reference to their injector, implying
+   * we'd leak memory.
+   */
+  private Map<Object, Multiset<Object>> backingSet;
+
+  /**
+   * This is already locked externally on add and getSources but we need it to handle clean up in
+   * the evictionCache's RemovalListener.
+   */
+  private final Object lock;
 
   /**
-   * We store strings rather than keys so we don't hold strong references.
-   *
-   * <p>One potential problem with this approach is that parent and child injectors cannot define
-   * keys whose class names are equal but class loaders are different. This shouldn't be an issue
-   * in practice.
+   * Tracks child injector lifetimes and evicts blacklisted keys/sources after the child injector is
+   * garbage collected.
    */
-  private Map<String, Set<Object>> backingSet;
+  private final Cache<State, Set<KeyAndSource>> evictionCache = CacheBuilder.newBuilder()
+      .weakKeys()
+      .removalListener(
+          new RemovalListener<State, Set<KeyAndSource>>() {
+            @Override
+            public void onRemoval(RemovalNotification<State, Set<KeyAndSource>> notification) {
+              Preconditions.checkState(RemovalCause.COLLECTED.equals(notification.getCause()));
 
-  public void add(Key<?> key, Object source) {
+              cleanUpForCollectedState(notification.getValue());
+            }
+          })
+      .build();
+
+  /**
+   * There may be multiple child injectors blacklisting a certain key so only remove the source
+   * that's relevant.
+   */
+  private void cleanUpForCollectedState(Set<KeyAndSource> keysAndSources) {
+    synchronized (lock) {
+
+      for (KeyAndSource keyAndSource : keysAndSources) {
+        Multiset<Object> set = backingSet.get(keyAndSource.mapKey);
+        if (set != null) {
+          set.remove(keyAndSource.source);
+          if (set.isEmpty()) {
+            backingSet.remove(keyAndSource.mapKey);
+          }
+        }
+      }
+    }
+  }
+
+  WeakKeySet(Object lock) {
+    this.lock = lock;
+  }
+
+  public void add(Key<?> key, State state, Object source) {
     if (backingSet == null) {
       backingSet = Maps.newHashMap();
     }
@@ -55,22 +104,74 @@ public void add(Key<?> key, Object source) {
     if (source instanceof Class || source == SourceProvider.UNKNOWN_SOURCE) {
       source = null;
     }
-    String k = key.toString();
-    Set<Object> sources = backingSet.get(k);
+    Object mapKey = toMapKey(key);
+    Multiset<Object> sources = backingSet.get(mapKey);
     if (sources == null) {
-      sources = Sets.newLinkedHashSet();
-      backingSet.put(k, sources);
+      sources = LinkedHashMultiset.create();
+      backingSet.put(mapKey, sources);
+    }
+    Object convertedSource = Errors.convert(source);
+    sources.add(convertedSource);
+
+    // Avoid all the extra work if we can.
+    if (state.parent() != State.NONE) {
+      Set<KeyAndSource> keyAndSources = evictionCache.getIfPresent(state);
+      if (keyAndSources == null) {
+        evictionCache.put(state, keyAndSources = Sets.newHashSet());
+      }
+      keyAndSources.add(new KeyAndSource(mapKey, convertedSource));
     }
-    sources.add(Errors.convert(source));
   }
 
   public boolean contains(Key<?> key) {
-    // avoid calling key.toString() if the backing set is empty. toString is expensive in aggregate,
-    // and most WeakKeySets are empty in practice (because they're used by top-level injectors)
-    return backingSet != null && backingSet.containsKey(key.toString());
+    evictionCache.cleanUp();
+    return backingSet != null && backingSet.containsKey(key);
   }
 
   public Set<Object> getSources(Key<?> key) {
-    return backingSet.get(key.toString());
+    evictionCache.cleanUp();
+    Multiset<Object> sources = backingSet.get(key);
+    return (sources == null) ? null : sources.elementSet();
+  }
+
+  private static Object toMapKey(Key<?> key) {
+    Annotation annotation = key.getAnnotation();
+    if (annotation != null
+        // HACK: See comment on backingSet for more info. This is tested in MultibinderTest,
+        // MapBinderTest, and OptionalBinderTest in the multibinder test suite.
+        && "com.google.inject.multibindings.RealElement".equals(annotation.getClass().getName())) {
+      return key.toString();
+    }
+    return key;
+  }
+
+  private static final class KeyAndSource {
+    final Object mapKey;
+    final Object source;
+
+    KeyAndSource(Object mapKey, Object source) {
+      this.mapKey = mapKey;
+      this.source = source;
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hashCode(mapKey, source);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (this == obj) {
+        return true;
+      }
+
+      if (!(obj instanceof KeyAndSource)) {
+        return false;
+      }
+
+      KeyAndSource other = (KeyAndSource) obj;
+      return Objects.equal(mapKey, other.mapKey)
+          && Objects.equal(source, other.source);
+    }
   }
 }

core/test/com/google/inject/AllTests.java

@@ -20,6 +20,7 @@
 import com.google.inject.internal.MoreTypesTest;
 import com.google.inject.internal.RehashableKeysTest;
 import com.google.inject.internal.UniqueAnnotationsTest;
+import com.google.inject.internal.WeakKeySetTest;
 import com.google.inject.internal.util.LineNumbersTest;
 import com.google.inject.matcher.MatcherTest;
 import com.google.inject.name.NamedEquivalanceTest;
@@ -106,6 +107,7 @@ public static Test suite() {
     suite.addTestSuite(TypeLiteralInjectionTest.class);
     suite.addTestSuite(TypeLiteralTest.class);
     suite.addTestSuite(TypeLiteralTypeResolutionTest.class);
+    suite.addTestSuite(WeakKeySetTest.class);
 
     // internal
     suite.addTestSuite(LineNumbersTest.class);