IDEA-352355 Asynchronous stack traces for flows in the IDEA debugger

Patched during cherry-pick to 1.10.1. File with conflicts: BufferedChannel.kt

Squashed commits as of version 1.8.0-intellij-11

Don't make unnecessary NULL unboxings

(cherry picked from commit 0267812)

Don't wrap suspend function

(cherry picked from commit 8684cad)

Get rid of changes in the public API -- bring back CoroutineChannel support

(cherry picked from commit 9836a6b)

Get rid of changes in the public API

(cherry picked from commit f852554)

Support SelectImplementation

This brings support for select-based flow operators, such as `timeout`.

(cherry picked from commit db906d8)

Support BufferedChannel

Before, `buffer` operations and such were only supported partially, and in some cases async stack traces were working only in 50% collects. For example, when several flows are merged with `flattenMerge` and emit values simultaneously.

This change seems large, changing a lot of lines in BufferedChannel.kt, but most of them are effectively refactoring (propagation of a wrapped value).

(cherry picked from commit 8be4def)

Discard strict double-wrapping check

We decided not to go with it, as it may dump a lot of error messages to a clueless user's console.

(cherry picked from commit 0efa558)

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.

(cherry picked from commit 00cb4e5)

Prepare shared flows for the debugger agent to support async stack traces

The agent needs three entities to establish a proper asynchronous stack traces connection:
  - a capture point -- method that indicates the stack trace that precedes the current stack trace;
  - an insertion point -- method within the current stack trace;
  - a key -- an object that is present in both points and is unique enough to bridge two points properly.

This change tweaks the code a bit to introduce the three entities in MutableSharedFlow and MutableStateFlow.

The key for MutableSharedFlow 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.

(cherry picked from commit 75107bc)

# Conflicts:
#	kotlinx-coroutines-core/common/src/channels/BufferedChannel.kt

# Conflicts:
#	IntelliJ-patches.md

Author: AlexVanGogen

IntelliJ-patches.md

@@ -50,3 +50,42 @@ exceeding the parallelism limit would eliminate this (likely expected) side effe
   associated coroutine dispatcher when it decides to park the thread
 - `CoroutineDispatcher.softLimitedParallelism` – an analogue of `.limitedParallelism` which supports
   parallelism compensation
+
+## Asynchronous stack traces for flows in the IDEA debugger
+
+The agent needs three entities to establish a proper asynchronous stack traces connection:
+- a capture point — method that indicates the stack trace that precedes the current stack trace;
+- an insertion point — method within the current stack trace;
+- a key — an object that is present in both points and is unique enough to bridge two stack traces properly.
+
+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.
+
+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
+
+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.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,8 @@ 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.*
 import kotlinx.coroutines.selects.TrySelectDetailedResult.*
@@ -684,7 +686,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,
@@ -693,7 +695,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") },
@@ -726,7 +728,7 @@ internal open class BufferedChannel<E>(
             // specified, we need to invoke it in the latter case.
             onElementRetrieved = { element ->
                 val onCancellation = onUndeliveredElement?.bindCancellationFun()
-                cont.resume(element, onCancellation)
+                cont.resume(unwrapTyped(element), onCancellation)
             },
             onClosed = { onClosedReceiveOnNoWaiterSuspend(cont) },
         )
@@ -752,7 +754,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) },
@@ -769,7 +771,8 @@ internal open class BufferedChannel<E>(
             segment, index, r,
             waiter = waiter,
             onElementRetrieved = { element ->
-                cont.resume(success(element), onUndeliveredElement?.bindCancellationFunResult())
+                val unwrapped = unwrapTyped<E>(element)
+                cont.resume(success(unwrapped), onUndeliveredElement?.bindCancellationFunResult())
             },
             onClosed = { onClosedReceiveCatchingOnNoWaiterSuspend(cont) }
         )
@@ -802,7 +805,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 ->
@@ -866,7 +869,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 }
                 }
             }
         }
@@ -884,7 +887,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,
@@ -954,7 +957,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)
                 }
             }
         }
@@ -972,7 +975,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. */
@@ -1561,7 +1564,7 @@ internal open class BufferedChannel<E>(
     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)
             }
         }
     }
@@ -1572,6 +1575,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]
@@ -1666,7 +1676,7 @@ internal open class BufferedChannel<E>(
                 onElementRetrieved = { element ->
                     this.receiveResult = element
                     this.continuation = null
-                    cont.resume(true, onUndeliveredElement?.bindCancellationFun(element))
+                    cont.resume(true, onUndeliveredElement?.bindCancellationFun(unwrapTyped(element)))
                 },
                 onClosed = { onClosedHasNextNoWaiterSuspend() }
             )
@@ -1694,8 +1704,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" }
@@ -1712,7 +1731,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.
@@ -2815,16 +2834,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

@@ -369,7 +369,7 @@ internal open class SharedFlowImpl<T>(
             val result = ArrayList<T>(replaySize)
             val buffer = buffer!! // must be allocated, because replaySize > 0
             @Suppress("UNCHECKED_CAST")
-            for (i in 0 until replaySize) result += buffer.getBufferAt(replayIndex + i) as T
+            for (i in 0 until replaySize) result += unwrapInternal(buffer.getBufferAt(replayIndex + i)) as T
             result
         }
 
@@ -378,7 +378,7 @@ internal open class SharedFlowImpl<T>(
      */
     @Suppress("UNCHECKED_CAST")
     protected val lastReplayedLocked: T
-        get() = buffer!!.getBufferAt(replayIndex + replaySize - 1) as T
+        get() = unwrapInternal(buffer!!.getBufferAt(replayIndex + replaySize - 1)) as T
 
     @Suppress("UNCHECKED_CAST")
     override suspend fun collect(collector: FlowCollector<T>): Nothing {
@@ -394,7 +394,7 @@ internal open class SharedFlowImpl<T>(
                     awaitValue(slot) // await signal that the new value is available
                 }
                 collectorJob?.ensureActive()
-                collector.emit(newValue as T)
+                collector.emitInternal(newValue as T)
             }
         } finally {
             freeSlot(slot)
@@ -474,8 +474,16 @@ internal open class SharedFlowImpl<T>(
         minCollectorIndex = newHead
     }
 
-    // enqueues item to buffer array, caller shall increment either bufferSize or queueSize
     private fun enqueueLocked(item: Any?) {
+        enqueueLockedInner(wrapInternal(item))
+    }
+
+    private fun enqueueEmitterLocked(emitter: Emitter) {
+        enqueueLockedInner(emitter)
+    }
+
+    // enqueues item to buffer array, caller shall increment either bufferSize or queueSize
+    private fun enqueueLockedInner(item: Any?) {
         val curSize = totalSize
         val buffer = when (val curBuffer = buffer) {
             null -> growBuffer(null, 0, 2)
@@ -505,8 +513,8 @@ internal open class SharedFlowImpl<T>(
                 return@lock null
             }
             // add suspended emitter to the buffer
-            Emitter(this, head + totalSize, value, cont).also {
-                enqueueLocked(it)
+            Emitter(this, head + totalSize, wrapInternal(value), cont).also {
+                enqueueEmitterLocked(it)
                 queueSize++ // added to queue of waiting emitters
                 // synchronous shared flow might rendezvous with waiting emitter
                 if (bufferCapacity == 0) resumes = findSlotsToResumeLocked(resumes)

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

@@ -330,7 +330,7 @@ private class StateFlowImpl<T>(
             val oldState = _state.value
             if (expectedState != null && oldState != expectedState) return false // CAS support
             if (oldState == newState) return true // Don't do anything if value is not changing, but CAS -> true
-            _state.value = newState
+            updateInner(newState)
             curSequence = sequence
             if (curSequence and 1 == 0) { // even sequence means quiescent state flow (no ongoing update)
                 curSequence++ // make it odd
@@ -366,6 +366,11 @@ private class StateFlowImpl<T>(
         }
     }
 
+    // Shouldn't be inlined, the method is instrumented by the IDEA debugger agent
+    private fun updateInner(newState: Any) {
+        _state.value = newState
+    }
+
     override val replayCache: List<T>
         get() = listOf(value)
 
@@ -398,7 +403,7 @@ private class StateFlowImpl<T>(
                 collectorJob?.ensureActive()
                 // Conflate value emissions using equality
                 if (oldState == null || oldState != newState) {
-                    collector.emit(NULL.unbox(newState))
+                    collector.emitInternal(newState)
                     oldState = newState
                 }
                 // Note: if awaitPending is cancelled, then it bails out of this loop and calls freeSlot