Skip to content

Stop restarting a long-lived Source of Truth reader on every write #765

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
djain2405 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Stop restarting a long-lived Source of Truth reader on every write #765

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
205 changes: 143 additions & 62 deletions ...ommonMain/kotlin/org/mobilenativefoundation/store/store5/impl/SourceOfTruthWithBarrier.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -71,81 +69,147 @@ internal class SourceOfTruthWithBarrier<Key : Any, Network : Any, Output : Any,
val readerVersion: Long = versionCounter.incrementAndGet()
try {
lock.await()
emitAll(
barrier
.flatMapLatest { barrierMessage ->
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<ReaderSignal<Output>>(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<StoreReadResponse<Output?>> =
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<Output?>
}
}.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.
*
Expand Down Expand Up @@ -198,6 +262,23 @@ internal class SourceOfTruthWithBarrier<Key : Any, Network : Any, Output : Any,
delegate.deleteAll()
}

private sealed interface ReaderSignal<out Output> {
class Barrier(
val message: BarrierMsg,
) : ReaderSignal<Nothing>

class Value<Output>(
val value: Output?,
val snapshot: BarrierMsg,
) : ReaderSignal<Output>

class Failed(
val cause: Throwable,
) : ReaderSignal<Nothing>

data object Completed : ReaderSignal<Nothing>
}

private sealed class BarrierMsg(
val version: Long,
) {
Expand Down
66 changes: 66 additions & 0 deletions store/src/commonTest/kotlin/org/mobilenativefoundation/store/store5/FlowStoreTests.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -1338,4 +1342,66 @@ class FlowStoreTests {
cancelAndIgnoreRemainingEvents()
}
}

@Test
fun givenLongLivedSourceOfTruthWhenCachedRefreshThenReaderIsCollectedOnce() =
testScope.runTest {
var subscriptions = 0
var readerCancelled = false
val storage = mutableMapOf<Int, String>()
val updates =
MutableSharedFlow<Unit>(
extraBufferCapacity = 16,
onBufferOverflow = BufferOverflow.DROP_OLDEST,
)
val pipeline =
StoreBuilder.from(
fetcher = Fetcher.of { _: Int -> "from-network" },
sourceOfTruth =
SourceOfTruth.of<Int, String, String>(
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<StoreReadResponse<String>>()
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)
}
}
Loading
Loading