-
Notifications
You must be signed in to change notification settings - Fork 443
/
CountDownLatch.kt
48 lines (42 loc) · 1.33 KB
/
CountDownLatch.kt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
package arrow.fx.coroutines
import arrow.core.continuations.AtomicRef
import arrow.core.continuations.loop
import kotlinx.coroutines.CompletableDeferred
/**
* [CountDownLatch] allows for awaiting a given number of countdown signals.
* Models the behavior of java.util.concurrent.CountDownLatch in Kotlin with `suspend`.
*
* Must be initialised with an [initial] value of 1 or higher,
* if constructed with 0 or negative value then it throws [IllegalArgumentException].
*/
public class CountDownLatch(private val initial: Long) {
private val signal = CompletableDeferred<Unit>()
private val count = AtomicRef(initial)
init {
require(initial > 0) {
"CountDownLatch must be constructed with positive non-zero initial count, but was $initial"
}
}
/** Remaining count */
public fun count(): Long = count.get()
/** Await [count] to reach zero */
public suspend fun await(): Unit =
when (count.get()) {
0L -> { }
else -> signal.await()
}
/** Decrement [count] by one */
@Suppress("ReturnCount")
public fun countDown() {
count.loop { current ->
when {
current == 0L -> return
current == 1L && count.compareAndSet(1L, 0L) -> {
signal.complete(Unit)
return
}
count.compareAndSet(current, current - 1) -> return
}
}
}
}