From 7f736410a6d4a78c693ade70bfa90081d3e464a7 Mon Sep 17 00:00:00 2001
From: Rafael Winterhalter <rafael.wth@gmail.com>
Date: Thu, 3 Sep 2020 22:47:13 +0200
Subject: [PATCH] Escape mock during method dispatch on mock to avoid premature
 garbage collection.

Fixes #1802.
---
 .../bytebuddy/MockMethodInterceptor.java      | 34 +++++++++++----
 .../PrematureGarbageCollectionTest.java       | 43 +++++++++++++++++++
 2 files changed, 69 insertions(+), 8 deletions(-)
 create mode 100644 src/test/java/org/mockito/PrematureGarbageCollectionTest.java

diff --git a/src/main/java/org/mockito/internal/creation/bytebuddy/MockMethodInterceptor.java b/src/main/java/org/mockito/internal/creation/bytebuddy/MockMethodInterceptor.java
index 532ecc804d..b506e3e415 100644
--- a/src/main/java/org/mockito/internal/creation/bytebuddy/MockMethodInterceptor.java
+++ b/src/main/java/org/mockito/internal/creation/bytebuddy/MockMethodInterceptor.java
@@ -36,6 +36,8 @@ public class MockMethodInterceptor implements Serializable {
 
     private final ByteBuddyCrossClassLoaderSerializationSupport serializationSupport;
 
+    private final ThreadLocal<Object> weakReferenceHatch = new ThreadLocal<>();
+
     public MockMethodInterceptor(MockHandler handler, MockCreationSettings mockCreationSettings) {
         this.handler = handler;
         this.mockCreationSettings = mockCreationSettings;
@@ -54,14 +56,30 @@ 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.
+        weakReferenceHatch.set(mock);
+        try {
+            return handler.handle(
+                    createInvocation(
+                            mock,
+                            invokedMethod,
+                            arguments,
+                            realMethod,
+                            mockCreationSettings,
+                            location));
+        } finally {
+            weakReferenceHatch.remove();
+        }
     }
 
     public MockHandler getMockHandler() {
diff --git a/src/test/java/org/mockito/PrematureGarbageCollectionTest.java b/src/test/java/org/mockito/PrematureGarbageCollectionTest.java
new file mode 100644
index 0000000000..cb10a9298a
--- /dev/null
+++ b/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);
+    }
+}