Fixes mockito#1614: Clean up mocks when session finishes.

Due to the introduction of map from weak reference from mock instance to
its invocation handler, Mockito became vunerable to memory leaks as
there are multiple situations where Mockito could unintentionally hold
strong references to mock instances in the map record. The strong
references could be through spiedInstance for spies, and arguments used
to faciliate method stubbing.

Mockito could never know if the arguments passed in for method stubbing
are also strongly referenced somewhere else or not, so Mockito needs to
save a strong reference to these arguments to avoid premature GC.
Therefore to solve cyclic strong references through arguments Mockito
needs to explicitly know when mocks are out of their lifespan, and clean
up all internal strong references associated with them.

This commit also fixes mockito#1532 and fixes mockito#1533.

src/main/java/org/mockito/Mockito.java

@@ -2978,7 +2978,7 @@ static MockitoDebugger debug() {
     @Incubating
     @CheckReturnValue
     public static MockitoFramework framework() {
-        return new DefaultMockitoFramework();
+        return MOCKITO_CORE.framework();
     }
 
     /**

src/main/java/org/mockito/exceptions/misusing/MultipleTrackingMockSessionException.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright (c) 2019 Mockito contributors
+ * This program is made available under the terms of the MIT License.
+ */
+
+package org.mockito.exceptions.misusing;
+
+import org.mockito.MockitoSession;
+import org.mockito.exceptions.base.MockitoException;
+
+/**
+ * Reports the misuse where user tries to open more than one {@link MockitoSession} that tracks and cleans up mocks.
+ * User needs to finish previous session that tracks and cleans up mocks before trying to open a new one. Note this
+ * doesn't prevent user from opening sessions that doesn't track or clean up mocks.
+ *
+ * @since 2.24.4
+ */
+public class MultipleTrackingMockSessionException extends MockitoException {
+    public MultipleTrackingMockSessionException(String message) {
+        super(message);
+    }
+}

src/main/java/org/mockito/internal/MockitoCore.java

@@ -7,8 +7,10 @@
 import org.mockito.InOrder;
 import org.mockito.MockSettings;
 import org.mockito.MockingDetails;
+import org.mockito.MockitoFramework;
 import org.mockito.exceptions.misusing.NotAMockException;
 import org.mockito.internal.creation.MockSettingsImpl;
+import org.mockito.internal.framework.DefaultMockitoFramework;
 import org.mockito.internal.invocation.finder.VerifiableInvocationsFinder;
 import org.mockito.internal.listeners.VerificationStartedNotifier;
 import org.mockito.internal.progress.MockingProgress;
@@ -49,6 +51,8 @@
 @SuppressWarnings("unchecked")
 public class MockitoCore {
 
+    private DefaultMockitoFramework framework = new DefaultMockitoFramework();
+
     public boolean isTypeMockable(Class<?> typeToMock) {
         return typeMockabilityOf(typeToMock).mockable();
     }
@@ -60,7 +64,7 @@ public <T> T mock(Class<T> typeToMock, MockSettings settings) {
         MockSettingsImpl impl = MockSettingsImpl.class.cast(settings);
         MockCreationSettings<T> creationSettings = impl.build(typeToMock);
         T mock = createMock(creationSettings);
-        mockingProgress().mockingStarted(mock, creationSettings);
+        framework.mockingStarted(mock, creationSettings);
         return mock;
     }
 
@@ -91,7 +95,7 @@ public <T> T verify(T mock, VerificationMode mode) {
 
         MockingProgress mockingProgress = mockingProgress();
         VerificationMode actualMode = mockingProgress.maybeVerifyLazily(mode);
-        mockingProgress.verificationStarted(new MockAwareVerificationMode(mock, actualMode, mockingProgress.verificationListeners()));
+        mockingProgress.verificationStarted(new MockAwareVerificationMode(mock, actualMode, framework.verificationListeners()));
         return mock;
     }
 
@@ -215,4 +219,8 @@ public MockingDetails mockingDetails(Object toInspect) {
     public LenientStubber lenient() {
         return new DefaultLenientStubber();
     }
+
+    public DefaultMockitoFramework framework() {
+        return framework;
+    }
 }

src/main/java/org/mockito/internal/creation/bytebuddy/ByteBuddyMockMaker.java

@@ -40,6 +40,12 @@ public void resetMock(Object mock, MockHandler newHandler, MockCreationSettings
         defaultByteBuddyMockMaker.resetMock(mock, newHandler, settings);
     }
 
+    @Override
+    @Incubating
+    public void cleanUpMock(Object mock) {
+        defaultByteBuddyMockMaker.cleanUpMock(mock);
+    }
+
     @Override
     @Incubating
     public TypeMockability isTypeMockable(Class<?> type) {

src/main/java/org/mockito/internal/creation/bytebuddy/InlineByteBuddyMockMaker.java

@@ -271,6 +271,11 @@ public void resetMock(Object mock, MockHandler newHandler, MockCreationSettings
         }
     }
 
+    @Override
+    public void cleanUpMock(Object mock) {
+        mocks.remove(mock);
+    }
+
     @Override
     public TypeMockability isTypeMockable(final Class<?> type) {
         return new TypeMockability() {

src/main/java/org/mockito/internal/creation/bytebuddy/SubclassByteBuddyMockMaker.java

@@ -140,6 +140,11 @@ public void resetMock(Object mock, MockHandler newHandler, MockCreationSettings
         );
     }
 
+    @Override
+    public void cleanUpMock(Object mock) {
+        // We don't need to do anything. When there is no strong reference to the mock GC will just clean it up.
+    }
+
     @Override
     public TypeMockability isTypeMockable(final Class<?> type) {
         return new TypeMockability() {

src/main/java/org/mockito/internal/exceptions/Reporter.java

@@ -913,4 +913,11 @@ public static void unfinishedMockingSession() {
                 "Previous MockitoSession was not concluded with 'finishMocking()'.",
                 "For examples of correct usage see javadoc for MockitoSession class."));
     }
+
+    public static void multipleTrackingMockSession() {
+        throw new MultipleTrackingMockSessionException(join(
+            "Multiple open sessions that track and clean up mocks detected.",
+            "There can only be at most one open session that tracks and cleans up mocks,",
+            "please finish previous session that tracks and clean up mocks before opening a new one."));
+    }
 }

src/main/java/org/mockito/internal/framework/DefaultMockitoFramework.java

@@ -4,18 +4,28 @@
  */
 package org.mockito.internal.framework;
 
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
 import org.mockito.MockitoFramework;
 import org.mockito.internal.configuration.plugins.Plugins;
 import org.mockito.internal.invocation.DefaultInvocationFactory;
+import org.mockito.internal.listeners.MockitoListenerUtil;
 import org.mockito.internal.util.Checks;
 import org.mockito.invocation.InvocationFactory;
+import org.mockito.listeners.MockCreationListener;
 import org.mockito.listeners.MockitoListener;
+import org.mockito.listeners.VerificationListener;
+import org.mockito.mock.MockCreationSettings;
 import org.mockito.plugins.MockitoPlugins;
 
 import static org.mockito.internal.progress.ThreadSafeMockingProgress.mockingProgress;
 
 public class DefaultMockitoFramework implements MockitoFramework {
 
+    private final Set<MockitoListener> listeners = new LinkedHashSet<MockitoListener>();
+
     public MockitoFramework addListener(MockitoListener listener) {
         Checks.checkNotNull(listener, "listener");
         mockingProgress().addListener(listener);
@@ -28,6 +38,64 @@ public MockitoFramework removeListener(MockitoListener listener) {
         return this;
     }
 
+    /**
+     * Adds a listener to Mockito.
+     *
+     * <p>It's different from {@link #addListener(MockitoListener)} in the way that this listener will receive
+     * notifications from all threads.
+     *
+     * <p>It also throws {@link org.mockito.exceptions.misusing.RedundantListenerException} if caller tries to add a
+     * listener of the same type as an added listener.
+     *
+     * <p>The listener will be called in the thread where the notification is sent. It's the responsibility of listeners
+     * to ensure thread-safety.
+     *
+     * @param listener to add to Mockito
+     * @return this instance of mockito framework (fluent builder pattern)
+     */
+    public MockitoFramework addGlobalListener(MockitoListener listener) {
+        Checks.checkNotNull(listener, "listener");
+        synchronized (listeners) {
+            MockitoListenerUtil.addListener(listener, listeners);
+        }
+        return this;
+    }
+
+    public MockitoFramework removeGlobalListener(MockitoListener listener) {
+        Checks.checkNotNull(listener, "listener");
+        synchronized (listeners) {
+            listeners.remove(listener);
+        }
+        return this;
+    }
+
+    public Set<VerificationListener> verificationListeners() {
+        Set<VerificationListener> verificationListeners = mockingProgress().verificationListeners();
+        synchronized (listeners) {
+            for (MockitoListener listener : listeners) {
+                if (listener instanceof VerificationListener) {
+                    verificationListeners.add((VerificationListener) listener);
+                }
+            }
+        }
+        return verificationListeners;
+    }
+
+    public void mockingStarted(Object mock, MockCreationSettings settings) {
+        mockingProgress().mockingStarted(mock, settings);
+        List<MockCreationListener> creationListeners = new ArrayList<MockCreationListener>();
+        synchronized (listeners) {
+            for (MockitoListener listener : listeners) {
+                if (listener instanceof MockCreationListener) {
+                    creationListeners.add((MockCreationListener) listener);
+                }
+            }
+        }
+        for (MockCreationListener listener : creationListeners) {
+            listener.onMockCreated(mock, settings);
+        }
+    }
+
     @Override
     public MockitoPlugins getPlugins() {
         return Plugins.getPlugins();

src/main/java/org/mockito/internal/framework/DefaultMockitoSession.java

@@ -6,9 +6,12 @@
 
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
+import org.mockito.MockitoFramework;
 import org.mockito.MockitoSession;
 import org.mockito.exceptions.misusing.RedundantListenerException;
+import org.mockito.internal.configuration.plugins.Plugins;
 import org.mockito.internal.exceptions.Reporter;
+import org.mockito.internal.junit.MockTracker;
 import org.mockito.internal.junit.TestFinishedEvent;
 import org.mockito.internal.junit.UniversalTestListener;
 import org.mockito.plugins.MockitoLogger;
@@ -20,13 +23,16 @@ public class DefaultMockitoSession implements MockitoSession {
 
     private final String name;
     private final UniversalTestListener listener;
+    private final MockTracker mockTracker;
 
-    public DefaultMockitoSession(List<Object> testClassInstances, String name, Strictness strictness, MockitoLogger logger) {
+    public DefaultMockitoSession(List<Object> testClassInstances, String name, Strictness strictness, MockitoLogger logger,
+                                 boolean trackAndCleanUpMocks) {
         this.name = name;
         listener = new UniversalTestListener(strictness, logger);
+        MockitoFramework framework = Mockito.framework();
         try {
             //So that the listener can capture mock creation events
-            Mockito.framework().addListener(listener);
+            framework.addListener(listener);
         } catch (RedundantListenerException e) {
             Reporter.unfinishedMockingSession();
         }
@@ -39,6 +45,18 @@ public DefaultMockitoSession(List<Object> testClassInstances, String name, Stric
             listener.setListenerDirty();
             throw e;
         }
+
+        if (trackAndCleanUpMocks && (framework instanceof DefaultMockitoFramework)) {
+            mockTracker = new MockTracker(Plugins.getMockMaker());
+            try {
+                ((DefaultMockitoFramework) framework).addGlobalListener(mockTracker);
+            } catch (RedundantListenerException e) {
+                listener.setListenerDirty();
+                Reporter.multipleTrackingMockSession();
+            }
+        } else {
+            mockTracker = null;
+        }
     }
 
     @Override
@@ -56,7 +74,11 @@ public void finishMocking(final Throwable failure) {
         //Cleaning up the state, we no longer need the listener hooked up
         //The listener implements MockCreationListener and at this point
         //we no longer need to listen on mock creation events. We are wrapping up the session
-        Mockito.framework().removeListener(listener);
+        MockitoFramework framework = Mockito.framework();
+        framework.removeListener(listener);
+        if (mockTracker != null) {
+            ((DefaultMockitoFramework) framework).removeGlobalListener(mockTracker);
+        }
 
         //Emit test finished event so that validation such as strict stubbing can take place
         listener.testFinished(new TestFinishedEvent() {
@@ -75,5 +97,9 @@ public String getTestName() {
             //Finally, validate user's misuse of Mockito framework.
             Mockito.validateMockitoUsage();
         }
+
+        if (mockTracker != null) {
+            mockTracker.testFinished();
+        }
     }
 }

src/main/java/org/mockito/internal/junit/MockTracker.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2019 Mockito contributors
+ * This program is made available under the terms of the MIT License.
+ */
+
+package org.mockito.internal.junit;
+
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.List;
+import org.mockito.internal.listeners.AutoCleanableListener;
+import org.mockito.listeners.MockCreationListener;
+import org.mockito.mock.MockCreationSettings;
+import org.mockito.plugins.MockMaker;
+
+public class MockTracker implements MockCreationListener, AutoCleanableListener {
+
+    private final List<WeakReference<Object>> mocks = new ArrayList<WeakReference<Object>>();
+    private final MockMaker mockMaker;
+    private boolean listenerDirty;
+
+    public MockTracker(MockMaker mockMaker) {
+        this.mockMaker = mockMaker;
+    }
+
+    @Override
+    public void onMockCreated(Object mock, MockCreationSettings settings) {
+        synchronized (mocks) {
+            mocks.add(new WeakReference<Object>(mock));
+        }
+    }
+
+    public void testFinished() {
+        WeakReference[] localMocks;
+        synchronized (mocks) {
+            localMocks = mocks.toArray(new WeakReference[0]);
+        }
+
+        for (WeakReference weakMock : localMocks) {
+            Object mock = weakMock.get();
+            if (mock != null) {
+                mockMaker.cleanUpMock(mock);
+            }
+        }
+    }
+
+    @Override
+    public boolean isListenerDirty() {
+        return listenerDirty;
+    }
+
+    public void setListenerDirty() {
+        listenerDirty = true;
+    }
+}

src/main/java/org/mockito/internal/listeners/MockitoListenerUtil.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2019 Mockito contributors
+ * This program is made available under the terms of the MIT License.
+ */
+
+package org.mockito.internal.listeners;
+
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Set;
+import org.mockito.internal.exceptions.Reporter;
+import org.mockito.listeners.MockitoListener;
+
+public class MockitoListenerUtil {
+    public static void addListener(MockitoListener listener, Set<MockitoListener> listeners) {
+        List<MockitoListener> delete = new LinkedList<MockitoListener>();
+        for (MockitoListener existing : listeners) {
+            if (existing.getClass().equals(listener.getClass())) {
+                if (existing instanceof AutoCleanableListener && ((AutoCleanableListener) existing).isListenerDirty()) {
+                    //dirty listener means that there was an exception even before the test started
+                    //if we fail here with redundant mockito listener exception there will be multiple failures causing confusion
+                    //so we simply remove the existing listener and move on
+                    delete.add(existing);
+                } else {
+                    Reporter.redundantMockitoListener(listener.getClass().getSimpleName());
+                }
+            }
+        }
+        //delete dirty listeners so they don't occupy state/memory and don't receive notifications
+        for (MockitoListener toDelete : delete) {
+            listeners.remove(toDelete);
+        }
+        listeners.add(listener);
+    }
+}