From 4243bafd1d28dfb708235ccdeb5ac14e307f7354 Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Fri, 20 Aug 2021 17:29:22 +0200
Subject: [PATCH] Scripting: implement correct disposal of REPL state on
 finalization

#KT-47927 fixed
---
 .../jsr223/test/KotlinJsr223ScriptEngineIT.kt | 25 +++++++++++++++++++
 .../jsr223/KotlinJsr223ScriptEngineImpl.kt    |  4 ++-
 .../plugin/repl/jvmReplCompilation.kt         |  4 +++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/libraries/scripting/jsr223-test/test/kotlin/script/experimental/jsr223/test/KotlinJsr223ScriptEngineIT.kt b/libraries/scripting/jsr223-test/test/kotlin/script/experimental/jsr223/test/KotlinJsr223ScriptEngineIT.kt
index 04226333890..5adb72a6392 100644
--- a/libraries/scripting/jsr223-test/test/kotlin/script/experimental/jsr223/test/KotlinJsr223ScriptEngineIT.kt
+++ b/libraries/scripting/jsr223-test/test/kotlin/script/experimental/jsr223/test/KotlinJsr223ScriptEngineIT.kt
@@ -8,7 +8,9 @@ package kotlin.script.experimental.jsr223.test
 import org.jetbrains.kotlin.cli.common.environment.setIdeaIoUseFallback
 import org.jetbrains.kotlin.config.KotlinCompilerVersion
 import org.junit.Assert
+import org.junit.Ignore
 import org.junit.Test
+import java.lang.management.ManagementFactory
 import javax.script.*
 import kotlin.script.experimental.jvmhost.jsr223.KotlinJsr223ScriptEngineImpl
 
@@ -71,6 +73,29 @@ class KotlinJsr223ScriptEngineIT {
         Assert.assertEquals(5, res2)
     }
 
+    @Test
+    @Ignore // Probably not possible to make it sensible on CI and with parallel run, so leaving it here for manual testing only
+    fun testMemory() {
+        val memoryMXBean = ManagementFactory.getMemoryMXBean()!!
+        var prevMem = memoryMXBean.getHeapMemoryUsage().getUsed()
+        for (i in 1..10) {
+            with(ScriptEngineManager().getEngineByExtension("kts")) {
+                val res1 = eval("val x = 3")
+                Assert.assertNull(res1)
+                val res2 = eval("x + 2")
+                Assert.assertEquals(5, res2)
+            }
+            System.gc()
+            val curMem = memoryMXBean.getHeapMemoryUsage().getUsed()
+            if (i > 3 && curMem > prevMem) {
+                Assert.assertTrue("Memory leak: iter: $i prev: $prevMem, cur: $curMem", (curMem - prevMem) < 1024*1024 )
+            }
+            println("${curMem/1024/1024}Mb")
+            prevMem = curMem
+        }
+    }
+
+
     @Test
     fun testIncomplete() {
         val engine = ScriptEngineManager().getEngineByExtension("kts")!!
diff --git a/libraries/scripting/jvm-host/src/kotlin/script/experimental/jvmhost/jsr223/KotlinJsr223ScriptEngineImpl.kt b/libraries/scripting/jvm-host/src/kotlin/script/experimental/jvmhost/jsr223/KotlinJsr223ScriptEngineImpl.kt
index bb0b2c45874..25e06408b97 100644
--- a/libraries/scripting/jvm-host/src/kotlin/script/experimental/jvmhost/jsr223/KotlinJsr223ScriptEngineImpl.kt
+++ b/libraries/scripting/jvm-host/src/kotlin/script/experimental/jvmhost/jsr223/KotlinJsr223ScriptEngineImpl.kt
@@ -6,6 +6,7 @@
 package kotlin.script.experimental.jvmhost.jsr223
 
 import org.jetbrains.kotlin.cli.common.repl.*
+import java.lang.ref.WeakReference
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import javax.script.ScriptContext
 import javax.script.ScriptEngineFactory
@@ -32,8 +33,9 @@ class KotlinJsr223ScriptEngineImpl(
     private var lastScriptContext: ScriptContext? = null
 
     val jsr223HostConfiguration = ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration) {
+        val weakThis = WeakReference(this@KotlinJsr223ScriptEngineImpl)
         jsr223 {
-            getScriptContext { lastScriptContext ?: getContext() }
+            getScriptContext { weakThis.get()?.let { it.lastScriptContext ?: it.getContext() } }
         }
     }
 
diff --git a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/jvmReplCompilation.kt b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/jvmReplCompilation.kt
index b7869c3d0f7..3582e01a40f 100644
--- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/jvmReplCompilation.kt
+++ b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/jvmReplCompilation.kt
@@ -37,6 +37,10 @@ class JvmReplCompilerState<CompilationT : JvmReplCompilerState.Compilation>(
         }
     }
 
+    protected fun finalize() {
+        dispose()
+    }
+
     fun getCompilationState(scriptCompilationConfiguration: ScriptCompilationConfiguration): CompilationT = lock.write {
         if (_compilation == null) {
             initializeCompilation(scriptCompilationConfiguration)