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.

Author: AlexVanGogen

IntelliJ-patches.md

@@ -51,3 +51,25 @@ 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.
+
+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.
+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:
+
+- `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-core/common/src/flow/SharedFlow.kt

@@ -364,7 +364,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
         }
 
@@ -373,7 +373,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 {
@@ -389,13 +389,18 @@ internal open class SharedFlowImpl<T>(
                     awaitValue(slot) // await signal that the new value is available
                 }
                 collectorJob?.ensureActive()
-                collector.emit(newValue as T)
+                emitInner(collector, 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) {
@@ -469,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)
@@ -500,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

@@ -321,7 +321,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
@@ -357,6 +357,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)
 
@@ -389,7 +394,7 @@ private class StateFlowImpl<T>(
                 collectorJob?.ensureActive()
                 // Conflate value emissions using equality
                 if (oldState == null || oldState != newState) {
-                    collector.emit(NULL.unbox(newState))
+                    emitInner(collector, newState)
                     oldState = newState
                 }
                 // Note: if awaitPending is cancelled, then it bails out of this loop and calls freeSlot
@@ -402,6 +407,11 @@ 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

@@ -0,0 +1,22 @@
+package kotlinx.coroutines.flow.internal
+
+/**
+ * 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,
+ * so, without a wrapper, if two flows, `f1` and `f2`, emit equal values,
+ * the agent could suggest `f1` as emitter for `f2.collect` and `f2` as emitter for `f1.collect`.
+ */
+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 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

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

@@ -201,7 +201,7 @@ private fun <T> Flow<T>.debounceInternal(timeoutMillisSelector: (T) -> Long): Fl
     scopedFlow { downstream ->
         // Produce the values using the default (rendezvous) channel
         val values = produce {
-            collect { value -> send(value ?: NULL) }
+            collect { value -> send(value?.let(::wrapInternal) ?: NULL) }
         }
         // Now consume the values
         var lastValue: Any? = null
@@ -212,7 +212,7 @@ private fun <T> Flow<T>.debounceInternal(timeoutMillisSelector: (T) -> Long): Fl
                 timeoutMillis = timeoutMillisSelector(NULL.unbox(lastValue))
                 require(timeoutMillis >= 0L) { "Debounce timeout should not be negative" }
                 if (timeoutMillis == 0L) {
-                    downstream.emit(NULL.unbox(lastValue))
+                    emitInner(downstream, 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) {
-                        downstream.emit(NULL.unbox(lastValue))
+                        emitInner<T>(downstream, lastValue)
                         lastValue = null // Consume the value
                     }
                 }
@@ -233,14 +233,19 @@ 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) downstream.emit(NULL.unbox(lastValue))
+                            if (lastValue != null) emitInner<T>(downstream, 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].
  *
@@ -270,7 +275,7 @@ public fun <T> Flow<T>.sample(periodMillis: Long): Flow<T> {
     require(periodMillis > 0) { "Sample period should be positive" }
     return scopedFlow { downstream ->
         val values = produce(capacity = Channel.CONFLATED) {
-            collect { value -> send(value ?: NULL) }
+            collect { value -> send(value?.let(::wrapInternal) ?: NULL) }
         }
         var lastValue: Any? = null
         val ticker = fixedPeriodTicker(periodMillis)
@@ -290,7 +295,7 @@ public fun <T> Flow<T>.sample(periodMillis: Long): Flow<T> {
                 ticker.onReceive {
                     val value = lastValue ?: return@onReceive
                     lastValue = null // Consume the value
-                    downstream.emit(NULL.unbox(value))
+                    emitInner(downstream, value)
                 }
             }
         }