Merge pull request #5 from AlexVanGogen/release/1.8.0-intellij-11

Enhance support for async stack traces in flows

Author: vsalavatov

IntelliJ-patches.md

@@ -61,15 +61,32 @@ 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`, `timeout`, `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`.
 
 ### 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.wrapInternalDebuggerCapture` -- wraps passed arguments into a `FlowValueWrapperInternal`; only used after transformation.
 - `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.unwrapTyped` -- utility function served to ease casting to a real underlying type
+- `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.emitInternal(FlowCollector, value)` -- alternative of a regular `FlowCollector.emit` that supports insertion points; if there is a `FlowCollector`, its `emit` call can be replaced with `emitInternal` so this case would also be supported for constructing async stack traces 
+- `kotlinx.coroutines.flow.internal.FlowValueWrapperInternalKt.debuggerCapture` -- common insertion point for a debugger agent; simplifies instrumentation; the value is always being unwrapped inside.
+
+One internal method was added to `BufferedChannel`: `emitAllInternal`. This method ensures the value will be unwrapped in an insertion point.
+
+One internal method was added to `flow/Channels.kt`: `emitAllInternal`. It emits all values, like usual, but also considers wrapping/unwrapping supported in `BufferedChannel`.
+
+One internal method was added to `ChannelCoroutine`: `emitAllInternal` serves to bridge its delegate and the method above.
+
+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.
+
+Changes were made to lambda parameter `onElementRetrieved` in `BufferedChannel<E>` methods: now they accept `Any?` instead of `E` because now they may be given a wrapped value.
+
+`SelectImplementation.complete` now uses `debuggerCapture` to properly propagate value that might come from flows. 

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.*
@@ -664,7 +665,7 @@ internal open class BufferedChannel<E>(
     protected open fun onReceiveDequeued() {}
 
     override suspend fun receive(): E =
-        receiveImpl( // <-- this is an inline function
+        receiveImpl<E>( // <-- this is an inline function
             // Do not create a continuation until it is required;
             // it is created later via [onNoWaiterSuspend], if needed.
             waiter = null,
@@ -673,7 +674,7 @@ internal open class BufferedChannel<E>(
             // Also, inform `BufferedChannel` extensions that
             // synchronization of this receive operation is completed.
             onElementRetrieved = { element ->
-                return element
+                return unwrapTyped<E>(element)
             },
             // As no waiter is provided, suspension is impossible.
             onSuspend = { _, _, _ -> error("unexpected") },
@@ -705,8 +706,9 @@ internal open class BufferedChannel<E>(
             // not dispatched yet. In case `onUndeliveredElement` is
             // specified, we need to invoke it in the latter case.
             onElementRetrieved = { element ->
-                val onCancellation = onUndeliveredElement?.bindCancellationFun(element, cont.context)
-                cont.resume(element, onCancellation)
+                val unwrapped: E = unwrapTyped(element)
+                val onCancellation = onUndeliveredElement?.bindCancellationFun(unwrapped, cont.context)
+                cont.resume(unwrapped, onCancellation)
             },
             onClosed = { onClosedReceiveOnNoWaiterSuspend(cont) },
         )
@@ -732,7 +734,7 @@ internal open class BufferedChannel<E>(
         receiveImpl( // <-- this is an inline function
             waiter = null,
             onElementRetrieved = { element ->
-                success(element)
+                success(unwrapTyped(element))
             },
             onSuspend = { _, _, _ -> error("unexpected") },
             onClosed = { closed(closeCause) },
@@ -749,7 +751,8 @@ internal open class BufferedChannel<E>(
             segment, index, r,
             waiter = waiter,
             onElementRetrieved = { element ->
-                cont.resume(success(element), onUndeliveredElement?.bindCancellationFun(element, cont.context))
+                val unwrapped = unwrapTyped<E>(element)
+                cont.resume(success(unwrapped), onUndeliveredElement?.bindCancellationFun(unwrapped, cont.context))
             },
             onClosed = { onClosedReceiveCatchingOnNoWaiterSuspend(cont) }
         )
@@ -782,7 +785,7 @@ internal open class BufferedChannel<E>(
             // Store an already interrupted receiver in case of suspension.
             waiter = INTERRUPTED_RCV,
             // Finish when an element is successfully retrieved.
-            onElementRetrieved = { element -> success(element) },
+            onElementRetrieved = { element -> success(unwrapTyped(element)) },
             // On suspension, the `INTERRUPTED_RCV` token has been
             // installed, and this `tryReceive()` must fail.
             onSuspend = { segm, _, globalIndex ->
@@ -846,7 +849,7 @@ internal open class BufferedChannel<E>(
                     // Clean the reference to the previous segment.
                     segment.cleanPrev()
                     @Suppress("UNCHECKED_CAST")
-                    onUndeliveredElement?.callUndeliveredElementCatchingException(updCellResult as E)?.let { throw it }
+                    onUndeliveredElement?.callUndeliveredElementCatchingException(unwrapTyped(updCellResult))?.let { throw it }
                 }
             }
         }
@@ -864,7 +867,7 @@ internal open class BufferedChannel<E>(
         /* This lambda is invoked when an element has been
         successfully retrieved, either from the buffer or
         by making a rendezvous with a suspended sender. */
-        onElementRetrieved: (element: E) -> R,
+        onElementRetrieved: (element: Any?) -> R,
         /* This lambda is called when the operation suspends in the cell
         specified by the segment and its global and in-segment indices. */
         onSuspend: (segm: ChannelSegment<E>, i: Int, r: Long) -> R,
@@ -934,7 +937,7 @@ internal open class BufferedChannel<E>(
                     // Clean the reference to the previous segment before finishing.
                     segment.cleanPrev()
                     @Suppress("UNCHECKED_CAST")
-                    onElementRetrieved(updCellResult as E)
+                    onElementRetrieved(updCellResult)
                 }
             }
         }
@@ -952,7 +955,7 @@ internal open class BufferedChannel<E>(
         /* This lambda is invoked when an element has been
         successfully retrieved, either from the buffer or
         by making a rendezvous with a suspended sender. */
-        onElementRetrieved: (element: E) -> Unit,
+        onElementRetrieved: (element: Any?) -> Unit,
         /* This lambda is called when the channel is observed
         in the closed state and no waiting senders is found,
         which means that it is closed for receiving. */
@@ -1540,7 +1543,7 @@ internal open class BufferedChannel<E>(
     @Suppress("UNCHECKED_CAST")
     private val onUndeliveredElementReceiveCancellationConstructor: OnCancellationConstructor? = onUndeliveredElement?.let {
         { select: SelectInstance<*>, _: Any?, element: Any? ->
-            { if (element !== CHANNEL_CLOSED) onUndeliveredElement.callUndeliveredElement(element as E, select.context) }
+            { if (element !== CHANNEL_CLOSED) onUndeliveredElement.callUndeliveredElement(unwrapTyped(element), select.context) }
         }
     }
 
@@ -1550,6 +1553,13 @@ internal open class BufferedChannel<E>(
 
     override fun iterator(): ChannelIterator<E> = BufferedChannelIterator()
 
+    internal suspend fun emitAllInternal(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]
@@ -1641,7 +1651,7 @@ internal open class BufferedChannel<E>(
                 onElementRetrieved = { element ->
                     this.receiveResult = element
                     this.continuation = null
-                    cont.resume(true, onUndeliveredElement?.bindCancellationFun(element, cont.context))
+                    cont.resume(true, onUndeliveredElement?.bindCancellationFun(unwrapTyped(element), cont.context))
                 },
                 onClosed = { onClosedHasNextNoWaiterSuspend() }
             )
@@ -1669,8 +1679,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,7 +1706,7 @@ internal open class BufferedChannel<E>(
             val cont = this.continuation!!
             this.continuation = null
             // Store the retrieved element in `receiveResult`.
-            this.receiveResult = element
+            this.receiveResult = wrapInternal(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.
@@ -2762,16 +2781,17 @@ internal class ChannelSegment<E>(id: Long, prev: ChannelSegment<E>?, channel: Bu
     }
 
     @Suppress("UNCHECKED_CAST")
-    internal fun getElement(index: Int) = data[index * 2].value as E
+    internal fun getElement(index: Int) = unwrapInternal(data[index * 2].value) as E
 
-    internal fun retrieveElement(index: Int): E = getElement(index).also { cleanElement(index) }
+    @Suppress("UNCHECKED_CAST")
+    internal fun retrieveElement(index: Int): E = (data[index * 2].value as E).also { cleanElement(index) }
 
     internal fun cleanElement(index: Int) {
-        setElementLazy(index, null)
+        data[index * 2].lazySet(null)
     }
 
     private fun setElementLazy(index: Int, value: Any?) {
-        data[index * 2].lazySet(value)
+        data[index * 2].lazySet(wrapInternal(value))
     }
 
     // ######################################

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.*

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

@@ -1,6 +1,8 @@
 package kotlinx.coroutines.channels
 
 import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.FlowCollector
+import kotlinx.coroutines.flow.emitAllInternal
 import kotlin.coroutines.*
 
 internal open class ChannelCoroutine<E>(
@@ -35,4 +37,8 @@ internal open class ChannelCoroutine<E>(
         _channel.cancel(exception) // cancel the channel
         cancelCoroutine(exception) // cancel the job
     }
+
+    internal suspend fun emitAllInternal(flowCollector: FlowCollector<E>) {
+        emitAllInternal(_channel, flowCollector)
+    }
 }

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)
-        }
+        emitAllInternal(channel, this)
     } catch (e: Throwable) {
         cause = e
         throw e
@@ -40,6 +38,22 @@ private suspend fun <T> FlowCollector<T>.emitAllImpl(channel: ReceiveChannel<T>,
     }
 }
 
+internal suspend fun <T> emitAllInternal(channel: ReceiveChannel<T>, collector: FlowCollector<T>) {
+    when (channel) {
+        is BufferedChannel<*> -> {
+            (channel as BufferedChannel<T>).emitAllInternal(collector)
+        }
+        is ChannelCoroutine<*> -> {
+            (channel as ChannelCoroutine<T>).emitAllInternal(collector)
+        }
+        else -> {
+            for (element in channel) {
+                collector.emit(element)
+            }
+        }
+    }
+}
+
 /**
  * Represents the given receive channel as a hot flow and [receives][ReceiveChannel.receive] from the channel
  * in fan-out fashion every time this flow is collected. One element will be emitted to one collector only.

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,
@@ -10,13 +14,32 @@ internal class FlowValueWrapperInternal<T>(val value: T)
 
 internal fun <T> wrapInternal(value: T): T = value
 internal fun <T> unwrapInternal(value: T): T = value
+internal fun <T> unwrapTyped(value: Any?): T = NULL.unbox(unwrapInternal(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 wrapInternalDebuggerCapture(value) instead of just value.
+private fun wrapInternalDebuggerCapture(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(unwrapTyped<T>(value))
+}
+
+// Shouldn't be inlined, the method is instrumented by the IDEA debugger agent
+internal suspend fun <Unwrapped, R> debuggerCapture(value: Any?, block: suspend (Unwrapped) -> R): R {
+    return block(unwrapInternal(value) as Unwrapped)
+}