RecompositionMode.Immediate with non-ui dispatcher can cause "missed" recompositions

[j2esu]

I've noticed that if molecule composition happens together with ui composition the molecule composition is "missed". It never happens with the ContextClock mode but it's quite easy to reproduce with the Immediate mode and non-ui dispatcher.

Below there is an example of counter that updates the local state and sends a message to update the state produced by molecule. In 5-15 clicks (no need to click fast though) it's possible to get the wrong state when the states are not equal. The next click after that usually makes them equal again (which I think means that the state in molecule gets updated, but it doesn't trigger a recomposition, probably because the ui recomposition is currently in progress).

Molecule 1.4.2 Compose compiler 1.5.11 Compose runtime 1.7.0-alpha06 (to get #396 fixed) Android SDK 34

class MainActivity : ComponentActivity() {

    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        val messages = MutableSharedFlow<Unit>(extraBufferCapacity = 10)
        val state = lifecycleScope.launchMolecule(RecompositionMode.Immediate, Dispatchers.Default) {
            var count by remember { mutableIntStateOf(0) }
            LaunchedEffect(Unit) {
                messages.collect { count++ }
            }
            count
        }
        setContent {
            val count by state.collectAsState()
            var countLocal by remember { mutableIntStateOf(0) }
            Column(modifier = Modifier.background(Color.Black)) {
                Text(text = "Count: $count", color = Color.White)
                Text(text = "Count local: $countLocal", color = Color.White)
                Text(
                    text = "Click me",
                    color = Color.White,
                    modifier = Modifier
                        .clickable {
                            countLocal++
                            messages.tryEmit(Unit)
                        }
                )
            }
        }
    }
}

[j2esu]

It's much harder to reproduce (takes 50-100 clicks) with molecule 2.0.0, compose 1.7.0-beta01 and compose compiler shipped with kotlin 2.0.0

[j2esu]

The issue is till there on compose 1.7.0-beta06

[abusalimov]

I can confirm that this reproduces for me as well, including the trunk version (eaa88891ecbe6d04bb58910494ede796283f224d).

Here's another reproducer that can run on desktop as well. To run it, just put it into main.kt next to molecule.kt and use green icon in the gutter.

package app.cash.molecule

import androidx.compose.runtime.collectAsState
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableIntStateOf
import androidx.compose.runtime.setValue
import app.cash.molecule.RecompositionMode.Immediate
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.first

// https://github.com/cashapp/molecule/issues/416
// RecompositionMode.Immediate with non-ui dispatcher can cause "missed" recompositions
public suspend fun main(): Unit = withContext(Dispatchers.Default) {
  val dispatcher = Dispatchers.Default.limitedParallelism(2)  // doesn't reproduce with 1

  val moleculeCountState = mutableIntStateOf(0)
  var moleculeCount by moleculeCountState

  val moleculeStateFlow = launchMolecule(Immediate, dispatcher) {
    var count by moleculeCountState
    count
  }

  launch {
    while (true) {
      if (moleculeCount % 1000 == 0){
        println("[$moleculeCount]")
      }
      try {
        withTimeout(1000) {
          moleculeStateFlow.first { it == moleculeCount }
        }
      } catch (_: TimeoutCancellationException) {
        println("[$moleculeCount] timeout / must not happen!")
        moleculeCount++
        continue
      }
      moleculeCount++
    }
  }

  // couldn't reproduce the issue without having a separate recomposer running
  launchMolecule(Immediate, dispatcher) {
    moleculeStateFlow.collectAsState().value
  }
}

Normally, it should never output "timeout / must not happen", but here's what I have:

> Task :molecule-runtime:jvmRun
[0]
[11] timeout / must not happen!
[35] timeout / must not happen!
[111] timeout / must not happen!
[218] timeout / must not happen!
[910] timeout / must not happen!
[1000]
[1775] timeout / must not happen!
[1844] timeout / must not happen!
[1864] timeout / must not happen!

[abusalimov]

FWIW, it's most likely a Compose runtime issue, not Molecule's.

Here's what I narrowed the Molecule down to, and it still reproduces the issue (using .stateIn(...) to turn the moleculeFlow() into a hot flow):

package com.example.helloandroid.molecule

import androidx.compose.runtime.AbstractApplier
import androidx.compose.runtime.Composable
import androidx.compose.runtime.Composition
import androidx.compose.runtime.MonotonicFrameClock
import androidx.compose.runtime.Recomposer
import androidx.compose.runtime.SideEffect
import androidx.compose.runtime.snapshots.Snapshot
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.async
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.conflate
import kotlinx.coroutines.launch
import java.util.concurrent.atomic.AtomicLong

fun <T> moleculeFlow(body: @Composable () -> T): Flow<T> = channelFlow {
  val clock = TrivialFrameClock(this)
  val finalContext = coroutineContext + clock

  val snapshotHandle = Snapshot.registerGlobalWriteObserver {
    launch {
      Snapshot.sendApplyNotifications()
    }
  }

  val recomposer = Recomposer(finalContext)
  val composition = Composition(UnitApplier, recomposer).apply {
    setContent {
      val compositionResult = body()
      SideEffect {
        trySend(compositionResult).getOrThrow()
      }
    }
  }
  launch(finalContext) {
    recomposer.runRecomposeAndApplyChanges()
  }

  awaitClose {
    snapshotHandle.dispose()
    composition.dispose()
  }
}.conflate()


private class TrivialFrameClock(private val cs: CoroutineScope) : MonotonicFrameClock {
  private val nanos = AtomicLong(0)

  override suspend fun <R> withFrameNanos(
    onFrame: (Long) -> R
  ): R {
    return cs.async {
      val currentNanos = nanos.incrementAndGet()
      onFrame(currentNanos)
    }.await()
  }
}

private object UnitApplier : AbstractApplier<Unit>(Unit) {
  override fun insertBottomUp(index: Int, instance: Unit) {}
  override fun insertTopDown(index: Int, instance: Unit) {}
  override fun move(from: Int, to: Int, count: Int) {}
  override fun remove(index: Int, count: Int) {}
  override fun onClear() {}
}

[abusalimov]

upstream bug: b/419527812 fix with tests: 3638709: Recomposer: remove incorrect resetting of snapshotInvalidations

FYI @JakeWharton

[j2esu]

The fix has been released in androidx.compose.runtime:runtime:1.9.0-alpha04 Can't reproduce it anymore, thank you

[abusalimov]

The fix in Compose Runtime has also been backported to JetBrains Compose Multiplatform 1.8.2: https://youtrack.jetbrains.com/issue/CMP-8239/Cherry-pick-changes-in-Recomposer-to-fix-race-condition-causing-missing-snapshot-invalidations

note: this only refers to org.jetbrains.compose.runtime:runtime-desktop:1.8.2 and doesn't include the mainline Compose artifacts built by Google for Android. The latter will be available in 1.9.0 indeed.