[native-mt] problem with MutableStateFlow

[dss99911]

Hi,
I'm using kotlin multiplatform.
and using MutableStateFlow which is accessed by swift code

when I run the code below, error occurred

class TestStateFlow {
    val scope = CoroutineScope(Dispatchers.Main)


    val state = MutableStateFlow("init")

    fun start() {
        scope.launch {
            state.collect {
                println("state $it")
            }
        }

        scope.launch(start = CoroutineStart.LAZY) {
            HttpClient().get<String>("https://test.com").also {
                println("result : $it")
            }
            stop()
        }.start()
    }

    fun stop() {
        scope.coroutineContext.cancel()
    }
}

Swift code

TestStateFlow().start()

The below is error log

Caused by: kotlin.native.concurrent.InvalidMutabilityException: mutation attempt of frozen kotlinx.coroutines.flow.StateFlowImpl@17d1ac8
    at 0   kotlinIOS                           0x000000010ac5e36d kfun:kotlin.Throwable#<init>(kotlin.String?){} + 93
    at 1   kotlinIOS                           0x000000010ac56f1b kfun:kotlin.Exception#<init>(kotlin.String?){} + 91
    at 2   kotlinIOS                           0x000000010ac5716b kfun:kotlin.RuntimeException#<init>(kotlin.String?){} + 91
    at 3   kotlinIOS                           0x000000010ac8d05b kfun:kotlin.native.concurrent.InvalidMutabilityException#<init>(kotlin.String){} + 91
    at 4   kotlinIOS                           0x000000010ac8e912 ThrowInvalidMutabilityException + 690
    at 5   kotlinIOS                           0x000000010ad892dc MutationCheck + 108
    at 6   kotlinIOS                           0x000000010ae2414d kfun:kotlinx.coroutines.flow.StateFlowImpl.<set-nSlots>#internal + 77
    at 7   kotlinIOS                           0x000000010ae26486 kfun:kotlinx.coroutines.flow.StateFlowImpl.freeSlot#internal + 294
    at 8   kotlinIOS                           0x000000010ae25aa4 kfun:kotlinx.coroutines.flow.StateFlowImpl.$collectCOROUTINE$127.invokeSuspend#internal + 3188
    at 9   kotlinIOS                           0x000000010ac7fe96 kfun:kotlin.coroutines.native.internal.BaseContinuationImpl#resumeWith(kotlin.Result<kotlin.Any?>){} + 758
    at 10  kotlinIOS                           0x000000010ae35aa6 kfun:kotlinx.coroutines.DispatchedTask#run(){} + 2470
    at 11  kotlinIOS                           0x000000010ae626e5 kfun:kotlinx.coroutines.DarwinMainDispatcher.dispatch$lambda-0#internal + 85
    ... and 14 more common stack frames skipped
(lldb) 

When I analyze the error.
the reason is that

• StateFlow.nSlot is frozen when background thread is started.
• when StateFlow.collect is cancelled, nSlot is needed to change.
• but, nSlot is frozen. crash occurred

So, I tested after changing nSlot's type to atomic, and It worked fine for me
Kindly review if this is bug or intended.
Thanks to read.

The code to test with atomic ( I just copy and pasted MutableStateFlow code and changed some internal code)

class TestStateFlow {
    val scope = CoroutineScope(Dispatchers.Main)


    val state = StateFlowImpl<String>("init")

    fun start() {
        println("start")
        scope.launch {
            state.collect {
                println("state $it")
            }
        }

        scope.launch(start = CoroutineStart.LAZY) {

            HttpClient().get<String>("https://sample.jeonghyeon.kim/ping").also {
                println("result : $it")
            }
            stop()
        }.start()
    }

    fun stop() {
        scope.coroutineContext.cancel()
    }
}

class StateFlowImpl<T>(initialValue: Any) : SynchronizedObject(), MutableStateFlow<T> {
    private val _state = kotlinx.atomicfu.atomic(initialValue) // T | NULL
    private var sequence = 0 // serializes updates, value update is in process when sequence is odd
    private var slots = arrayOfNulls<StateFlowSlot?>(2)
    private val nSlots = kotlinx.atomicfu.atomic(0) // T | NULL // number of allocated (!free) slots
    private var nextIndex = 0 // oracle for the next free slot index

    @OptIn(ExperimentalCoroutinesApi::class)
    @Suppress("UNCHECKED_CAST")
    public override var value: T
        get() = NULL2.unbox(_state.value)
        set(value) {
            var curSequence = 0
            var curSlots: Array<StateFlowSlot?> = this.slots // benign race, we will not use it
            val newState = value ?: NULL2
            kotlinx.atomicfu.locks.synchronized(this) {
                val oldState = _state.value
                if (oldState == newState) return // Don't do anything if value is not changing
                _state.value = newState
                curSequence = sequence
                if (curSequence and 1 == 0) { // even sequence means quiescent state flow (no ongoing update)
                    curSequence++ // make it odd
                    sequence = curSequence
                } else {
                    // update is already in process, notify it, and return
                    sequence = curSequence + 2 // change sequence to notify, keep it odd
                    return
                }
                curSlots = slots // read current reference to collectors under lock
            }
            /*
               Fire value updates outside of the lock to avoid deadlocks with unconfined coroutines
               Loop until we're done firing all the changes. This is sort of simple flat combining that
               ensures sequential firing of concurrent updates and avoids the storm of collector resumes
               when updates happen concurrently from many threads.
             */
            while (true) {
                // Benign race on element read from array
                for (col in curSlots) {
                    col?.makePending()
                }
                // check if the value was updated again while we were updating the old one
                kotlinx.atomicfu.locks.synchronized(this) {
                    if (sequence == curSequence) { // nothing changed, we are done
                        sequence = curSequence + 1 // make sequence even again
                        return // done
                    }
                    // reread everything for the next loop under the lock
                    curSequence = sequence
                    curSlots = slots
                }
            }
        }

    @OptIn(InternalCoroutinesApi::class)
    override suspend fun collect(collector: FlowCollector<T>) {
        val slot = allocateSlot()
        var prevState: Any? = null // previously emitted T!! | NULL (null -- nothing emitted yet)
        try {
            // The loop is arranged so that it starts delivering current value without waiting first
            while (true) {
                // Here the coroutine could have waited for a while to be dispatched,
                // so we use the most recent state here to ensure the best possible conflation of stale values
                val newState = _state.value
                // Conflate value emissions using equality
                if (prevState == null || newState != prevState) {
                    collector.emit(NULL2.unbox(newState))
                    prevState = newState
                }
                // Note: if awaitPending is cancelled, then it bails out of this loop and calls freeSlot
                if (!slot.takePending()) { // try fast-path without suspending first
                    slot.awaitPending() // only suspend for new values when needed
                }
            }
        } finally {
            freeSlot(slot)
        }
    }

    private fun allocateSlot(): StateFlowSlot = kotlinx.atomicfu.locks.synchronized(this) {
        val size = slots.size
        if (nSlots.value >= size) slots = slots.copyOf(2 * size)
        var index = nextIndex
        var slot: StateFlowSlot
        while (true) {
            slot = slots[index] ?: StateFlowSlot().also { slots[index] = it }
            index++
            if (index >= slots.size) index = 0
            if (slot.allocate()) break // break when found and allocated free slot
        }
        nextIndex = index
        nSlots.incrementAndGet()
        slot
    }

    private fun freeSlot(slot: StateFlowSlot): Unit = kotlinx.atomicfu.locks.synchronized(this) {
        slot.free()
        nSlots.decrementAndGet()
    }
}

class StateFlowSlot {
    /**
     * Each slot can have one of the following states:
     *
     * * `null` -- it is not used right now. Can [allocate] to new collector.
     * * `NONE` -- used by a collector, but neither suspended nor has pending value.
     * * `PENDING` -- pending to process new value.
     * * `CancellableContinuationImpl<Unit>` -- suspended waiting for new value.
     *
     * It is important that default `null` value is used, because there can be a race between allocation
     * of a new slot and trying to do [makePending] on this slot.
     */
    private val _state = kotlinx.atomicfu.atomic<Any?>(null)

    fun allocate(): Boolean {
        // No need for atomic check & update here, since allocated happens under StateFlow lock
        if (_state.value != null) return false // not free
        _state.value = NONE // allocated
        return true
    }

    fun free() {
        _state.value = null // free now
    }

    @Suppress("UNCHECKED_CAST")
    fun makePending() {
        _state.loop { state ->
            when {
                state == null -> return // this slot is free - skip it
                state === PENDING -> return // already pending, nothing to do
                state === NONE -> { // mark as pending
                    if (_state.compareAndSet(state, PENDING)) return
                }
                else -> { // must be a suspend continuation state
                    // we must still use CAS here since continuation may get cancelled and free the slot at any time
                    if (_state.compareAndSet(state, NONE)) {
                        (state as CancellableContinuation<Unit>).resume(Unit)
                        return
                    }
                }
            }
        }
    }

    fun takePending(): Boolean = _state.getAndSet(NONE)!!.let { state ->
        check(state !is CancellableContinuation<*>)
        return state === PENDING
    }

    @Suppress("UNCHECKED_CAST")
    suspend fun awaitPending(): Unit = suspendCancellableCoroutine sc@ { cont ->
        check(_state.value !is CancellableContinuation<*>) // can be NONE or PENDING
        if (_state.compareAndSet(NONE, cont)) return@sc // installed continuation, waiting for pending
        // CAS failed -- the only possible reason is that it is already in pending state now
        check(_state.value === PENDING)
        cont.resume(Unit)
    }
}

val NONE = Symbol("NONE")
val PENDING = Symbol("PENDING")
val NULL2 = Symbol("NULL")//error when using name NULL on swift
class Symbol(val symbol: String) {
    override fun toString(): String = symbol

    @Suppress("UNCHECKED_CAST", "NOTHING_TO_INLINE")
    inline fun <T> unbox(value: Any?): T = if (value === this) null as T else value as T
}

[dss99911]

I mentioned only nSlot but I think nextIndex, slots, sequence also need to be changed to atomic

[Thomas-Vos]

#2138

[qwwdfsad]

Fixed in 1.3.9-native-mt-2