Launching Molecule with Dispatchers.Main.immediate breaks state invalidations in some Compose code

[benkay]

The issue is demonstrated by the below code, which reproduces the issue without actually depending on Molecule. The original
problem manifested as images loaded by Coil not fading in correctly - Coil's CrossfadePainter functions in much the same way as the Painter in this code.

class MainActivity : ComponentActivity() {
  override fun onCreate(savedInstanceState: Bundle?) {
    super.onCreate(savedInstanceState)
    enableEdgeToEdge()
    setContent {
      MaterialTheme {
        Scaffold(modifier = Modifier.fillMaxSize()) { innerPadding ->
          TestImage(modifier = Modifier.padding(innerPadding))
        }
      }
    }

    // Replacing the below block of code with this also causes the same issue
    // lifecycle.coroutineScope.launchMolecule(RecompositionMode.Immediate) {}

    lifecycle.coroutineScope.launch {
      var applyScheduled = false
      Snapshot.registerGlobalWriteObserver {
        if (!applyScheduled) {
          applyScheduled = true
          launch {
            applyScheduled = false
            Snapshot.sendApplyNotifications()
          }
        }
      }

      awaitCancellation()
    }
  }
}

@Composable
fun TestImage(modifier: Modifier = Modifier) {
  Image(
    painter = remember { InvalidatingPainter() },
    modifier = modifier,
    contentDescription = ""
  )
}

class InvalidatingPainter : Painter() {

  private var invalidateTick by mutableIntStateOf(0)
  private var lastStartTime: TimeSource.Monotonic.ValueTimeMark? = null
  private var duration = 1.seconds
  private var reverse = false

  override val intrinsicSize: Size = Size(200f, 200f)

  override fun DrawScope.onDraw() {
    val startTime = lastStartTime ?: TimeSource.Monotonic.markNow().also { lastStartTime = it }
    val percent = (startTime.elapsedNow() / duration).toFloat().coerceIn(0f, 1f)
    val alpha = if (reverse) 1 - percent else percent
    drawRect(Color.Blue, alpha = alpha)

    if (percent >= 1f) {
      lastStartTime = null
      reverse = !reverse
    }
    invalidateTick++
  }
}

The cause seems to be

• The default coroutine scopes exposed by AndroidX use Dispatchers.Main.immediate
• Because of this, the launch within the global write observer doesn't actually dispatch, and apply notifications are sent immediately instead of in the next event loop
• The invalidateTick state write is happening within the draw phase. This is where my Compose internals knowledge gets a bit fuzzy, but I'm guessing Snapshot.sendApplyNotifications() needs to happen after read SnapshotStateObserver.observeReads finishes for any updates to the read state to actually cause a change notification.

The answer here seems to be "never use the immediate dispatcher", but it's very easy to do so without knowing at the moment.
e.g. viewModelScope.launchMolecule(..) will use the immediate dispatcher unless otherwise specified.

Perhaps Molecule could detect this and do the right thing?

[JakeWharton]

So there's a few paths forward here, and I'll probably do all of them:

1. We should provide the option to disable starting our own snapshot application callback. If you are in an Android app with Compose UI there's no reason to run a second snapshot applier.
2. We should provide the ability to pass a parent Composition (which likely implies 1). If you are hosting Molecule beneath Compose UI to drive the UI, we can make our composition a child of it.
3. We should detect this case and force dispatch. I assume this implies a CoroutineDispatcher is available as a key in the context, but if not perhaps we just hard fail?

[chris6647]

I'm experiencing, what I suspect, is the same issue. Though my molecule stream runs my own CoroutineScope, and the LaunchedEffects creating state for the composeUi still uses currentComposer.applyCoroutineContext.

In my Android app I noticed some MutableState updates, triggered from a LaunchedEffect, would never trigger recompositions, causing the compose UI to be out of sync with the underlying state.

Took ages to narrow down, but I believe that I've been able to make a test for it. "Most" of the time it fails. I took a Circuit test and modified it to run a launchMolecule StateFlow. It seems that the combination of running molecule, and updating compose states in LaunchedEffects, causes some state updates to be lost.
Running the test is sometimes successfull, sometimes fails due to being unable to find the expected compose nodes and a few times fails with timeout, due to the compose never going idle.
I've commited the test on my fork of Circuit here, in case its helpful to anyone.

If my issue really is the same issue, then path 1 could work, as long as I can enable/disable the snapshot callbacks while molecule is running and has subscribers, so I can use the UI one when I have it, and not when in background.

If you guys know of a workaround for now though, that would be great - thanks in advance! :)