1+ package kotlinx.coroutines.internal
2+
3+ import kotlinx.atomicfu.*
4+ import kotlinx.coroutines.*
5+ import kotlinx.coroutines.scheduling.ParallelismCompensation
6+ import kotlin.coroutines.*
7+
8+ /* *
9+ * Introduced as part of IntelliJ patches.
10+ *
11+ * CoroutineDispatchers may optionally implement this interface to declare an ability to construct [SoftLimitedDispatcher]
12+ * on top of themselves. This is not possible in general case, because the worker of the underlying dispatcher must
13+ * implement [ParallelismCompensation] and properly propagate such requests to the task it is running.
14+ */
15+ internal interface SoftLimitedParallelism {
16+ fun softLimitedParallelism (parallelism : Int ): CoroutineDispatcher
17+ }
18+
19+ /* *
20+ * Introduced as part of IntelliJ patches.
21+ */
22+ internal fun CoroutineDispatcher.softLimitedParallelism (parallelism : Int ): CoroutineDispatcher {
23+ if (this is SoftLimitedParallelism ) {
24+ return this .softLimitedParallelism(parallelism)
25+ }
26+ // SoftLimitedDispatcher cannot be used on top of LimitedDispatcher, because the latter doesn't propagate compensation requests
27+ throw UnsupportedOperationException (" CoroutineDispatcher.softLimitedParallelism cannot be applied to $this "
28+ }
29+
30+ /* *
31+ * Introduced as part of IntelliJ patches.
32+ *
33+ * Shamelessly copy-pasted from [LimitedDispatcher], but [ParallelismCompensation] is
34+ * implemented for [Worker] to allow compensation.
35+ *
36+ * [ParallelismCompensation] breaks the contract of [LimitedDispatcher] so a separate class is made to implement a
37+ * dispatcher that mostly behaves as limited, but can temporarily increase parallelism if necessary.
38+ */
39+ internal class SoftLimitedDispatcher (
40+ private val dispatcher : CoroutineDispatcher ,
41+ parallelism : Int
42+ ) : CoroutineDispatcher(), Delay by (dispatcher as ? Delay ? : DefaultDelay ), SoftLimitedParallelism {
43+ private val initialParallelism = parallelism
44+ // `parallelism limit - runningWorkers`; may be < 0 if decompensation is expected
45+ private val availablePermits = atomic(parallelism)
46+
47+ private val queue = LockFreeTaskQueue <Runnable >(singleConsumer = false )
48+
49+ private val workerAllocationLock = SynchronizedObject ()
50+
51+ override fun limitedParallelism (parallelism : Int ): CoroutineDispatcher {
52+ return super .limitedParallelism(parallelism)
53+ }
54+
55+ override fun softLimitedParallelism (parallelism : Int ): CoroutineDispatcher {
56+ parallelism.checkParallelism()
57+ if (parallelism >= initialParallelism) return this
58+ return SoftLimitedDispatcher (this , parallelism)
59+ }
60+
61+ override fun dispatch (context : CoroutineContext , block : Runnable ) {
62+ dispatchInternal(block) { worker ->
63+ dispatcher.dispatch(this , worker)
64+ }
65+ }
66+
67+ @InternalCoroutinesApi
68+ override fun dispatchYield (context : CoroutineContext , block : Runnable ) {
69+ dispatchInternal(block) { worker ->
70+ dispatcher.dispatchYield(this , worker)
71+ }
72+ }
73+
74+ /* *
75+ * Tries to dispatch the given [block].
76+ * If there are not enough workers, it starts a new one via [startWorker].
77+ */
78+ private inline fun dispatchInternal (block : Runnable , startWorker : (Worker ) -> Unit ) {
79+ queue.addLast(block)
80+ if (availablePermits.value <= 0 ) return
81+ if (! tryAllocateWorker()) return
82+ val task = obtainTaskOrDeallocateWorker() ? : return
83+ startWorker(Worker (task))
84+ }
85+
86+ /* *
87+ * Tries to obtain the permit to start a new worker.
88+ */
89+ private fun tryAllocateWorker (): Boolean {
90+ synchronized(workerAllocationLock) {
91+ val permits = availablePermits.value
92+ if (permits <= 0 ) return false
93+ return availablePermits.compareAndSet(permits, permits - 1 )
94+ }
95+ }
96+
97+ /* *
98+ * Obtains the next task from the queue, or logically deallocates the worker if the queue is empty.
99+ */
100+ private fun obtainTaskOrDeallocateWorker (): Runnable ? {
101+ val permits = availablePermits.value
102+ if (permits < 0 ) { // decompensation
103+ if (availablePermits.compareAndSet(permits, permits + 1 )) return null
104+ }
105+ while (true ) {
106+ when (val nextTask = queue.removeFirstOrNull()) {
107+ null -> synchronized(workerAllocationLock) {
108+ availablePermits.incrementAndGet()
109+ if (queue.size == 0 ) return null
110+ availablePermits.decrementAndGet()
111+ }
112+ else -> return nextTask
113+ }
114+ }
115+ }
116+
117+ /* *
118+ * Every running Worker holds a permit
119+ */
120+ private inner class Worker (private var currentTask : Runnable ) : Runnable, ParallelismCompensation {
121+ override fun run () {
122+ var fairnessCounter = 0
123+ while (true ) {
124+ try {
125+ currentTask.run ()
126+ } catch (e: Throwable ) {
127+ handleCoroutineException(EmptyCoroutineContext , e)
128+ }
129+ currentTask = obtainTaskOrDeallocateWorker() ? : return
130+ // 16 is our out-of-thin-air constant to emulate fairness. Used in JS dispatchers as well
131+ if (++ fairnessCounter >= 16 && dispatcher.isDispatchNeeded(this @SoftLimitedDispatcher)) {
132+ // Do "yield" to let other views execute their runnable as well
133+ // Note that we do not decrement 'runningWorkers' as we are still committed to our part of work
134+ dispatcher.dispatch(this @SoftLimitedDispatcher, this )
135+ return
136+ }
137+ }
138+ }
139+
140+ override fun increaseParallelismAndLimit () {
141+ val newTask = obtainTaskOrDeallocateWorker() // either increases the number of permits or we launch a new worker (which holds a permit)
142+ if (newTask != null ) {
143+ dispatcher.dispatch(this @SoftLimitedDispatcher, Worker (newTask))
144+ }
145+ (currentTask as ? ParallelismCompensation )?.increaseParallelismAndLimit()
146+ }
147+
148+ override fun decreaseParallelismLimit () {
149+ try {
150+ (currentTask as ? ParallelismCompensation )?.decreaseParallelismLimit()
151+ } finally {
152+ availablePermits.decrementAndGet()
153+ }
154+ }
155+ }
156+ }
0 commit comments