From 8c1c70412ffadb49dfe8753e2ee397899b83f163 Mon Sep 17 00:00:00 2001 From: djain2405 Date: Fri, 16 Oct 2026 20:02:46 -0700 Subject: [PATCH] Stop restarting a long-lived Source of Truth reader on every write (#572) Signed-off-by: Divya Jain --- .../store5/impl/SourceOfTruthWithBarrier.kt | 205 ++++++++++++------ .../store/store5/FlowStoreTests.kt | 66 ++++++ .../store5/SourceOfTruthWithBarrierTests.kt | 145 +++++++++++++ 3 files changed, 354 insertions(+), 62 deletions(-) diff --git a/store/src/commonMain/kotlin/org/mobilenativefoundation/store/store5/impl/SourceOfTruthWithBarrier.kt b/store/src/commonMain/kotlin/org/mobilenativefoundation/store/store5/impl/SourceOfTruthWithBarrier.kt index f9926a25f..6ae88f176 100644 --- a/store/src/commonMain/kotlin/org/mobilenativefoundation/store/store5/impl/SourceOfTruthWithBarrier.kt +++ b/store/src/commonMain/kotlin/org/mobilenativefoundation/store/store5/impl/SourceOfTruthWithBarrier.kt @@ -18,21 +18,19 @@ package org.mobilenativefoundation.store.store5.impl import kotlinx.atomicfu.atomic import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Job import kotlinx.coroutines.NonCancellable +import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.coroutineScope import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.MutableStateFlow -import kotlinx.coroutines.flow.catch -import kotlinx.coroutines.flow.emitAll -import kotlinx.coroutines.flow.flatMapLatest import kotlinx.coroutines.flow.flow -import kotlinx.coroutines.flow.flowOf -import kotlinx.coroutines.flow.onStart +import kotlinx.coroutines.launch import kotlinx.coroutines.withContext import org.mobilenativefoundation.store.store5.Converter import org.mobilenativefoundation.store.store5.SourceOfTruth import org.mobilenativefoundation.store.store5.StoreReadResponse import org.mobilenativefoundation.store.store5.StoreReadResponseOrigin -import org.mobilenativefoundation.store.store5.impl.operators.mapIndexed /** * Wraps a [SourceOfTruth] and blocks reads while a write is in progress. @@ -71,81 +69,147 @@ internal class SourceOfTruthWithBarrier - val messageArrivedAfterMe = readerVersion < barrierMessage.version - val writeError = - if (messageArrivedAfterMe && barrierMessage is BarrierMsg.Open) { - barrierMessage.writeError - } else { - null + // Keep one collection of a live reader across writes. flatMapLatest used to cancel + // it on every barrier transition, which restarted Room and channelFlow readers + // (#572). A one-shot reader completes after its value, so it is collected again + // only after that. The writer can emit the new row before this coroutine observes + // the barrier reopen; that value is the post-write read. + coroutineScope { + val signals = Channel>(Channel.UNLIMITED) + val barrierJob = + launch { + barrier.collect { message -> + signals.send(ReaderSignal.Barrier(message)) + } + } + var readerJob: Job? = null + + fun startReader() { + readerJob = + launch { + try { + delegate.reader(key).collect { value -> + signals.send(ReaderSignal.Value(value, barrier.value)) + } + signals.send(ReaderSignal.Completed) + } catch (cancellation: CancellationException) { + throw cancellation + } catch (throwable: Throwable) { + signals.send(ReaderSignal.Failed(throwable.cause ?: throwable)) } - val readFlow: Flow> = - when (barrierMessage) { - is BarrierMsg.Open -> - delegate.reader(key).mapIndexed { index, local: Output? -> - if (index == 0 && messageArrivedAfterMe) { - val firstMsgOrigin = - if (writeError == null) { - // restarted barrier without an error means write succeeded - StoreReadResponseOrigin.Fetcher() - } else { - // when a write fails, we still get a new reader because - // we've disabled the previous reader before starting the - // write operation. But since write has failed, we should - // use the SourceOfTruth as the origin - StoreReadResponseOrigin.SourceOfTruth - } - StoreReadResponse.Data( - origin = firstMsgOrigin, - value = local, + } + } + + var pendingValue: Output? = null + var hasPending = false + var announcedVersion = Long.MIN_VALUE + var emittedVersion = Long.MIN_VALUE + try { + for (signal in signals) { + when (signal) { + is ReaderSignal.Barrier -> { + val message = signal.message + if (message is BarrierMsg.Open) { + val afterMe = readerVersion < message.version + if (afterMe && announcedVersion < message.version) { + message.writeError?.let { writeError -> + emit( + StoreReadResponse.Error.Exception( + origin = StoreReadResponseOrigin.SourceOfTruth, + error = writeError, + ), ) - } else { - StoreReadResponse.Data( - origin = StoreReadResponseOrigin.SourceOfTruth, - value = local, - ) as StoreReadResponse } - }.catch { throwable -> - this.emit( - StoreReadResponse.Error.Exception( - error = - SourceOfTruth.ReadException( - key = key, - cause = throwable.cause ?: throwable, - ), - origin = StoreReadResponseOrigin.SourceOfTruth, - ), - ) + announcedVersion = message.version + if (hasPending && emittedVersion < message.version) { + emit( + StoreReadResponse.Data( + origin = postWriteOrigin(message), + value = pendingValue, + ), + ) + emittedVersion = message.version + hasPending = false + pendingValue = null + } + } + if (readerJob?.isActive != true) { + startReader() } - - is BarrierMsg.Blocked -> { - flowOf() } } - readFlow - .onStart { - // if we have a pending error, make sure to dispatch it first. - if (writeError != null) { + + is ReaderSignal.Value -> { + val snapshot = signal.snapshot + if (snapshot !is BarrierMsg.Open) { + pendingValue = signal.value + hasPending = true + } else if ( + readerVersion < snapshot.version && + emittedVersion < snapshot.version && + announcedVersion < snapshot.version + ) { + pendingValue = signal.value + hasPending = true + } else if ( + readerVersion < snapshot.version && + emittedVersion < snapshot.version + ) { + emit( + StoreReadResponse.Data( + origin = postWriteOrigin(snapshot), + value = signal.value, + ), + ) + emittedVersion = snapshot.version + hasPending = false + pendingValue = null + } else { emit( - StoreReadResponse.Error.Exception( + StoreReadResponse.Data( origin = StoreReadResponseOrigin.SourceOfTruth, - error = writeError, + value = signal.value, ), ) } } - }, - ) + + is ReaderSignal.Failed -> { + emit( + StoreReadResponse.Error.Exception( + origin = StoreReadResponseOrigin.SourceOfTruth, + error = + SourceOfTruth.ReadException( + key = key, + cause = signal.cause, + ), + ), + ) + } + + ReaderSignal.Completed -> Unit + } + } + } finally { + barrierJob.cancel() + readerJob?.cancel() + } + } } finally { // we are using a finally here instead of onCompletion as there might be a - // possibility where flow gets cancelled right before `emitAll`. + // possibility where flow gets cancelled right before collection starts. barriers.release(key, barrier) } } } + private fun postWriteOrigin(message: BarrierMsg.Open): StoreReadResponseOrigin = + if (message.writeError == null) { + StoreReadResponseOrigin.Fetcher() + } else { + StoreReadResponseOrigin.SourceOfTruth + } + /** * Writes a value to the underlying [SourceOfTruth] and returns any error that occurred. * @@ -198,6 +262,23 @@ internal class SourceOfTruthWithBarrier { + class Barrier( + val message: BarrierMsg, + ) : ReaderSignal + + class Value( + val value: Output?, + val snapshot: BarrierMsg, + ) : ReaderSignal + + class Failed( + val cause: Throwable, + ) : ReaderSignal + + data object Completed : ReaderSignal + } + private sealed class BarrierMsg( val version: Long, ) { diff --git a/store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/FlowStoreTests.kt b/store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/FlowStoreTests.kt index aacbbda83..6b38919a5 100644 --- a/store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/FlowStoreTests.kt +++ b/store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/FlowStoreTests.kt @@ -21,6 +21,9 @@ import kotlinx.coroutines.FlowPreview import kotlinx.coroutines.async import kotlinx.coroutines.cancelAndJoin import kotlinx.coroutines.delay +import kotlinx.coroutines.channels.BufferOverflow +import kotlinx.coroutines.flow.MutableSharedFlow +import kotlinx.coroutines.flow.channelFlow import kotlinx.coroutines.flow.filter import kotlinx.coroutines.flow.first import kotlinx.coroutines.flow.flow @@ -43,6 +46,7 @@ import org.mobilenativefoundation.store.store5.util.asSourceOfTruth import kotlin.test.Test import kotlin.test.assertContains import kotlin.test.assertEquals +import kotlin.test.assertFalse import kotlin.test.assertIs @FlowPreview @@ -1338,4 +1342,66 @@ class FlowStoreTests { cancelAndIgnoreRemainingEvents() } } + + @Test + fun givenLongLivedSourceOfTruthWhenCachedRefreshThenReaderIsCollectedOnce() = + testScope.runTest { + var subscriptions = 0 + var readerCancelled = false + val storage = mutableMapOf() + val updates = + MutableSharedFlow( + extraBufferCapacity = 16, + onBufferOverflow = BufferOverflow.DROP_OLDEST, + ) + val pipeline = + StoreBuilder.from( + fetcher = Fetcher.of { _: Int -> "from-network" }, + sourceOfTruth = + SourceOfTruth.of( + reader = { key -> + channelFlow { + subscriptions++ + try { + send(storage[key]) + updates.collect { send(storage[key]) } + } finally { + readerCancelled = true + } + } + }, + writer = { key, value -> + storage[key] = value + updates.tryEmit(Unit) + }, + ), + ).disableCache() + .buildWithTestScope() + + val collected = mutableListOf>() + val collection = + launch { + pipeline.stream(StoreReadRequest.cached(3, refresh = true)).collect { + collected.add(it) + } + } + advanceUntilIdle() + + assertEquals(1, subscriptions) + assertFalse(readerCancelled) + assertEquals( + listOf( + Loading(origin = StoreReadResponseOrigin.Fetcher()), + Data( + origin = StoreReadResponseOrigin.Fetcher(), + value = "from-network", + ), + ), + collected, + ) + + collection.cancelAndJoin() + advanceUntilIdle() + assertEquals(true, readerCancelled) + } } diff --git a/store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/SourceOfTruthWithBarrierTests.kt b/store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/SourceOfTruthWithBarrierTests.kt index 645a21856..13a842fb4 100644 --- a/store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/SourceOfTruthWithBarrierTests.kt +++ b/store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/SourceOfTruthWithBarrierTests.kt @@ -23,6 +23,9 @@ import kotlinx.coroutines.FlowPreview import kotlinx.coroutines.async import kotlinx.coroutines.cancelAndJoin import kotlinx.coroutines.delay +import kotlinx.coroutines.channels.BufferOverflow +import kotlinx.coroutines.flow.MutableSharedFlow +import kotlinx.coroutines.flow.channelFlow import kotlinx.coroutines.flow.flow import kotlinx.coroutines.isActive import kotlinx.coroutines.currentCoroutineContext @@ -39,6 +42,8 @@ import org.mobilenativefoundation.store.store5.util.InMemoryPersister import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertIs import kotlin.test.assertNull import kotlin.test.assertTrue @@ -334,4 +339,144 @@ class SourceOfTruthWithBarrierTests { collection.cancelAndJoin() } } + + @Test + fun longLivedReaderIsCollectedOnceAcrossWritesAndKeepsFetcherOrigin() = + testScope.runTest { + val fixture = LongLivedReader() + val collected = mutableListOf>() + val reader = + launch { + fixture.source.reader(1, CompletableDeferred(Unit)).collect { + collected.add(it) + } + } + advanceUntilIdle() + assertEquals(1, fixture.subscriptions) + assertFalse(fixture.readerCancelled) + assertEquals( + listOf>( + StoreReadResponse.Data( + origin = StoreReadResponseOrigin.SourceOfTruth, + value = null, + ), + ), + collected, + ) + + fixture.source.write(1, "a") + advanceUntilIdle() + fixture.source.write(1, "b") + advanceUntilIdle() + + assertEquals(1, fixture.subscriptions) + assertFalse(fixture.readerCancelled) + assertTrue(reader.isActive) + assertEquals( + listOf>( + StoreReadResponse.Data( + origin = StoreReadResponseOrigin.SourceOfTruth, + value = null, + ), + StoreReadResponse.Data( + origin = StoreReadResponseOrigin.Fetcher(), + value = "a", + ), + StoreReadResponse.Data( + origin = StoreReadResponseOrigin.Fetcher(), + value = "b", + ), + ), + collected, + ) + + reader.cancelAndJoin() + advanceUntilIdle() + assertTrue(fixture.readerCancelled) + assertEquals(0, fixture.source.barrierCount()) + } + + @Test + fun longLivedReaderIsNotRestartedWhenWriteFails() = + testScope.runTest { + val exception = RuntimeException("write fails") + val fixture = LongLivedReader(writeFailure = exception) + val collected = mutableListOf>() + val reader = + launch { + fixture.source.reader(1, CompletableDeferred(Unit)).collect { + collected.add(it) + } + } + advanceUntilIdle() + + fixture.source.write(1, "nope") + advanceUntilIdle() + + assertEquals(1, fixture.subscriptions) + assertFalse(fixture.readerCancelled) + assertTrue(reader.isActive) + assertEquals( + listOf>( + StoreReadResponse.Data( + origin = StoreReadResponseOrigin.SourceOfTruth, + value = null, + ), + StoreReadResponse.Error.Exception( + origin = StoreReadResponseOrigin.SourceOfTruth, + error = + WriteException( + key = 1, + value = "nope", + cause = exception, + ), + ), + ), + collected, + ) + assertIs( + (collected.last() as StoreReadResponse.Error.Exception).error, + ) + + reader.cancelAndJoin() + advanceUntilIdle() + assertEquals(0, fixture.source.barrierCount()) + } + + /** + * A reader that emits the current row and then stays open, the way Room or a + * [channelFlow] does. Writes notify the same collection instead of completing it. + */ + private class LongLivedReader( + writeFailure: Throwable? = null, + ) { + var subscriptions: Int = 0 + var readerCancelled: Boolean = false + private val storage = mutableMapOf() + private val updates = + MutableSharedFlow( + extraBufferCapacity = 16, + onBufferOverflow = BufferOverflow.DROP_OLDEST, + ) + private val delegate = + PersistentSourceOfTruth( + realReader = { key -> + channelFlow { + subscriptions++ + try { + send(storage[key]) + updates.collect { send(storage[key]) } + } finally { + readerCancelled = true + } + } + }, + realWriter = { key, value -> + writeFailure?.let { throw it } + storage[key] = value + updates.tryEmit(Unit) + }, + ) + val source = SourceOfTruthWithBarrier(delegate) + } }