Properly cleanup thread locals for non-CoroutineDispatcher-intercepte…

kotlinx-coroutines-core/common/src/AbstractCoroutine.kt

@@ -5,6 +5,7 @@ package kotlinx.coroutines
 import kotlinx.coroutines.CoroutineStart.*
 import kotlinx.coroutines.intrinsics.*
 import kotlin.coroutines.*
+import kotlinx.coroutines.internal.ScopeCoroutine
 
 /**
  * Abstract base class for implementation of coroutines in coroutine builders.
@@ -100,6 +101,15 @@ public abstract class AbstractCoroutine<in T>(
         afterResume(state)
     }
 
+    /**
+     * Invoked when the corresponding `AbstractCoroutine` was **conceptually** resumed, but not mechanically.
+     * Currently, this function only invokes `resume` on the underlying continuation for [ScopeCoroutine]
+     * or does nothing otherwise.
+     *
+     * Examples of resumes:
+     * - `afterCompletion` calls when the corresponding `Job` changed its state (i.e. got cancelled)
+     * - [AbstractCoroutine.resumeWith] was invoked
+     */
     protected open fun afterResume(state: Any?): Unit = afterCompletion(state)
 
     internal final override fun handleOnCompletionException(exception: Throwable) {

kotlinx-coroutines-core/common/src/internal/Scopes.kt

@@ -23,6 +23,13 @@ internal open class ScopeCoroutine<in T>(
         uCont.intercepted().resumeCancellableWith(recoverResult(state, uCont))
     }
 
+    /**
+     * Invoked when a scoped coorutine was completed in an undispatched manner directly
+     * at the place of its start because it never suspended.
+     */
+    open fun afterCompletionUndispatched() {
+    }
+
     override fun afterResume(state: Any?) {
         // Resume direct because scope is already in the correct context
         uCont.resumeWith(recoverResult(state, uCont))

kotlinx-coroutines-core/common/src/intrinsics/Undispatched.kt

@@ -79,6 +79,7 @@ private inline fun <T> ScopeCoroutine<T>.undispatchedResult(
     if (result === COROUTINE_SUSPENDED) return COROUTINE_SUSPENDED // (1)
     val state = makeCompletingOnce(result)
     if (state === COMPLETING_WAITING_CHILDREN) return COROUTINE_SUSPENDED // (2)
+    afterCompletionUndispatched()
     return if (state is CompletedExceptionally) { // (3)
         when {
             shouldThrow(state.cause) -> throw recoverStackTrace(state.cause, uCont)

kotlinx-coroutines-core/jvm/src/CoroutineContext.kt

@@ -253,18 +253,26 @@ internal actual class UndispatchedCoroutine<in T>actual constructor (
         }
     }
 
+    override fun afterCompletionUndispatched() {
+        clearThreadLocal()
+    }
+
     override fun afterResume(state: Any?) {
+        clearThreadLocal()
+        // resume undispatched -- update context but stay on the same dispatcher
+        val result = recoverResult(state, uCont)
+        withContinuationContext(uCont, null) {
+            uCont.resumeWith(result)
+        }
+    }
+
+    private fun clearThreadLocal() {
         if (threadLocalIsSet) {
             threadStateToRecover.get()?.let { (ctx, value) ->
                 restoreThreadContext(ctx, value)
             }
             threadStateToRecover.remove()
         }
-        // resume undispatched -- update context but stay on the same dispatcher
-        val result = recoverResult(state, uCont)
-        withContinuationContext(uCont, null) {
-            uCont.resumeWith(result)
-        }
     }
 }
 

kotlinx-coroutines-core/jvm/test/ThreadLocalsLeaksTest.kt

@@ -0,0 +1,95 @@
+package kotlinx.coroutines
+
+import kotlinx.coroutines.testing.TestBase
+import java.lang.ref.WeakReference
+import kotlin.coroutines.AbstractCoroutineContextElement
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.ContinuationInterceptor
+import kotlin.coroutines.CoroutineContext
+import kotlin.test.Test
+import kotlin.test.assertNull
+
+/*
+ * This is an adapted verion of test from #4296.
+ *
+ * qwwdfsad: the test relies on System.gc() actually collecting the garbage.
+ * If these tests flake on CI, first check that JDK/GC setup in not an issue.
+ */
+class ThreadLocalCustomContinuationInterceptorTest : TestBase() {
+
+    private class CustomContinuationInterceptor(private val delegate: ContinuationInterceptor) :
+        AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor {
+
+        override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> {
+            return delegate.interceptContinuation(continuation)
+        }
+    }
+
+    private class CustomNeverEqualContinuationInterceptor(private val delegate: ContinuationInterceptor) :
+        AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor {
+
+        override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> {
+            return delegate.interceptContinuation(continuation)
+        }
+
+        override fun equals(other: Any?) = false
+    }
+
+    @Test
+    fun testDefaultDispatcherNoSuspension() = ensureCoroutineContextGCed(Dispatchers.Default, suspend = false)
+
+    @Test
+    fun testDefaultDispatcher() = ensureCoroutineContextGCed(Dispatchers.Default, suspend = true)
+
+
+    @Test
+    fun testNonCoroutineDispatcher() = ensureCoroutineContextGCed(
+        CustomContinuationInterceptor(Dispatchers.Default),
+        suspend = true
+    )
+
+    @Test
+    fun testNonCoroutineDispatcherSuspension() = ensureCoroutineContextGCed(
+        CustomContinuationInterceptor(Dispatchers.Default),
+        suspend = false
+    )
+
+    // Note asymmetric equals codepath never goes through the undispatched withContext, thus the separate test case
+
+    @Test
+    fun testNonCoroutineDispatcherAsymmetricEquals() =
+        ensureCoroutineContextGCed(
+            CustomNeverEqualContinuationInterceptor(Dispatchers.Default),
+            suspend = true
+        )
+
+    @Test
+    fun testNonCoroutineDispatcherAsymmetricEqualsSuspension() =
+        ensureCoroutineContextGCed(
+            CustomNeverEqualContinuationInterceptor(Dispatchers.Default),
+            suspend = false
+        )
+
+
+    private fun ensureCoroutineContextGCed(coroutineContext: CoroutineContext, suspend: Boolean) {
+        runTest {
+            lateinit var ref: WeakReference<CoroutineName>
+            val job = GlobalScope.launch(coroutineContext) {
+                val coroutineName = CoroutineName("Yo")
+                ref = WeakReference(coroutineName)
+                withContext(coroutineName) {
+                    if (suspend) {
+                        delay(1)
+                    }
+                }
+            }
+            job.join()
+
+            // Twice is enough to ensure
+            System.gc()
+            System.gc()
+            assertNull(ref.get())
+        }
+    }
+
+}