~ Use WeakReference in Remove.ref to avoid leaks of removed chains

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

@@ -58,7 +58,7 @@ public suspend fun <T> withTimeoutOrNull(timeMillis: Long, block: suspend Corout
         }
     } catch (e: TimeoutCancellationException) {
         // Return null if it's our exception, otherwise propagate it upstream (e.g. in case of nested withTimeouts)
-        if (weakReferenceUnwrap(e.coroutine) === coroutine) {
+        if (e.coroutine.unweakRef() === coroutine) {
             return null
         }
         throw e
@@ -112,4 +112,4 @@ public class TimeoutCancellationException internal constructor(
 internal fun TimeoutCancellationException(
     time: Long,
     coroutine: Job
-) : TimeoutCancellationException = TimeoutCancellationException("Timed out waiting for $time ms", weakReference(coroutine))
+) : TimeoutCancellationException = TimeoutCancellationException("Timed out waiting for $time ms", coroutine.weakRef())

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

@@ -22,5 +22,5 @@ internal expect fun <T> CancellableContinuationImpl<T>.shareableResume(delegate:
 internal expect fun isReuseSupportedInPlatform(): Boolean
 internal expect fun <T> ArrayList<T>.addOrUpdate(element: T, update: (ArrayList<T>) -> Unit)
 internal expect fun <T> ArrayList<T>.addOrUpdate(index: Int, element: T, update: (ArrayList<T>) -> Unit)
-internal expect fun weakReference(obj: Any): Any
-internal expect fun weakReferenceUnwrap(ref: Any?): Any?
+internal expect fun Any.weakRef(): Any
+internal expect fun Any?.unweakRef(): Any?

kotlinx-coroutines-core/concurrent/src/internal/LockFreeLinkedList.kt

@@ -265,7 +265,8 @@ public actual open class LockFreeLinkedListNode {
     public actual fun helpRemove() {
         val next = this.next ?: return // unlinked node on Kotlin/Native
         val removed = next as? Removed ?: error("Must be invoked on a removed node")
-        finishRemove(removed.ref)
+        val ref = removed.ref ?: return // unlinked node on Kotlin/Native
+        finishRemove(ref)
     }
 
     public actual fun removeFirstOrNull(): Node? {
@@ -608,15 +609,16 @@ public actual open class LockFreeLinkedListNode {
             val nextNext = next.next ?: return // could be unlinked on Kotlin/Native
             if (nextNext is Removed) {
                 next.markPrev() ?: return // could be unlinked on Kotlin/Native
-                next = nextNext.ref
+                next = nextNext.ref ?: return // could be unlinked on Kotlin/Native
                 continue
             }
             // move the the left until first non-removed node
             val prevNext = prev.next ?: return // could be unlinked on Kotlin/Native
             if (prevNext is Removed) {
                 if (last != null) {
                     prev.markPrev() ?: return // could be unlinked on Kotlin/Native
-                    last._next.compareAndSet(prev, prevNext.ref)
+                    val ref = prevNext.ref ?: return // could be unlinked on Kotlin/Native
+                    last._next.compareAndSet(prev, ref)
                     prev = last
                     last = null
                 } else {
@@ -692,7 +694,9 @@ public actual open class LockFreeLinkedListNode {
     override fun toString(): String = "$classSimpleName@$hexAddress"
 }
 
-private class Removed(@JvmField val ref: Node) {
+private class Removed(ref: Node) {
+    private val wRef: Any = ref.weakRef()
+    val ref: Node? get() = wRef.unweakRef() as Node?
     override fun toString(): String = "Removed[$ref]"
 }
 

kotlinx-coroutines-core/js/src/internal/Sharing.kt

@@ -4,9 +4,11 @@
 
 package kotlinx.coroutines.internal
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
+import kotlin.internal.*
 
 @Suppress("ACTUAL_WITHOUT_EXPECT") // visibility different
 internal actual typealias ShareableRefHolder = Any
@@ -59,7 +61,7 @@ internal actual inline fun <T> ArrayList<T>.addOrUpdate(index: Int, element: T,
 }
 
 @Suppress("NOTHING_TO_INLINE") // Should be NOP
-internal actual inline fun weakReference(obj: Any): Any = obj
+internal actual inline fun Any.weakRef(): Any = this
 
 @Suppress("NOTHING_TO_INLINE") // Should be NOP
-internal actual inline fun weakReferenceUnwrap(ref: Any?): Any? = ref
+internal actual inline fun Any?.unweakRef(): Any? = this

kotlinx-coroutines-core/jvm/src/internal/Sharing.kt

@@ -6,6 +6,7 @@
 
 package kotlinx.coroutines.internal
 
+import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
@@ -77,8 +78,8 @@ internal actual inline fun <T> ArrayList<T>.addOrUpdate(index: Int, element: T,
 
 @InlineOnly
 @Suppress("NOTHING_TO_INLINE") // Should be NOP
-internal actual inline fun weakReference(obj: Any): Any = obj
+internal actual inline fun Any.weakRef(): Any = this
 
 @InlineOnly
 @Suppress("NOTHING_TO_INLINE") // Should be NOP
-internal actual inline fun weakReferenceUnwrap(ref: Any?): Any? = ref
+internal actual inline fun Any?.unweakRef(): Any? = this

kotlinx-coroutines-core/native/src/internal/Sharing.kt

@@ -110,9 +110,9 @@ internal actual inline fun <T> ArrayList<T>.addOrUpdate(index: Int, element: T,
 }
 
 @Suppress("NOTHING_TO_INLINE")
-internal actual inline fun weakReference(obj: Any): Any = WeakReference(obj)
+internal actual inline fun Any.weakRef(): Any = WeakReference(this)
 
-internal actual fun weakReferenceUnwrap(ref: Any?): Any? = (ref as WeakReference<Any>?)?.get()
+internal actual fun Any?.unweakRef(): Any? = (this as WeakReference<Any>?)?.get()
 
 internal open class ShareableObject<T : Any>(obj: T) {
     val thread: Thread = currentThread()