Coroutines debugger should keep weak references to running coroutines

It should not prevent garbage-collection of coroutines that were otherwise lost, which included the following practically-useful cases:
* Synchronous coroutines (iterator/sequence).
* Lazy coroutines that were not started.
* Abandoned coroutines that suspend forever without strong references to them in GlobalScope.

Two kinds of tests cover this functionality:
* A test via FieldWalker ensures that debugger impl does not keep a strong reference. This tests works fast and provides good diagnostics if anything goes wrong, but it is fragile, as futures changes to debugger my introduce static references to running coroutines elsewhere.
* A stress-test that ensures that no OOM indeed happens when you run a lot of such lost coroutines. Longer-running, more stable to code change, but fragile in a difference sense as it may accidentally start passing in the future if lots of memory get allocated for tests.

Fixes #2117

Author: elizarov

kotlinx-coroutines-core/jvm/src/debug/internal/ConcurrentWeakMap.kt

@@ -0,0 +1,280 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.debug.internal
+
+import kotlinx.atomicfu.*
+import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.internal.*
+import java.lang.ref.*
+import java.util.concurrent.atomic.*
+
+// This is very limited implementation, not suitable as a generic map replacement.
+// It has lock-free get and put with synchronized rehash for simplicity (and better CPU usage on contention)
+@OptIn(ExperimentalStdlibApi::class)
+@Suppress("UNCHECKED_CAST")
+internal class ConcurrentWeakMap<K : Any, V: Any> : AbstractMutableMap<K, V>() {
+    private val _size = atomic(0)
+    private val core = atomic(Core(MIN_CAPACITY))
+    private val queue = ReferenceQueue<K>()
+
+    override val size: Int
+        get() = _size.value
+
+    private fun decrementSize() { _size.decrementAndGet() }
+
+    /**
+     * Proactively checks if any already-garbage-collected memory can be released.
+     * Calling it periodically is not strictly required, as it is also checked on every put and rehash,
+     * but calling it periodically helps in reducing memory utilization.
+     */
+    fun cleanWeakRefs() {
+        core.value.cleanWeakRefs()
+    }
+
+    override fun get(key: K): V? = core.value.getImpl(key)
+
+    override fun put(key: K, value: V): V? {
+        var oldValue = core.value.putImpl(key, value)
+        if (oldValue === REHASH) oldValue = putSynchronized(key, value)
+        if (oldValue == null) _size.incrementAndGet()
+        return oldValue as V?
+    }
+
+    @Synchronized
+    private fun putSynchronized(key: K, value: V): V? {
+        // Note: concurrent put leaves chance that we fail to put even after rehash, we retry until successful
+        var curCore = core.value
+        while (true) {
+            val oldValue = curCore.putImpl(key, value)
+            if (oldValue !== REHASH) return oldValue as V?
+            curCore = curCore.rehash()
+            core.value = curCore
+        }
+    }
+
+    override fun remove(key: K): V? {
+        val oldValue = core.value.putImpl(key, null)
+        assert(oldValue !== REHASH) // cannot happen when removing
+        if (oldValue != null) _size.decrementAndGet()
+        return oldValue as V?
+    }
+
+    override val keys: MutableSet<K>
+        get() = KeyValueSet { k, _ -> k }
+
+    override val entries: MutableSet<MutableMap.MutableEntry<K, V>>
+        get() = KeyValueSet { k, v -> Entry(k, v) }
+
+    // We don't care much about clear's efficiency
+    override fun clear() {
+        for (k in keys) remove(k)
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    private inner class Core(private val allocated: Int) {
+        private val shift = allocated.countLeadingZeroBits() + 1
+        private val threshold = 2 * allocated / 3 // max fill factor at 66% to ensure speedy lookups
+        private val load = atomic(0) // counts how many slots are occupied in this core
+        private val keys = AtomicReferenceArray<HashedWeakRef<K>?>(allocated)
+        private val values = AtomicReferenceArray<Any?>(allocated)
+
+        private fun index(hash: Int) = (hash * MAGIC) ushr shift
+
+        // get is always lock-free, unwraps the value that was marked by concurrent rehash
+        fun getImpl(key: K): V? {
+            var index = index(key.hashCode())
+            while (true) {
+                val w = keys[index] ?: return null // not found
+                val k = w.get()
+                if (k == key) {
+                    val value = values[index]
+                    return (if (value is Marked) value.ref else value) as V?
+                }
+                if (k == null) removeCleanedAt(index) // weak ref was here, but collected
+                if (index == 0) index = allocated
+                index--
+            }
+        }
+
+        private fun removeCleanedAt(index: Int) {
+            while (true) {
+                val oldValue = values[index] ?: return // return when already removed
+                if (oldValue is Marked) return // cannot remove marked (rehash is working on it, will not copy)
+                if (values.compareAndSet(index, oldValue, null)) { // removed
+                    decrementSize()
+                    return
+                }
+            }
+        }
+
+        // returns REHASH when rehash is needed (the value was not put)
+        fun putImpl(key: K, value: V?, weakKey0: HashedWeakRef<K>? = null): Any? {
+            cleanWeakRefs()
+            var index = index(key.hashCode())
+            var loadIncremented = false
+            var weakKey: HashedWeakRef<K>? = weakKey0
+            while (true) {
+                val w = keys[index]
+                if (w == null) { // slot empty => not found => try reserving slot
+                    if (value == null) return null // removing missing value, nothing to do here
+                    if (!loadIncremented) {
+                        // We must increment load before we even try to occupy a slot to avoid overfill during concurrent put
+                        if (load.incrementAndGet() >= threshold) return REHASH
+                        loadIncremented = true
+                    }
+                    if (weakKey == null) weakKey = HashedWeakRef(key, queue)
+                    if (keys.compareAndSet(index, null, weakKey)) break // slot reserved !!!
+                    continue // retry at this slot on CAS failure (somebody already reserved this slot)
+                }
+                val k = w.get()
+                if (k == key) { // found already reserved slot at index
+                    if (loadIncremented) load.decrementAndGet() // undo increment, because found a slot
+                    break
+                }
+                if (k == null) removeCleanedAt(index) // weak ref was here, but collected
+                if (index == 0) index = allocated
+                index--
+            }
+            // update value
+            var oldValue: Any?
+            while (true) {
+                oldValue = values[index]
+                if (oldValue is Marked) return REHASH // rehash started, cannot work here
+                if (values.compareAndSet(index, oldValue, value)) break
+            }
+            return oldValue as V?
+        }
+
+        // only one thread can rehash, put maybe concurrent puts/gets
+        fun rehash(): Core {
+            // use size to approximate new required capacity to have at least 25-50% fill factor,
+            // may fail due to concurrent modification, will retry
+            retry@while (true) {
+                val newCapacity = size.coerceAtLeast(MIN_CAPACITY / 4).takeHighestOneBit() * 4
+                val newCore = Core(newCapacity)
+                for (index in 0 until allocated) {
+                    // load the key
+                    val w = keys[index]
+                    val k = w?.get()
+                    if (w != null && k == null) removeCleanedAt(index) // weak ref was here, but collected
+                    // mark value so that it cannot be changed while we rehash to new core
+                    var value: Any?
+                    while (true) {
+                        value = values[index]
+                        if (value is Marked) { // already marked -- good
+                            value = value.ref
+                            break
+                        }
+                        // try mark
+                        if (values.compareAndSet(index, value, value.mark())) break
+                    }
+                    if (k != null && value != null) {
+                        val oldValue = newCore.putImpl(k, value as V, w)
+                        if (oldValue === REHASH) continue@retry // retry if we underestimated capacity
+                        assert(oldValue == null)
+                    }
+                }
+                return newCore // rehashed everything successfully
+            }
+        }
+
+        fun cleanWeakRefs() {
+            while (true) {
+                cleanRef(queue.poll() as HashedWeakRef<K>? ?: return)
+            }
+        }
+
+        private fun cleanRef(weakRef: HashedWeakRef<K>) {
+            var index = index(weakRef.hash)
+            while (true) {
+                val w = keys[index] ?: return // return when slots are over
+                if (w === weakRef) { // found
+                    removeCleanedAt(index)
+                    return
+                }
+                if (index == 0) index = allocated
+                index--
+            }
+        }
+
+        fun <E> keyValueIterator(factory: (K, V) -> E): MutableIterator<E> {
+            cleanWeakRefs()
+            return KeyValueIterator(factory)
+        }
+
+        private inner class KeyValueIterator<E>(private val factory: (K, V) -> E) : MutableIterator<E> {
+            private var index = -1
+            private lateinit var key: K
+            private lateinit var value: V
+
+            init { findNext() }
+
+            private fun findNext() {
+                while (++index < allocated) {
+                    key = keys[index]?.get() ?: continue
+                    var value = values[index]
+                    if (value is Marked) value = value.ref
+                    if (value != null) {
+                        this.value = value as V
+                        return
+                    }
+                }
+            }
+
+            override fun hasNext(): Boolean = index < allocated
+
+            override fun next(): E {
+                if (index >= allocated) throw NoSuchElementException()
+                return factory(key, value).also { findNext() }
+            }
+
+            override fun remove() = noImpl()
+        }
+    }
+
+    private class Entry<K, V>(override val key: K, override val value: V) : MutableMap.MutableEntry<K, V> {
+        override fun setValue(newValue: V): V = noImpl()
+    }
+
+    private inner class KeyValueSet<E>(
+        private val factory: (K, V) -> E
+    ) : AbstractMutableSet<E>() {
+        override val size: Int get() = this@ConcurrentWeakMap.size
+        override fun add(element: E): Boolean = noImpl()
+        override fun iterator(): MutableIterator<E> = core.value.keyValueIterator(factory)
+    }
+
+    private fun noImpl(): Nothing {
+        throw UnsupportedOperationException("no implemented")
+    }
+}
+
+private const val MAGIC = 2654435769L.toInt() // golden ratio
+private const val MIN_CAPACITY = 16
+private val REHASH = Symbol("REHASH")
+private val MARKED_NULL = Marked(null)
+private val MARKED_TRUE = Marked(true) // When using map as set "true" used as value, optimize its mark allocation
+
+/**
+ * Weak reference that stores the original hash code so that we can use reference queue to promptly clean them up
+ * from the hashtable even in the absence of ongoing modifications.
+ */
+private class HashedWeakRef<T>(ref: T, queue: ReferenceQueue<T>) : WeakReference<T>(ref, queue) {
+    @JvmField
+    val hash = ref.hashCode()
+}
+
+/**
+ * Marked values cannot be modified. The marking is performed when rehash has started to ensure that concurrent
+ * modifications (that are lock-free) cannot perform any changes and are forced to synchronize with ongoing rehash.
+ */
+private class Marked(@JvmField val ref: Any?)
+
+private fun Any?.mark(): Marked = when(this) {
+    null -> MARKED_NULL
+    true -> MARKED_TRUE
+    else -> Marked(this)
+}
+

kotlinx-coroutines-core/jvm/src/debug/internal/DebugCoroutineInfo.kt

@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines.debug.internal
 
+import java.lang.ref.*
 import kotlin.coroutines.*
 import kotlin.coroutines.jvm.internal.*
 
@@ -12,10 +13,18 @@ internal const val RUNNING = "RUNNING"
 internal const val SUSPENDED = "SUSPENDED"
 
 internal class DebugCoroutineInfo(
-    public val context: CoroutineContext,
+    context: CoroutineContext?,
     public val creationStackBottom: CoroutineStackFrame?,
     @JvmField internal val sequenceNumber: Long
 ) {
+    /**
+     * We cannot keep a strong reference to the context, because with the [Job] in the context it will indirectly
+     * keep a reference to the last frame of an abandoned coroutine which the debugger should not be preventing
+     * garbage-collection of. The reference to context will not disappear as long as the coroutine itself is not lost.
+     */
+    private val _context = WeakReference(context)
+    public val context: CoroutineContext?
+        get() = _context.get()
 
     public val creationStackTrace: List<StackTraceElement> get() = creationStackTrace()
 
@@ -28,8 +37,15 @@ internal class DebugCoroutineInfo(
 
     @JvmField
     internal var lastObservedThread: Thread? = null
-    @JvmField
-    internal var lastObservedFrame: CoroutineStackFrame? = null
+
+    /**
+     * We cannot keep a strong reference to the last observed frame of the coroutine, because this will
+     * prevent garbage-collection of a coroutine that was lost.
+     */
+    private var _lastObservedFrame: WeakReference<CoroutineStackFrame>? = null
+    internal var lastObservedFrame: CoroutineStackFrame?
+        get() = _lastObservedFrame?.get()
+        set(value) { _lastObservedFrame = value?.let { WeakReference(it) } }
 
     public fun copy(): DebugCoroutineInfo = DebugCoroutineInfo(
         context,