The bug in the test that became flaky has been fixed. Note that this is not a strict rollback of the rollback - I've added protection against the NPE that would happen if there's a GC between the isBlacklisted call and the getSources call.

*** Original change description ***

Enhance WeakKeySet to auto evict keys and avoid calling toString on Keys.

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=64507759

5 files changed, 161 insertions, 51 deletions

diff --git a/core/src/com/google/inject/internal/AbstractBindingProcessor.java b/core/src/com/google/inject/internal/AbstractBindingProcessor.java
index 72dae2bc..e1a0bba9 100644
--- a/core/src/com/google/inject/internal/AbstractBindingProcessor.java
+++ b/core/src/com/google/inject/internal/AbstractBindingProcessor.java
@@ -98,7 +98,7 @@ abstract class AbstractBindingProcessor extends AbstractProcessor {
     }
 
     // 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);
   }
 
diff --git a/core/src/com/google/inject/internal/InheritingState.java b/core/src/com/google/inject/internal/InheritingState.java
index 2d694cb4..06e0ec47 100644
--- a/core/src/com/google/inject/internal/InheritingState.java
+++ b/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 @@ final class InheritingState implements State {
     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) {
diff --git a/core/src/com/google/inject/internal/InjectorImpl.java b/core/src/com/google/inject/internal/InjectorImpl.java
index 6faab973..eff9c8fc 100644
--- a/core/src/com/google/inject/internal/InjectorImpl.java
+++ b/core/src/com/google/inject/internal/InjectorImpl.java
@@ -65,7 +65,7 @@ import java.util.Set;
  */
 final class InjectorImpl implements Injector, Lookups {
   public static final TypeLiteral<String> STRING_TYPE = TypeLiteral.get(String.class);
-  
+
   /** Options that control how the injector behaves. */
   static class InjectorOptions {
     final Stage stage;
@@ -82,7 +82,7 @@ final class InjectorImpl implements Injector, Lookups {
       this.atInjectRequired = atInjectRequired;
       this.exactBindingAnnotationsRequired = exactBindingAnnotationsRequired;
     }
-    
+
     @Override
     public String toString() {
       return Objects.toStringHelper(getClass())
@@ -94,9 +94,9 @@ final class InjectorImpl implements Injector, Lookups {
           .toString();
     }
   }
-  
+
   /** some limitations on what just in time bindings are allowed. */
-  enum JitLimitation { 
+  enum JitLimitation {
     /** does not allow just in time bindings */
     NO_JIT,
     /** allows existing just in time bindings, but does not allow new ones */
@@ -175,7 +175,7 @@ final class InjectorImpl implements Injector, Lookups {
         }
       }
     }
-    
+
     // If Key is a Provider, we have to see if the type it is providing exists,
     // and, if so, we have to create the binding for the provider.
     if(isProvider(key)) {
@@ -290,7 +290,7 @@ final class InjectorImpl implements Injector, Lookups {
   private static boolean isTypeLiteral(Key<?> key) {
     return key.getTypeLiteral().getRawType().equals(TypeLiteral.class);
   }
-  
+
   private static <T> Key<T> getProvidedKey(Key<Provider<T>> key, Errors errors) throws ErrorsException {
     Type providerType = key.getTypeLiteral().getType();
 
@@ -384,7 +384,7 @@ final class InjectorImpl implements Injector, Lookups {
     public Set<Dependency<?>> getDependencies() {
       return ImmutableSet.<Dependency<?>>of(Dependency.get(getProvidedKey()));
     }
-    
+
     @Override
     public boolean equals(Object obj) {
       if(obj instanceof ProviderBindingImpl) {
@@ -396,7 +396,7 @@ final class InjectorImpl implements Injector, Lookups {
         return false;
       }
     }
-    
+
     @Override
     public int hashCode() {
       return Objects.hashCode(getKey(), getScoping(), providedBinding);
@@ -508,7 +508,7 @@ final class InjectorImpl implements Injector, Lookups {
           .add("value", value)
           .toString();
     }
-    
+
     @Override
     public boolean equals(Object obj) {
       if(obj instanceof ConvertedConstantBindingImpl) {
@@ -520,7 +520,7 @@ final class InjectorImpl implements Injector, Lookups {
         return false;
       }
     }
-    
+
     @Override
     public int hashCode() {
       return Objects.hashCode(getKey(), getScoping(), value);
@@ -536,7 +536,7 @@ final class InjectorImpl implements Injector, Lookups {
   <T> void initializeJitBinding(BindingImpl<T> binding, Errors errors) throws ErrorsException {
     // Put the partially constructed binding in the map a little early. This enables us to handle
     // circular dependencies. Example: FooImpl -> BarImpl -> FooImpl.
-    // Note: We don't need to synchronize on state.lock() during injector creation.    
+    // Note: We don't need to synchronize on state.lock() during injector creation.
     if (binding instanceof DelayedInitialize) {
       Key<T> key = binding.getKey();
       jitBindings.put(key, binding);
@@ -556,7 +556,7 @@ final class InjectorImpl implements Injector, Lookups {
       }
     }
   }
-  
+
   /**
    * Iterates through the binding's dependencies to clean up any stray bindings that were leftover
    * from a failed JIT binding. This is required because the bindings are eagerly &
@@ -593,7 +593,7 @@ final class InjectorImpl implements Injector, Lookups {
     }
     return bindingFailed;
   }
-  
+
   /** Cleans up any state that may have been cached when constructing the JIT binding. */
   private void removeFailedJitBinding(Binding<?> binding, InjectionPoint ip) {
     failedJitBindings.add(binding.getKey());
@@ -604,7 +604,7 @@ final class InjectorImpl implements Injector, Lookups {
       constructors.remove(ip);
     }
   }
-  
+
   /** Safely gets the dependencies of possibly not initialized bindings. */
   @SuppressWarnings("unchecked")
   private Set<Dependency<?>> getInternalDependencies(BindingImpl<?> binding) {
@@ -652,7 +652,7 @@ final class InjectorImpl implements Injector, Lookups {
       return createProvidedByBinding(key, scoping, providedBy, errors);
     }
 
-    
+
     return ConstructorBindingImpl.create(this,
         key,
         null, /* use default constructor */
@@ -785,13 +785,17 @@ final class InjectorImpl implements Injector, Lookups {
       }
     }
 
+    // Retrieve the sources before checking for blacklisting to guard against sources becoming null
+    // due to a full GC happening after calling state.isBlacklisted and
+    // state.getSourcesForBlacklistedKey.
+    // TODO(user): Consolidate these two APIs.
+    Set<Object> sources = state.getSourcesForBlacklistedKey(key);
     if (state.isBlacklisted(key)) {
-      Set<Object> sources = state.getSourcesForBlacklistedKey(key);
       throw errors.childBindingAlreadySet(key, sources).toException();
     }
 
     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;
   }
@@ -813,8 +817,12 @@ final class InjectorImpl implements Injector, Lookups {
       boolean jitDisabled, JitLimitation jitType) throws ErrorsException {
     int numErrorsBefore = errors.size();
 
+    // Retrieve the sources before checking for blacklisting to guard against sources becoming null
+    // due to a full GC happening after calling state.isBlacklisted and
+    // state.getSourcesForBlacklistedKey.
+    // TODO(user): Consolidate these two APIs.
+    Set<Object> sources = state.getSourcesForBlacklistedKey(key);
     if (state.isBlacklisted(key)) {
-      Set<Object> sources = state.getSourcesForBlacklistedKey(key);
       throw errors.childBindingAlreadySet(key, sources).toException();
     }
 
@@ -842,7 +850,7 @@ final class InjectorImpl implements Injector, Lookups {
       return convertedBinding;
     }
 
-    if (!isTypeLiteral(key) 
+    if (!isTypeLiteral(key)
         && jitDisabled
         && jitType != JitLimitation.NEW_OR_EXISTING_JIT) {
       throw errors.jitDisabled(key).toException();
@@ -912,7 +920,7 @@ final class InjectorImpl implements Injector, Lookups {
     <T> List<Binding<T>> getAll(TypeLiteral<T> type) {
       List<Binding<?>> bindings = multimap.get(type);
       return bindings != null
-          ? Collections.<Binding<T>>unmodifiableList((List) multimap.get(type)) 
+          ? Collections.<Binding<T>>unmodifiableList((List) multimap.get(type))
           : ImmutableList.<Binding<T>>of();
     }
   }
@@ -958,7 +966,7 @@ final class InjectorImpl implements Injector, Lookups {
 
   /** Cached field and method injectors for each type. */
   MembersInjectorStore membersInjectorStore;
-  
+
   /** Cached provision listener callbacks for each key. */
   ProvisionListenerCallbackStore provisionListenerStore;
 
diff --git a/core/src/com/google/inject/internal/State.java b/core/src/com/google/inject/internal/State.java
index d824cc1b..3c9e0817 100644
--- a/core/src/com/google/inject/internal/State.java
+++ b/core/src/com/google/inject/internal/State.java
@@ -105,7 +105,7 @@ interface State {
       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 @@ interface State {
   /**
    * 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
diff --git a/core/src/com/google/inject/internal/WeakKeySet.java b/core/src/com/google/inject/internal/WeakKeySet.java
index 761cd35e..43bbab1f 100644
--- a/core/src/com/google/inject/internal/WeakKeySet.java
+++ b/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 @@ final class WeakKeySet {
     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 == null) ? null : 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);
+    }
   }
 }