Escape mock during method dispatch on mock to avoid premature garbage collection.

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

@@ -6,6 +6,8 @@
 
 import static org.mockito.internal.invocation.DefaultInvocationFactory.createInvocation;
 
+import java.io.IOException;
+import java.io.ObjectInputStream;
 import java.io.ObjectStreamException;
 import java.io.Serializable;
 import java.lang.reflect.Method;
@@ -36,12 +38,19 @@ public class MockMethodInterceptor implements Serializable {
 
     private final ByteBuddyCrossClassLoaderSerializationSupport serializationSupport;
 
+    private transient ThreadLocal<Object> weakReferenceHatch = new ThreadLocal<>();
+
     public MockMethodInterceptor(MockHandler handler, MockCreationSettings mockCreationSettings) {
         this.handler = handler;
         this.mockCreationSettings = mockCreationSettings;
         serializationSupport = new ByteBuddyCrossClassLoaderSerializationSupport();
     }
 
+    private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
+        stream.defaultReadObject();
+        weakReferenceHatch = new ThreadLocal<>();
+    }
+
     Object doIntercept(Object mock, Method invokedMethod, Object[] arguments, RealMethod realMethod)
             throws Throwable {
         return doIntercept(mock, invokedMethod, arguments, realMethod, new LocationImpl());
@@ -54,14 +63,33 @@ Object doIntercept(
             RealMethod realMethod,
             Location location)
             throws Throwable {
-        return handler.handle(
-                createInvocation(
-                        mock,
-                        invokedMethod,
-                        arguments,
-                        realMethod,
-                        mockCreationSettings,
-                        location));
+        // If the currently dispatched method is used in a hot path, typically a tight loop and if
+        // the mock is not used after the currently dispatched method, the JVM might attempt a
+        // garbage collection of the mock instance even before the execution of the current
+        // method is completed. Since we only reference the mock weakly from hereon after to avoid
+        // leaking the instance, it might therefore be garbage collected before the
+        // handler.handle(...) method completes. Since the handler method expects the mock to be
+        // present while a method call onto the mock is dispatched, this can lead to the problem
+        // described in GitHub #1802.
+        //
+        // To avoid this problem, we distract the JVM JIT by escaping the mock instance to a thread
+        // local field for the duration of the handler's dispatch.
+        //
+        // When dropping support for Java 8, instead of this hatch we should use an explicit fence
+        // https://docs.oracle.com/javase/9/docs/api/java/lang/ref/Reference.html#reachabilityFence-java.lang.Object-
+        weakReferenceHatch.set(mock);
+        try {
+            return handler.handle(
+                    createInvocation(
+                            mock,
+                            invokedMethod,
+                            arguments,
+                            realMethod,
+                            mockCreationSettings,
+                            location));
+        } finally {
+            weakReferenceHatch.remove();
+        }
     }
 
     public MockHandler getMockHandler() {

src/test/java/org/mockito/PrematureGarbageCollectionTest.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2020 Mockito contributors
+ * This program is made available under the terms of the MIT License.
+ */
+package org.mockito;
+
+import org.junit.Test;
+
+public class PrematureGarbageCollectionTest {
+
+    @Test
+    public void provoke_premature_garbage_collection() {
+        for (int i = 0; i < 500; i++) {
+            populateNodeList();
+        }
+    }
+
+    private static void populateNodeList() {
+        Node node = nodes();
+        while (node != null) {
+            Node next = node.next;
+            node.object.run();
+            node = next;
+        }
+    }
+
+    private static Node nodes() {
+        Node node = null;
+        for (int i = 0; i < 1_000; ++i) {
+            Node next = new Node();
+            next.next = node;
+            node = next;
+        }
+        return node;
+    }
+
+    private static class Node {
+
+        private Node next;
+
+        private final Runnable object = Mockito.mock(Runnable.class);
+    }
+}