Merge pull request #1018 from bakomchik/bugfix/1013/fix-oom

#1013 uncouple proxy objects from cancellable to make proxy available…
mockk · Jan 15, 2023 · 1419517 · 1419517
2 parents 9f236a9 + c19753c
commit 1419517
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 8 deletions.
diff --git a/modules/mockk-agent-api/src/jvmMain/kotlin/io/mockk/proxy/common/CancelableResult.kt b/modules/mockk-agent-api/src/jvmMain/kotlin/io/mockk/proxy/common/CancelableResult.kt
@@ -2,16 +2,19 @@ package io.mockk.proxy.common
 
 import io.mockk.proxy.Cancelable
 import io.mockk.proxy.MockKAgentException
+import java.lang.ref.WeakReference
 import java.util.concurrent.atomic.AtomicBoolean
 
 open class CancelableResult<T : Any>(
-    private val value: T? = null,
+    input:T?=null,
     private val cancelBlock: () -> Unit = {}
 ) : Cancelable<T> {
 
+    private val weakValue: WeakReference<T>? = input?.let { WeakReference(it) }
+
     val fired = AtomicBoolean()
 
-    override fun get() = value
+    override fun get() = weakValue?.get()
             ?: throw MockKAgentException("Value for this result is not assigned")
 
     override fun cancel() {
@@ -23,7 +26,7 @@ open class CancelableResult<T : Any>(
     fun <R : Any> withValue(value: R) = CancelableResult(value, cancelBlock)
 
     fun alsoOnCancel(block: () -> Unit) =
-        CancelableResult(value) {
+        CancelableResult(weakValue?.get()) {
             cancel()
             block()
         }

diff --git a/modules/mockk-agent-api/src/jvmMain/kotlin/io/mockk/proxy/common/ProxyMaker.kt b/modules/mockk-agent-api/src/jvmMain/kotlin/io/mockk/proxy/common/ProxyMaker.kt
@@ -21,6 +21,8 @@ import io.mockk.proxy.common.transformation.InlineInstrumentation
 import io.mockk.proxy.common.transformation.SubclassInstrumentation
 import io.mockk.proxy.common.transformation.TransformationRequest
 import io.mockk.proxy.common.transformation.TransformationType
+import java.lang.ref.SoftReference
+import java.lang.ref.WeakReference
 import java.lang.reflect.Method
 import java.lang.reflect.Modifier
 import java.lang.reflect.Proxy
@@ -68,11 +70,14 @@ class ProxyMaker(
 
             subclasser.setProxyHandler(proxy, handler)
             handlers[proxy] = handler
-
+            val callbackRef = WeakReference(proxy)
             return result
                 .withValue(proxy)
                 .alsoOnCancel {
-                    handlers.remove(proxy)
+                    callbackRef.get()?.let {
+                        handlers.remove(it)
+                    }
+
                 }
         } catch (e: Exception) {
             result.cancel()

diff --git a/modules/mockk-agent/src/jvmMain/kotlin/io/mockk/proxy/jvm/ProxyMaker.kt b/modules/mockk-agent/src/jvmMain/kotlin/io/mockk/proxy/jvm/ProxyMaker.kt
@@ -6,6 +6,8 @@ import io.mockk.proxy.common.transformation.InlineInstrumentation
 import io.mockk.proxy.common.transformation.TransformationRequest
 import io.mockk.proxy.common.transformation.TransformationType.SIMPLE
 import io.mockk.proxy.jvm.transformation.SubclassInstrumentation
+import java.lang.ref.SoftReference
+import java.lang.ref.WeakReference
 import java.lang.reflect.Method
 import java.lang.reflect.Modifier
 
@@ -46,15 +48,16 @@ internal class ProxyMaker(
             val proxy = instantiate(actualClass, proxyClass, useDefaultConstructor, instance)
 
             handlers[proxy] = handler
-
+            val callbackRef = WeakReference(proxy)
             return result
                 .withValue(proxy)
                 .alsoOnCancel {
-                    handlers.remove(proxy)
+                    callbackRef.get()?.let {
+                        handlers.remove(it)
+                    }
                 }
         } catch (e: Exception) {
             result.cancel()
-
             throw MockKAgentException("Instantiation exception", e)
         }
     }

diff --git a/modules/mockk/src/jvmTest/kotlin/io/mockk/it/ProxyWeakReferenceTest.kt b/modules/mockk/src/jvmTest/kotlin/io/mockk/it/ProxyWeakReferenceTest.kt
@@ -0,0 +1,24 @@
+package io.mockk.it
+
+import io.mockk.clearMocks
+import io.mockk.every
+import io.mockk.spyk
+import org.junit.jupiter.api.Test
+
+class ProxyWeakReferenceTest {
+    @Test
+    fun test() {
+        for (i in 0..1000) {
+            val spyk = spyk(LazySpringBeanWhichHoldsReferenceForBeanFactory(ByteArray(10 * 1024 * 1024)))
+            every { spyk.doSmth() } returns "Wow"
+            spyk.doSmth()
+            clearMocks(spyk)
+        }
+    }
+
+    class LazySpringBeanWhichHoldsReferenceForBeanFactory(val beanFactory: ByteArray) {
+        fun doSmth(): String {
+            return "Smth"
+        }
+    }
+}