Fix suspend-resume race in DebugProbes

qwwdfsad · qwwdfsad · commit bb3f0a1d7be3 · 2023-02-02T15:16:20.000+01:00
### Reproducing scenario For the following code: ``` suspend fun foo() = yield() ``` the following bytecode is produced: ``` fun foo(...) { uCont.intercepted().dispatchUsingDispatcher() // 1 probeCoroutineSuspended() // 2 return COROUTINE_SUSPENDED // Unroll the stack } ``` And it is possible to observe the next 'probeCoroutineSuspended' **prior** to receiving 'probeCoroutineSuspended'. ### Steps taken To address this, a dedicated 'unmatchedResumes' field is introduced to the coroutine state, which keeps track of consecutive 'probeCoroutineResumed' without matching 'probeCoroutineSuspended' and attributes lately arrived probes to it. Unfortunately, it introduces a much more unlikely race when **two** 'probeCoroutineSuspended' are reordered, then we misattribute 'lastObservedFrame', but it is still better than misattributing the actual coroutine state. @volatile and @synchronized are also introduced to DebugCoroutineInfoImpl as previously they have been subject to data races as well Fixes #3193
diff --git a/kotlinx-coroutines-core/jvm/src/debug/internal/DebugCoroutineInfoImpl.kt b/kotlinx-coroutines-core/jvm/src/debug/internal/DebugCoroutineInfoImpl.kt
@@ -41,9 +41,86 @@ internal class DebugCoroutineInfoImpl(
      * Can be CREATED, RUNNING, SUSPENDED.
      */
     public val state: String get() = _state
+
     @Volatile
     private var _state: String = CREATED
 
+    /*
+     * How many consecutive unmatched 'updateState(RESUMED)' this object has received.
+     * It can be `> 1` in two cases:
+     *
+     * * The coroutine is finishing and its state is being unrolled in BaseContinuationImpl, see comment to DebugProbesImpl#callerInfoCache
+     *   Such resumes are not expected to be matched and are ignored.
+     * * We encountered suspend-resume race explained above, and we do wait for a match.
+     */
+    private var unmatchedResume = 0
+
+    /**
+     * Here we orchestrate overlapping state updates that are coming asynchronously.
+     * In a nutshell, `probeCoroutineSuspended` can arrive **later** than its matching `probeCoroutineResumed`,
+     * e.g. for the following code:
+     * ```
+     * suspend fun foo() = yield()
+     * ```
+     *
+     * we have this sequence:
+     * ```
+     * fun foo(...) {
+     *     uCont.intercepted().dispatchUsingDispatcher() // 1
+     *     // Notify the debugger the coroutine is suspended
+     *     probeCoroutineSuspended() // 2
+     *     return COROUTINE_SUSPENDED // Unroll the stack
+     * }
+     * ```
+     * Nothing prevents coroutine to be dispatched and invoke `probeCoroutineResumed` right between '1' and '2'.
+     * See also: https://github.com/Kotlin/kotlinx.coroutines/issues/3193
+     *
+     * [shouldBeMatched] -- `false` if it is an expected consecutive `probeCoroutineResumed` from BaseContinuationImpl,
+     * `true` otherwise.
+     */
+    @Synchronized
+    internal fun updateState(state: String, frame: Continuation<*>, shouldBeMatched: Boolean) {
+        /**
+         * We observe consecutive resume that had to be matched, but it wasn't,
+         * increment
+         */
+        if (_state == RUNNING && state == RUNNING && shouldBeMatched) {
+            ++unmatchedResume
+        } else if (unmatchedResume > 0 && state == SUSPENDED) {
+            /*
+             * We received late 'suspend' probe for unmatched resume, skip it.
+             * Here we deliberately allow the very unlikely race;
+             * Consider the following scenario ('[r:a]' means "probeCoroutineResumed at a()"):
+             * ```
+             * [r:a] a() -> b() [s:b] [r:b] -> (back to a) a() -> c() [s:c]
+             * ```
+             * We can, in theory, observe the following probes interleaving:
+             * ```
+             * r:a
+             * r:b // Unmatched resume
+             * s:c // Matched suspend, discard
+             * s:b
+             * ```
+             * Thus mis-attributing 'lastObservedFrame' to a previously-observed.
+             * It is possible in theory (though I've failed to reproduce it), yet
+             * is more preferred than indefinitely mismatched state (-> mismatched real/enhanced stacktrace)
+             */
+            --unmatchedResume
+            return
+        }
+
+        // Propagate only non-duplicating transitions to running, see KT-29997
+        if (_state == state && state == SUSPENDED && lastObservedFrame != null) return
+
+        _state = state
+        lastObservedFrame = frame as? CoroutineStackFrame
+        lastObservedThread = if (state == RUNNING) {
+            Thread.currentThread()
+        } else {
+            null
+        }
+    }
+
     @JvmField
     @Volatile
     internal var lastObservedThread: Thread? = null
@@ -56,7 +133,9 @@ internal class DebugCoroutineInfoImpl(
     private var _lastObservedFrame: WeakReference<CoroutineStackFrame>? = null
     internal var lastObservedFrame: CoroutineStackFrame?
         get() = _lastObservedFrame?.get()
-        set(value) { _lastObservedFrame = value?.let { WeakReference(it) } }
+        set(value) {
+            _lastObservedFrame = value?.let { WeakReference(it) }
+        }
 
     /**
      * Last observed stacktrace of the coroutine captured on its suspension or resumption point.
@@ -88,17 +167,5 @@ internal class DebugCoroutineInfoImpl(
         }
     }
 
-    internal fun updateState(state: String, frame: Continuation<*>) {
-        // Propagate only duplicating transitions to running for KT-29997
-        if (_state == state && state == SUSPENDED && lastObservedFrame != null) return
-        _state = state
-        lastObservedFrame = frame as? CoroutineStackFrame
-        lastObservedThread = if (state == RUNNING) {
-            Thread.currentThread()
-        } else {
-            null
-        }
-    }
-
     override fun toString(): String = "DebugCoroutineInfo(state=$state,context=$context)"
 }
diff --git a/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt b/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt
@@ -436,16 +436,18 @@ internal object DebugProbesImpl {
         // Lookup coroutine info in cache or by traversing stack frame
         val info: DebugCoroutineInfoImpl
         val cached = callerInfoCache.remove(frame)
+        val shouldBeMatchedWithProbeSuspended: Boolean
         if (cached != null) {
             info = cached
+            shouldBeMatchedWithProbeSuspended = false
         } else {
             info = frame.owner()?.info ?: return
+            shouldBeMatchedWithProbeSuspended = true
             // Guard against improper implementations of CoroutineStackFrame and bugs in the compiler
             val realCaller = info.lastObservedFrame?.realCaller()
             if (realCaller != null) callerInfoCache.remove(realCaller)
         }
-
-        info.updateState(state, frame as Continuation<*>)
+        info.updateState(state, frame as Continuation<*>, shouldBeMatchedWithProbeSuspended)
         // Do not cache it for proxy-classes such as ScopeCoroutines
         val caller = frame.realCaller() ?: return
         callerInfoCache[caller] = info
@@ -458,7 +460,7 @@ internal object DebugProbesImpl {
 
     private fun updateState(owner: CoroutineOwner<*>, frame: Continuation<*>, state: String) {
         if (!isInstalled) return
-        owner.info.updateState(state, frame)
+        owner.info.updateState(state, frame, true)
     }
 
     private fun Continuation<*>.owner(): CoroutineOwner<*>? = (this as? CoroutineStackFrame)?.owner()
diff --git a/kotlinx-coroutines-debug/test/DebugProbesTest.kt b/kotlinx-coroutines-debug/test/DebugProbesTest.kt
@@ -6,6 +6,7 @@ package kotlinx.coroutines.debug
 import kotlinx.coroutines.*
 import org.junit.Test
 import java.util.concurrent.*
+import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.test.*
 
 class DebugProbesTest : DebugTestBase() {
@@ -100,4 +101,59 @@ class DebugProbesTest : DebugTestBase() {
             verifyStackTrace(e, traces)
         }
     }
+
+    @Test
+    fun testMultipleConsecutiveProbeResumed() = runTest {
+        val job = launch {
+            expect(1)
+            foo()
+            expect(4)
+            delay(Long.MAX_VALUE)
+            expectUnreached()
+        }
+        yield()
+        yield()
+        expect(5)
+        val infos = DebugProbes.dumpCoroutinesInfo()
+        assertEquals(2, infos.size)
+        assertEquals(setOf(State.RUNNING, State.SUSPENDED), infos.map { it.state }.toSet())
+        job.cancel()
+        finish(6)
+    }
+
+    @Test
+    fun testMultipleConsecutiveProbeResumedAndLaterRunning() = runTest {
+        val reachedActiveStage = AtomicBoolean(false)
+        val job = launch(Dispatchers.Default) {
+            expect(1)
+            foo()
+            expect(4)
+            yield()
+            reachedActiveStage.set(true)
+            while (isActive) {
+                // Spin until test is done
+            }
+        }
+        while (!reachedActiveStage.get()) {
+            delay(10)
+        }
+        expect(5)
+        val infos = DebugProbes.dumpCoroutinesInfo()
+        assertEquals(2, infos.size)
+        assertEquals(setOf(State.RUNNING, State.RUNNING), infos.map { it.state }.toSet())
+        job.cancel()
+        finish(6)
+    }
+
+    private suspend fun foo() {
+        bar()
+        // Kill TCO
+        expect(3)
+    }
+
+
+    private suspend fun bar() {
+        yield()
+        expect(2)
+    }
 }
