forked from Kotlin/kotlinx.coroutines
/
CancellationPoint.kt
145 lines (135 loc) · 5.37 KB
/
CancellationPoint.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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
/*
* Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
*/
package kotlinx.coroutines
import kotlinx.atomicfu.AtomicRef
import kotlinx.atomicfu.atomic
import kotlinx.atomicfu.loop
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext
import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
/**
* Makes a blocking code block cancellable (become a cancellation point of the coroutine).
*
* The blocking code block will be interrupted and this function will throw [CancellationException]
* if the coroutine is cancelled.
*
* Example:
* ```
* GlobalScope.launch(Dispatchers.IO) {
* async {
* // This function will throw [CancellationException].
* runInterruptible {
* doSomethingUseful()
*
* // This blocking procedure will be interrupted when this coroutine is canceled
* // by Exception thrown by the below async block.
* doSomethingElseUsefulInterruptible()
* }
* }
*
* async {
* delay(500L)
* throw Exception()
* }
* }
* ```
*
* There is also an optional context parameter to this function to enable single-call conversion of
* interruptible Java methods into main-safe suspending functions like this:
* ```
* // With one call here we are moving the call to Dispatchers.IO and supporting interruption.
* suspend fun <T> BlockingQueue<T>.awaitTake(): T =
* runInterruptible(Dispatchers.IO) { queue.take() }
* ```
*
* @param context additional to [CoroutineScope.coroutineContext] context of the coroutine.
* @param block regular blocking block that will be interrupted on coroutine cancellation.
*/
public suspend fun <T> runInterruptible(
context: CoroutineContext = EmptyCoroutineContext,
block: () -> T
): T {
// fast path: empty context
if (context === EmptyCoroutineContext) { return runInterruptibleInExpectedContext(block) }
// slow path:
return withContext(context) { runInterruptibleInExpectedContext(block) }
}
private suspend fun <T> runInterruptibleInExpectedContext(block: () -> T): T =
suspendCoroutineUninterceptedOrReturn sc@{ uCont ->
try {
// fast path: no job
val job = uCont.context[Job] ?: return@sc block()
// slow path
val threadState = ThreadState().apply { initInterrupt(job) }
try {
block()
} finally {
threadState.clearInterrupt()
}
} catch (e: InterruptedException) {
throw CancellationException()
}
}
private class ThreadState {
fun initInterrupt(job: Job) {
// starts with Init
if (state.value !== Init) throw IllegalStateException("impossible state")
// remembers this running thread
state.value = Working(Thread.currentThread(), null)
// watches the job for cancellation
val cancelHandle =
job.invokeOnCompletion(onCancelling = true, invokeImmediately = true, handler = CancelHandler())
// remembers the cancel handle or drops it
state.loop { s ->
when {
s is Working -> if (state.compareAndSet(s, Working(s.thread, cancelHandle))) return
s === Interrupting || s === Interrupted -> return
s === Init || s === Finish -> throw IllegalStateException("impossible state")
else -> throw IllegalStateException("unknown state")
}
}
}
fun clearInterrupt() {
state.loop { s ->
when {
s is Working -> if (state.compareAndSet(s, Finish)) { s.cancelHandle!!.dispose(); return }
s === Interrupting -> Thread.yield() // eases the thread
s === Interrupted -> { Thread.interrupted(); return } // no interrupt leak
s === Init || s === Finish -> throw IllegalStateException("impossible state")
else -> throw IllegalStateException("unknown state")
}
}
}
private inner class CancelHandler : CompletionHandler {
override fun invoke(cause: Throwable?) {
state.loop { s ->
when {
s is Working -> {
if (state.compareAndSet(s, Interrupting)) {
s.thread!!.interrupt()
state.value = Interrupted
return
}
}
s === Finish -> return
s === Interrupting || s === Interrupted -> return
s === Init -> throw IllegalStateException("impossible state")
else -> throw IllegalStateException("unknown state")
}
}
}
}
private val state: AtomicRef<State> = atomic(Init)
private interface State
// initial state
private object Init : State
// cancellation watching is setup and/or the continuation is running
private data class Working(val thread: Thread?, val cancelHandle: DisposableHandle?) : State
// the continuation done running without interruption
private object Finish : State
// interrupting this thread
private object Interrupting: State
// done interrupting
private object Interrupted: State
}