Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Helpers to control concurrency for one shot requests using Kotlin coroutines.
import kotlinx.coroutines.CoroutineStart.LAZY
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.async
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.yield
import java.util.concurrent.atomic.AtomicReference
import kotlin.DeprecationLevel.ERROR
/**
* A helper class to execute tasks sequentially in coroutines.
*
* Calling [afterPrevious] will always ensure that all previously requested work completes prior to
* calling the block passed. Any future calls to [afterPrevious] while the current block is running
* will wait for the current block to complete before starting.
*/
class SingleRunner {
/**
* A coroutine mutex implements a lock that may only be taken by one coroutine at a time.
*/
private val mutex = Mutex()
/**
* Ensure that the block will only be executed after all previous work has completed.
*
* When several coroutines call afterPrevious at the same time, they will queue up in the order
* that they call afterPrevious. Then, one coroutine will enter the block at a time.
*
* In the following example, only one save operation (user or song) will be executing at a time.
*
* ```
* class UserAndSongSaver {
* val singleRunner = SingleRunner()
*
* fun saveUser(user: User) {
* singleRunner.afterPrevious { api.post(user) }
* }
*
* fun saveSong(song: Song) {
* singleRunner.afterPrevious { api.post(song) }
* }
* }
* ```
*
* @param block the code to run after previous work is complete.
*/
suspend fun <T> afterPrevious(block: suspend () -> T): T {
// Before running the block, ensure that no other blocks are running by taking a lock on the
// mutex.
// The mutex will be released automatically when we return.
// If any other block were already running when we get here, it will wait for it to complete
// before entering the `withLock` block.
mutex.withLock {
return block()
}
}
}
/**
* A controlled runner decides what to do when new tasks are run.
*
* Note: This implementation is for example only. It will not work in the presence of
* multi-threading and is not safe to call from Dispatchers.IO or Dispatchers.Default. In
* real code use the thread-safe implementation of [ControlledRunner] code listed below.
*
* By calling [joinPreviousOrRun], the new task will be discarded and the result of the previous task
* will be returned. This is useful when you want to ensure that a network request to the same
* resource does not flood.
*
* By calling [cancelPreviousThenRun], the old task will *always* be cancelled and then the new task will
* be run. This is useful in situations where a new event implies that the previous work is no
* longer relevant such as sorting or filtering a list.
*/
@Deprecated("This code is not thread-safe and should not be used. Use " +
"the ControlledRunner implementation below instead.", level = ERROR)
class ControlledRunnerExampleImplementation<T> {
private var activeTask: Deferred<T>? = null
/**
* Cancel all previous tasks before calling block.
*
* When several coroutines call cancelPreviousThenRun at the same time, only one will run and
* the others will be cancelled.
*/
@Deprecated("This code is not thread-safe. Use ControlledRunner below instead.",
level = ERROR)
suspend fun cancelPreviousThenRun(block: suspend () -> T): T {
// If there is an activeTask, cancel it because it's result is no longer needed
//
// By waiting for the cancellation to complete with `cancelAndJoin` we know that activeTask
// has stopped executing before continuing.
activeTask?.cancelAndJoin()
// use a coroutineScope builder to safely start a new coroutine in a suspend function
return coroutineScope {
// create a new task to call the block
val newTask = async {
block()
}
// when the new task completes, reset activeTask to null
// this will be called by cancellation as well as normal completion
newTask.invokeOnCompletion {
activeTask = null
}
// save the newTask to activeTask, then wait for it to complete and return the result
activeTask = newTask
newTask.await()
}
}
/**
* Don't run the new block if a previous block is running, instead wait for the previous block
* and return it's result.
*
* When several coroutines call joinPreviousOrRun at the same time, only one will run and
* the others will return the result from the winner.
*/
@Deprecated("This code is not thread-safe. Use ControlledRunner below instead.",
level = ERROR)
suspend fun joinPreviousOrRun(block: suspend () -> T): T {
// if there is an activeTask, return it's result and don't run the block
activeTask?.let {
return it.await()
}
// use a coroutineScope builder to safely start a new coroutine in a suspend function
return coroutineScope {
// create a new task to call the block
val newTask = async {
block()
}
// when the task completes, reset activeTask to null
newTask.invokeOnCompletion {
activeTask = null
}
// save newTask to activeTask, then wait for it to complete and return the result
activeTask = newTask
newTask.await()
}
}
}
/**
* A controlled runner decides what to do when new tasks are run.
*
* By calling [joinPreviousOrRun], the new task will be discarded and the result of the previous task
* will be returned. This is useful when you want to ensure that a network request to the same
* resource does not flood.
*
* By calling [cancelPreviousThenRun], the old task will *always* be cancelled and then the new task will
* be run. This is useful in situations where a new event implies that the previous work is no
* longer relevant such as sorting or filtering a list.
*/
class ControlledRunner<T> {
/**
* The currently active task.
*
* This uses an atomic reference to ensure that it's safe to update activeTask on both
* Dispatchers.Default and Dispatchers.Main which will execute coroutines on multiple threads at
* the same time.
*/
private val activeTask = AtomicReference<Deferred<T>?>(null)
/**
* Cancel all previous tasks before calling block.
*
* When several coroutines call cancelPreviousThenRun at the same time, only one will run and
* the others will be cancelled.
*
* In the following example, only one sort operation will execute and any previous sorts will be
* cancelled.
*
* ```
* class Products {
* val controlledRunner = ControlledRunner<Product>()
*
* fun sortAscending(): List<Product> {
* return controlledRunner.cancelPreviousThenRun { dao.loadSortedAscending() }
* }
*
* fun sortDescending(): List<Product> {
* return controlledRunner.cancelPreviousThenRun { dao.loadSortedDescending() }
* }
* }
* ```
*
* @param block the code to run after previous work is cancelled.
* @return the result of block, if this call was not cancelled prior to returning.
*/
suspend fun cancelPreviousThenRun(block: suspend() -> T): T {
// fast path: if we already know about an active task, just cancel it right away.
activeTask.get()?.cancelAndJoin()
return coroutineScope {
// Create a new coroutine, but don't start it until it's decided that this block should
// execute. In the code below, calling await() on newTask will cause this coroutine to
// start.
val newTask = async(start = LAZY) {
block()
}
// When newTask completes, ensure that it resets activeTask to null (if it was the
// current activeTask).
newTask.invokeOnCompletion {
activeTask.compareAndSet(newTask, null)
}
// Kotlin ensures that we only set result once since it's a val, even though it's set
// inside the while(true) loop.
val result: T
// Loop until we are sure that newTask is ready to execute (all previous tasks are
// cancelled)
while(true) {
if (!activeTask.compareAndSet(null, newTask)) {
// some other task started before newTask got set to activeTask, so see if it's
// still running when we call get() here. If so, we can cancel it.
// we will always start the loop again to see if we can set activeTask before
// starting newTask.
activeTask.get()?.cancelAndJoin()
// yield here to avoid a possible tight loop on a single threaded dispatcher
yield()
} else {
// happy path - we set activeTask so we are ready to run newTask
result = newTask.await()
break
}
}
// Kotlin ensures that the above loop always sets result exactly once, so we can return
// it here!
result
}
}
/**
* Don't run the new block if a previous block is running, instead wait for the previous block
* and return it's result.
*
* When several coroutines call jonPreviousOrRun at the same time, only one will run and
* the others will return the result from the winner.
*
* In the following example, only one network operation will execute at a time and any other
* requests will return the result from the "in flight" request.
*
* ```
* class Products {
* val controlledRunner = ControlledRunner<Product>()
*
* fun fetchProducts(): List<Product> {
* return controlledRunner.joinPreviousOrRun {
* val results = api.fetchProducts()
* dao.insert(results)
* results
* }
* }
* }
* ```
*
* @param block the code to run if and only if no other task is currently running
* @return the result of block, or if another task was running the result of that task instead.
*/
suspend fun joinPreviousOrRun(block: suspend () -> T): T {
// fast path: if there's already an active task, just wait for it and return the result
activeTask.get()?.let {
return it.await()
}
return coroutineScope {
// Create a new coroutine, but don't start it until it's decided that this block should
// execute. In the code below, calling await() on newTask will cause this coroutine to
// start.
val newTask = async(start = LAZY) {
block()
}
newTask.invokeOnCompletion {
activeTask.compareAndSet(newTask, null)
}
// Kotlin ensures that we only set result once since it's a val, even though it's set
// inside the while(true) loop.
val result: T
// Loop until we figure out if we need to run newTask, or if there is a task that's
// already running we can join.
while(true) {
if (!activeTask.compareAndSet(null, newTask)) {
// some other task started before newTask got set to activeTask, so see if it's
// still running when we call get() here. There is a chance that it's already
// been completed before the call to get, in which case we need to start the
// loop over and try again.
val currentTask = activeTask.get()
if (currentTask != null) {
// happy path - we found the other task so use that one instead of newTask
newTask.cancel()
result = currentTask.await()
break
} else {
// retry path - the other task completed before we could get it, loop to try
// setting activeTask again.
// call yield here in case we're executing on a single threaded dispatcher
// like Dispatchers.Main to allow other work to happen.
yield()
}
} else {
// happy path - we were able to set activeTask, so start newTask and return its
// result
result = newTask.await()
break
}
}
// Kotlin ensures that the above loop always sets result exactly once, so we can return
// it here!
result
}
}
}
@mrleolink

This comment has been minimized.

Copy link

commented Jun 6, 2019

I'm having a hard time understand why in function joinPreviousOrRun, we need to use activeTask.compareAndSet(newTask, null) instead of activeTask.set(null) when we are already cancelling all the new tasks (at line 300) when current task hasn't completed.

Do you mind to explain that?

@ivanbartsov

This comment has been minimized.

Copy link

commented Aug 27, 2019

@mrleolink

I'm having a hard time understand why in function joinPreviousOrRun, we need to use activeTask.compareAndSet(newTask, null) instead of activeTask.set(null) when we are already cancelling all the new tasks (at line 300) when current task hasn't completed.

Do you mind to explain that?

This may be late but, I think I get it, although I'm not the author so I may have missed something:

the compareAndSet thing makes sure that we only submit our newTask to be the current atomically-referenced one if there was none already: another coroutine with this code may have beaten us by a hair to this line and set the atomic reference before us -- in which case, we should wait for it and return its result. Calling activeTask.set(null) we would lose reference to whatever task there was before us each time -- effectively it would be the same as not using all this atomic logic at all: each time we would overwrite the atomic reference without ever checking if there was another task, i.e. the atomic reference could just be removed altogether and what we'd be left with would just be a weirdly convoluted way to call async{}.await().

What does seem a tad weird to me is the newTask.cancel() call at line 300: at that point, we never really started it, because of the lazy start mode ((start = LAZY) at line 277) and since it's in the "setting atomic reference failed" block, we didn't pass the task to any external code, so nothing else could've called .await() on it and cause it to start. Is it to untie newTask's job from the parent coroutineScope{}?

PS Sean, kudos on the great article on coroutines on Android, one of the most detailed and to the point I've read so far

@objcode

This comment has been minimized.

Copy link
Owner Author

commented Aug 27, 2019

That explanation is spot on - and I think you're right that the call to cancel in 300 is unnecessary. It won't hurt anything, but it's doing a bit of extra work that's not needed.

@ivanbartsov

This comment has been minimized.

Copy link

commented Aug 27, 2019

That explanation is spot on - and I think you're right that the call to cancel in 300 is unnecessary. It won't hurt anything, but it's doing a bit of extra work that's not needed.

Now I'm wondering what will happen in regard to job inheritance if it's removed -- this is beyond my current depth of understanding of coroutines. From this answer here it seems lazy async builders still get their jobs added to parent's job, so removing that .cancel() call might actually end up causing the coroutineScope{} block wait for newTask's completion indefinitely, in which case, sorry for the disruption

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.