Enhance support for async stack traces in flows

* simplify instrumentation by making a single insertion point source instead of having one in every class
* handle a double-wrapping case which leads to errors; allow agent to choose how to handle it
* support more commonly used operators (such as `scan`, `buffer`, `debounce` with dynamic timeout)

Unfortunately, this change doesn't cover all possible scenarios of using flows, as many of them interoperate with `Channel`s, and it should be addressed separately.

Author: AlexVanGogen

IntelliJ-patches.md

@@ -61,15 +61,29 @@ The agent needs three entities to establish a proper asynchronous stack traces c
 
 The key for MutableStateFlow is the element itself. For MutableSharedFlow, the element is wrapped into a unique object to prevent bridging mistakes when two equal elements are emitted from different places.
 
-Also, operators `debounce` and `sample` are supported. As they use an intermediary flow inside, the transferred values are wrapped and unwrapped the same way as in MutableSharedFlow.
+Most of the operators applicable to flows (such as `map`, `scan`, `debounce`, `buffer`) are supported. As some of them use an intermediary flow inside, the transferred values are wrapped and unwrapped the same way as in MutableSharedFlow.
 It means there may be all-library async stack traces between a stack trace containing `emit` and a stack trace containing `collect`.
 
+There is no support yet for many operators that heavily use `Channel`s inside (such as `timeout`), as well as for functions that convert flows to channels and vice versa (such as `produceIn`).
+
 ### API
 
-No new public methods are introduced; some logic was extracted to separate methods so that the debugger agent could instrument it properly:
+Some logic related to instrumentation was extracted to separate methods so that the debugger agent could instrument it properly:
 
 - `kotlinx.coroutines.flow.internal.FlowValueWrapperInternal` -- wrapper class used to create a unique object for the debugger agent
 - `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.wrapInternal` -- returns passed argument by default; the agent instruments it to call `wrapInternalDebuggerCapture` instead
-- `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.wrapInternalDebuggerCapture` -- wraps passed arguments into a `FlowValueWrapperInternal`; only used after transformation
+- `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.wrapInternalDebuggerCaptureX` -- wraps passed arguments into a `FlowValueWrapperInternal`; only used after transformation.
+  `X` may mean `Strict` or `Lenient`. Both methods handle double-wrapping, which is always an error that is hard to investigate if it arises naturally.
+    - `Strict` throws an exception, thus allowing fail-fast strategy
+    - `Lenient` returns its argument without wrapping it again
+  Debugger agent decides which version to use based on IDE settings.
 - `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.unwrapInternal` -- returns passed argument by default; the agent instruments it to call `unwrapInternalDebuggerCapture` instead
 - `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.unwrapInternalDebuggerCapture` -- unwraps passed argument so it returns the original value; only used after transformation
+- `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.emitInternal(FlowCollector, value)` -- common insertion point for a debugger agent; simplifies instrumentation; the value is always being unwrapped inside
+
+One internal method was added to `BufferedChannelIterator`: `nextInternal` -- same as `next` but may return a wrapped value. It should only be used with a function that is capable of unwrapping the value (see `BufferedChannel.emitAll` and `BufferedChannelIterator.next`), so there's a guarantee a wrapped value will always unwrap before emitting.
+
+Why not just let `next` return a maybe wrapped value? That's because it is heavily used outside a currently supported scope. For example, one may just indirectly call it from a for-loop. In this case, unwrapping will never happen, and a user will get a handful of `ClassCastException`s.
+
+One public method was added to support `buffer` and operators that use it inside:
+- `ReceiveChannel.emitAll`. It encapsulates emitting values in `FlowCollector.emitAllImpl` and has a special implementation in `BufferedChannel`.

kotlinx-coroutines-core/common/src/channels/BufferedChannel.kt

@@ -7,6 +7,7 @@ import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.ChannelResult.Companion.closed
 import kotlinx.coroutines.channels.ChannelResult.Companion.failure
 import kotlinx.coroutines.channels.ChannelResult.Companion.success
+import kotlinx.coroutines.flow.FlowCollector
 import kotlinx.coroutines.flow.internal.*
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.selects.*
@@ -1550,6 +1551,13 @@ internal open class BufferedChannel<E>(
 
     override fun iterator(): ChannelIterator<E> = BufferedChannelIterator()
 
+    override suspend fun emitAll(collector: FlowCollector<E>) {
+        val iterator = iterator() as BufferedChannel.BufferedChannelIterator
+        while (iterator.hasNext()) {
+            collector.emitInternal(iterator.nextInternal())
+        }
+    }
+
     /**
      * The key idea is that an iterator is a special receiver type,
      * which should be resumed differently to [receive] and [onReceive]
@@ -1601,7 +1609,7 @@ internal open class BufferedChannel<E>(
                 // Also, inform the `BufferedChannel` extensions that
                 // the synchronization of this receive operation is completed.
                 onElementRetrieved = { element ->
-                    this.receiveResult = element
+                    saveReceiveResult(element)
                     true
                 },
                 // As no waiter is provided, suspension is impossible.
@@ -1639,7 +1647,7 @@ internal open class BufferedChannel<E>(
                 // In case `onUndeliveredElement` is present, we must
                 // invoke it in the latter case.
                 onElementRetrieved = { element ->
-                    this.receiveResult = element
+                    saveReceiveResult(element)
                     this.continuation = null
                     cont.resume(true, onUndeliveredElement?.bindCancellationFun(element, cont.context))
                 },
@@ -1669,8 +1677,17 @@ internal open class BufferedChannel<E>(
             }
         }
 
-        @Suppress("UNCHECKED_CAST")
         override fun next(): E {
+            return unwrapInternal(nextInternal())
+        }
+
+        /**
+         * Result may be wrapped by debugger agent; use this method only with [unwrapInternal] or [emitInternal]!
+         *
+         * @see [next], [emitAll]
+         */
+        @Suppress("UNCHECKED_CAST")
+        internal fun nextInternal(): E {
             // Read the already received result, or [NO_RECEIVE_RESULT] if [hasNext] has not been invoked yet.
             val result = receiveResult
             check(result !== NO_RECEIVE_RESULT) { "`hasNext()` has not been invoked" }
@@ -1687,13 +1704,17 @@ internal open class BufferedChannel<E>(
             val cont = this.continuation!!
             this.continuation = null
             // Store the retrieved element in `receiveResult`.
-            this.receiveResult = element
+            saveReceiveResult(element)
             // Try to resume this `hasNext()`. Importantly, the receiver coroutine
             // may be cancelled after it is successfully resumed but not dispatched yet.
             // In case `onUndeliveredElement` is specified, we need to invoke it in the latter case.
             return cont.tryResume0(true, onUndeliveredElement?.bindCancellationFun(element, cont.context))
         }
 
+        private fun saveReceiveResult(element: E) {
+            this.receiveResult = wrapInternal(element)
+        }
+
         fun tryResumeHasNextOnClosedChannel() {
             /*
              * Read the current continuation of the suspended `hasNext()` call and clean the corresponding field to avoid memory leaks.

kotlinx-coroutines-core/common/src/channels/Channel.kt

@@ -8,6 +8,7 @@ import kotlinx.coroutines.channels.Channel.Factory.CHANNEL_DEFAULT_CAPACITY
 import kotlinx.coroutines.channels.Channel.Factory.CONFLATED
 import kotlinx.coroutines.channels.Channel.Factory.RENDEZVOUS
 import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
+import kotlinx.coroutines.flow.FlowCollector
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.selects.*
 import kotlin.contracts.*
@@ -288,6 +289,15 @@ public interface ReceiveChannel<out E> {
      */
     public operator fun iterator(): ChannelIterator<E>
 
+    /**
+     * Emits all elements of this channel using [collector].
+     */
+    public suspend fun emitAll(collector: FlowCollector<E>) {
+        for (element in this) {
+            collector.emit(element)
+        }
+    }
+
     /**
      * Cancels reception of remaining elements from this channel with an optional [cause].
      * This function closes the channel and removes all buffered sent elements from it.

kotlinx-coroutines-core/common/src/flow/Channels.kt

@@ -29,9 +29,7 @@ private suspend fun <T> FlowCollector<T>.emitAllImpl(channel: ReceiveChannel<T>,
     ensureActive()
     var cause: Throwable? = null
     try {
-        for (element in channel) {
-            emit(element)
-        }
+        channel.emitAll(this)
     } catch (e: Throwable) {
         cause = e
         throw e

kotlinx-coroutines-core/common/src/flow/SharedFlow.kt

@@ -389,18 +389,13 @@ internal open class SharedFlowImpl<T>(
                     awaitValue(slot) // await signal that the new value is available
                 }
                 collectorJob?.ensureActive()
-                emitInner(collector, newValue as T)
+                collector.emitInternal(newValue as T)
             }
         } finally {
             freeSlot(slot)
         }
     }
 
-    // Shouldn't be inlined, the method is instrumented by the IDEA debugger agent
-    private suspend fun emitInner(collector: FlowCollector<T>, value: T) {
-        collector.emit(unwrapInternal(value))
-    }
-
     override fun tryEmit(value: T): Boolean {
         var resumes: Array<Continuation<Unit>?> = EMPTY_RESUMES
         val emitted = synchronized(this) {

kotlinx-coroutines-core/common/src/flow/StateFlow.kt

@@ -394,7 +394,7 @@ private class StateFlowImpl<T>(
                 collectorJob?.ensureActive()
                 // Conflate value emissions using equality
                 if (oldState == null || oldState != newState) {
-                    emitInner(collector, newState)
+                    collector.emitInternal(newState)
                     oldState = newState
                 }
                 // Note: if awaitPending is cancelled, then it bails out of this loop and calls freeSlot
@@ -407,11 +407,6 @@ private class StateFlowImpl<T>(
         }
     }
 
-    // Shouldn't be inlined, the method is instrumented by the IDEA debugger agent
-    private suspend fun emitInner(collector: FlowCollector<T>, newState: Any) {
-        collector.emit(NULL.unbox(newState))
-    }
-
     override fun createSlot() = StateFlowSlot()
     override fun createSlotArray(size: Int): Array<StateFlowSlot?> = arrayOfNulls(size)
 

kotlinx-coroutines-core/common/src/flow/internal/FlowValueWrapperInternal.kt

@@ -1,5 +1,9 @@
+@file:Suppress("UNCHECKED_CAST")
+
 package kotlinx.coroutines.flow.internal
 
+import kotlinx.coroutines.flow.FlowCollector
+
 /**
  * Used by IDEA debugger agent to support asynchronous stack traces in flows.
  * The agent requires a unique object present in both current and async stack traces,
@@ -11,12 +15,40 @@ internal class FlowValueWrapperInternal<T>(val value: T)
 internal fun <T> wrapInternal(value: T): T = value
 internal fun <T> unwrapInternal(value: T): T = value
 
-// debugger agent transforms wrapInternal so it returns wrapInternalDebuggerCapture(value) instead of just value
-private fun wrapInternalDebuggerCapture(value: Any?): Any = FlowValueWrapperInternal(value)
+// debugger agent transforms wrapInternal so it returns wrapInternalDebuggerCaptureX(value) instead of just value.
+// "X" may be Strict or Lenient, debugger agent picks one depending on its arguments.
+// Both versions are aimed at handling double wrapping, which is always an error;
+// a lenient version swallows it, allowing the application to proceed normally
+// at the cost of not working async stack traces in the place;
+// a strict version throws an exception so that the issue could be traced at its earliest point
+
+private fun wrapInternalDebuggerCaptureStrict(value: Any?): Any {
+    if (value is FlowValueWrapperInternal<*>) {
+        throw DoubleWrappingException("Double-wrapping detected; failing fast. This should never happen!")
+    }
+    return FlowValueWrapperInternal(value)
+}
+
+private fun wrapInternalDebuggerCaptureLenient(value: Any?): Any {
+    if (value is FlowValueWrapperInternal<*>) {
+        return value
+    }
+    return FlowValueWrapperInternal(value)
+}
 
 // debugger agent transforms unwrapInternal so it returns unwrapInternalDebuggerCapture(value) instead of just value
 //
 // normally, value is always FlowValueWrapperInternal, but potentially instrumentation may start
 // in the middle of the execution (for example, when the debugger was attached to a running application),
 // and the emitted value hadn't been wrapped
-private fun unwrapInternalDebuggerCapture(value: Any?): Any? = (value as? FlowValueWrapperInternal<*>)?.value ?: value
+private fun unwrapInternalDebuggerCapture(value: Any?): Any? {
+    val wrapper = value as? FlowValueWrapperInternal<*> ?: return value
+    return wrapper.value
+}
+
+// Shouldn't be inlined, the method is instrumented by the IDEA debugger agent
+internal suspend fun <T> FlowCollector<T>.emitInternal(value: Any?) {
+    emit(NULL.unbox(unwrapInternal(value)))
+}
+
+private class DoubleWrappingException(message: String) : RuntimeException(message)

kotlinx-coroutines-core/common/src/flow/operators/Delay.kt

@@ -209,10 +209,10 @@ private fun <T> Flow<T>.debounceInternal(timeoutMillisSelector: (T) -> Long): Fl
             var timeoutMillis = 0L // will be always computed when lastValue != null
             // Compute timeout for this value
             if (lastValue != null) {
-                timeoutMillis = timeoutMillisSelector(NULL.unbox(lastValue))
+                timeoutMillis = timeoutMillisSelector(NULL.unbox(unwrapInternal(lastValue)))
                 require(timeoutMillis >= 0L) { "Debounce timeout should not be negative" }
                 if (timeoutMillis == 0L) {
-                    emitInner(downstream, lastValue)
+                    downstream.emitInternal(lastValue)
                     lastValue = null // Consume the value
                 }
             }
@@ -223,7 +223,7 @@ private fun <T> Flow<T>.debounceInternal(timeoutMillisSelector: (T) -> Long): Fl
                 // Set timeout when lastValue exists and is not consumed yet
                 if (lastValue != null) {
                     onTimeout(timeoutMillis) {
-                        emitInner<T>(downstream, lastValue)
+                        downstream.emitInternal<T>(lastValue)
                         lastValue = null // Consume the value
                     }
                 }
@@ -233,19 +233,14 @@ private fun <T> Flow<T>.debounceInternal(timeoutMillisSelector: (T) -> Long): Fl
                         .onFailure {
                             it?.let { throw it }
                             // If closed normally, emit the latest value
-                            if (lastValue != null) emitInner<T>(downstream, lastValue)
+                            if (lastValue != null) downstream.emitInternal<T>(lastValue)
                             lastValue = DONE
                         }
                 }
             }
         }
     }
 
-// Shouldn't be inlined, the method is instrumented by the IDEA debugger agent
-private suspend fun <T> emitInner(downstream: FlowCollector<T>, value: Any?) {
-    downstream.emit(NULL.unbox(unwrapInternal(value)))
-}
-
 /**
  * Returns a flow that emits only the latest value emitted by the original flow during the given sampling [period][periodMillis].
  *
@@ -295,7 +290,7 @@ public fun <T> Flow<T>.sample(periodMillis: Long): Flow<T> {
                 ticker.onReceive {
                     val value = lastValue ?: return@onReceive
                     lastValue = null // Consume the value
-                    emitInner(downstream, value)
+                    downstream.emitInternal(value)
                 }
             }
         }