SingleProcessDataStore.kt
/*
* Copyright 2020 The Android Open Source Project
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package androidx.datastore.core
import androidx.datastore.core.handlers.NoOpCorruptionHandler
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.completeWith
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.dropWhile
import kotlinx.coroutines.flow.emitAll
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.coroutineContext
/**
* Single process implementation of DataStore. This is NOT multi-process safe.
*/
internal class SingleProcessDataStore<T>(
private val storage: Storage<T>,
/**
* The list of initialization tasks to perform. These tasks will be completed before any data
* is published to the data and before any read-modify-writes execute in updateData. If
* any of the tasks fail, the tasks will be run again the next time data is collected or
* updateData is called. Init tasks should not wait on results from data - this will
* result in deadlock.
*/
initTasksList: List<suspend (api: InitializerApi<T>) -> Unit> = emptyList(),
private val corruptionHandler: CorruptionHandler<T> = NoOpCorruptionHandler<T>(),
private val scope: CoroutineScope = CoroutineScope(ioDispatcher() + SupervisorJob())
) : DataStore<T> {
val connection: StorageConnection<T> by lazy {
storage.createConnection()
}
override val data: Flow<T> = flow {
/**
* If downstream flow is UnInitialized, no data has been read yet, we need to trigger a new
* read then start emitting values once we have seen a new value (or exception).
*
* If downstream flow has a ReadException, there was an exception last time we tried to read
* data. We need to trigger a new read then start emitting values once we have seen a new
* value (or exception).
*
* If downstream flow has Data, we should just start emitting from downstream flow.
*
* If Downstream flow is Final, the scope has been cancelled so the data store is no
* longer usable. We should just propagate this exception.
*
* State always starts at null. null can transition to ReadException, Data or
* Final. ReadException can transition to another ReadException, Data or Final.
* Data can transition to another Data or Final. Final will not change.
*/
val currentDownStreamFlowState = downstreamFlow.value
if (currentDownStreamFlowState !is Data) {
// We need to send a read request because we don't have data yet.
actor.offer(Message.Read(currentDownStreamFlowState))
}
emitAll(
downstreamFlow.dropWhile {
if (currentDownStreamFlowState is Data<T> ||
currentDownStreamFlowState is Final<T>
) {
// We don't need to drop any Data or Final values.
false
} else {
// we need to drop the last seen state since it was either an exception or
// wasn't yet initialized. Since we sent a message to actor, we *will* see a
// new value.
it === currentDownStreamFlowState
}
}.map {
when (it) {
is ReadException<T> -> throw it.readException
is Final<T> -> throw it.finalException
is Data<T> -> it.value
is UnInitialized -> error(
"This is a bug in DataStore. Please file a bug at: " +
"https://issuetracker.google.com/issues/new?" +
"component=907884&template=1466542"
)
}
}
)
}
override suspend fun updateData(transform: suspend (t: T) -> T): T {
/**
* The states here are the same as the states for reads. Additionally we send an ack that
* the actor *must* respond to (even if it is cancelled).
*/
val ack = CompletableDeferred<T>()
val currentDownStreamFlowState = downstreamFlow.value
val updateMsg =
Message.Update(transform, ack, currentDownStreamFlowState, coroutineContext)
actor.offer(updateMsg)
return ack.await()
}
@Suppress("UNCHECKED_CAST")
private val downstreamFlow = MutableStateFlow(UnInitialized as State<T>)
private var initTasks: List<suspend (api: InitializerApi<T>) -> Unit>? =
initTasksList.toList()
private val actor = SimpleActor<Message<T>>(
scope = scope,
onComplete = {
it?.let {
downstreamFlow.value = Final(it)
}
// We expect it to always be non-null but we will leave the alternative as a no-op
// just in case.
connection.close()
},
onUndeliveredElement = { msg, ex ->
if (msg is Message.Update) {
// TODO(rohitsat): should we instead use scope.ensureActive() to get the original
// cancellation cause? Should we instead have something like
// UndeliveredElementException?
msg.ack.completeExceptionally(
ex ?: CancellationException(
"DataStore scope was cancelled before updateData could complete"
)
)
}
}
) { msg ->
when (msg) {
is Message.Read -> {
handleRead(msg)
}
is Message.Update -> {
handleUpdate(msg)
}
}
}
private suspend fun handleRead(read: Message.Read<T>) {
when (val currentState = downstreamFlow.value) {
is Data -> {
// We already have data so just return...
}
is ReadException -> {
if (currentState === read.lastState) {
readAndInitOrPropagateFailure()
}
// Someone else beat us but also failed. The collector has already
// been signalled so we don't need to do anything.
}
UnInitialized -> {
readAndInitOrPropagateFailure()
}
is Final -> error("Can't read in final state.") // won't happen
}
}
private suspend fun handleUpdate(update: Message.Update<T>) {
// All branches of this *must* complete ack either successfully or exceptionally.
// We must *not* throw an exception, just propagate it to the ack.
update.ack.completeWith(
runCatching {
when (val currentState = downstreamFlow.value) {
is Data -> {
// We are already initialized, we just need to perform the update
transformAndWrite(update.transform, update.callerContext)
}
is ReadException, is UnInitialized -> {
if (currentState === update.lastState) {
// we need to try to read again
readAndInitOrPropagateAndThrowFailure()
// We've successfully read, now we need to perform the update
transformAndWrite(update.transform, update.callerContext)
} else {
// Someone else beat us to read but also failed. We just need to
// signal the writer that is waiting on ack.
// This cast is safe because we can't be in the UnInitialized
// state if the state has changed.
throw (currentState as ReadException).readException
}
}
is Final -> throw currentState.finalException // won't happen
}
}
)
}
private suspend fun readAndInitOrPropagateAndThrowFailure() {
try {
readAndInit()
} catch (throwable: Throwable) {
downstreamFlow.value = ReadException(throwable)
throw throwable
}
}
private suspend fun readAndInitOrPropagateFailure() {
try {
readAndInit()
} catch (throwable: Throwable) {
downstreamFlow.value = ReadException(throwable)
}
}
private suspend fun readAndInit() {
// This should only be called if we don't already have cached data.
check(downstreamFlow.value == UnInitialized || downstreamFlow.value is ReadException)
val updateLock = Mutex()
var initData = readDataOrHandleCorruption()
var initializationComplete: Boolean = false
// TODO(b/151635324): Consider using Context Element to throw an error on re-entrance.
val api = object : InitializerApi<T> {
override suspend fun updateData(transform: suspend (t: T) -> T): T {
return updateLock.withLock() {
if (initializationComplete) {
throw IllegalStateException(
"InitializerApi.updateData should not be " +
"called after initialization is complete."
)
}
val newData = transform(initData)
if (newData != initData) {
connection.writeData(newData)
initData = newData
}
initData
}
}
}
initTasks?.forEach { it(api) }
initTasks = null // Init tasks have run successfully, we don't need them anymore.
updateLock.withLock {
initializationComplete = true
}
downstreamFlow.value = Data(initData, initData.hashCode())
}
private suspend fun readDataOrHandleCorruption(): T {
try {
return connection.readData()
} catch (ex: CorruptionException) {
val newData: T = corruptionHandler.handleCorruption(ex)
try {
connection.writeData(newData)
} catch (writeEx: IOException) {
// If we fail to write the handled data, add the new exception as a suppressed
// exception.
ex.addSuppressed(writeEx)
throw ex
}
// If we reach this point, we've successfully replaced the data on disk with newData.
return newData
}
}
// downstreamFlow.value must be successfully set to data before calling this
private suspend fun transformAndWrite(
transform: suspend (t: T) -> T,
callerContext: CoroutineContext
): T {
// value is not null or an exception because we must have the value set by now so this cast
// is safe.
val curDataAndHash = downstreamFlow.value as Data<T>
curDataAndHash.checkHashCode()
val curData = curDataAndHash.value
val newData = withContext(callerContext) { transform(curData) }
// Check that curData has not changed...
curDataAndHash.checkHashCode()
return if (curData == newData) {
curData
} else {
connection.writeData(newData)
downstreamFlow.value = Data(newData, newData.hashCode())
newData
}
}
}