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 , name : String? ): CoroutineDispatcher
17+ }
18+
19+ /* *
20+ * Introduced as part of IntelliJ patches.
21+ */
22+ internal fun CoroutineDispatcher.softLimitedParallelism (parallelism : Int , name : String? ): CoroutineDispatcher {
23+ if (this is SoftLimitedParallelism ) {
24+ return this .softLimitedParallelism(parallelism, name)
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+ private val name : String?
43+ ) : CoroutineDispatcher(), Delay by (dispatcher as ? Delay ? : DefaultDelay ), SoftLimitedParallelism {
44+ private val initialParallelism = parallelism
45+ // `parallelism limit - runningWorkers`; may be < 0 if decompensation is expected
46+ private val availablePermits = atomic(parallelism)
47+
48+ private val queue = LockFreeTaskQueue <Runnable >(singleConsumer = false )
49+
50+ private val workerAllocationLock = SynchronizedObject ()
51+
52+ override fun limitedParallelism (parallelism : Int , name : String? ): CoroutineDispatcher {
53+ return super .limitedParallelism(parallelism, name)
54+ }
55+
56+ override fun softLimitedParallelism (parallelism : Int , name : String? ): CoroutineDispatcher {
57+ parallelism.checkParallelism()
58+ if (parallelism >= initialParallelism) return namedOrThis(name)
59+ return SoftLimitedDispatcher (this , parallelism, name)
60+ }
61+
62+ override fun dispatch (context : CoroutineContext , block : Runnable ) {
63+ dispatchInternal(block) { worker ->
64+ dispatcher.safeDispatch(this , worker)
65+ }
66+ }
67+
68+ @InternalCoroutinesApi
69+ override fun dispatchYield (context : CoroutineContext , block : Runnable ) {
70+ dispatchInternal(block) { worker ->
71+ dispatcher.dispatchYield(this , worker)
72+ }
73+ }
74+
75+ /* *
76+ * Tries to dispatch the given [block].
77+ * If there are not enough workers, it starts a new one via [startWorker].
78+ */
79+ private inline fun dispatchInternal (block : Runnable , startWorker : (Worker ) -> Unit ) {
80+ queue.addLast(block)
81+ if (availablePermits.value <= 0 ) return
82+ if (! tryAllocateWorker()) return
83+ val task = obtainTaskOrDeallocateWorker() ? : return
84+ startWorker(Worker (task))
85+ }
86+
87+ /* *
88+ * Tries to obtain the permit to start a new worker.
89+ */
90+ private fun tryAllocateWorker (): Boolean {
91+ synchronized(workerAllocationLock) {
92+ val permits = availablePermits.value
93+ if (permits <= 0 ) return false
94+ return availablePermits.compareAndSet(permits, permits - 1 )
95+ }
96+ }
97+
98+ /* *
99+ * Obtains the next task from the queue, or logically deallocates the worker if the queue is empty.
100+ */
101+ private fun obtainTaskOrDeallocateWorker (): Runnable ? {
102+ val permits = availablePermits.value
103+ if (permits < 0 ) { // decompensation
104+ if (availablePermits.compareAndSet(permits, permits + 1 )) {
105+ return null
106+ }
107+ }
108+ while (true ) {
109+ when (val nextTask = queue.removeFirstOrNull()) {
110+ null -> synchronized(workerAllocationLock) {
111+ availablePermits.incrementAndGet()
112+ if (queue.size == 0 ) return null
113+ availablePermits.decrementAndGet()
114+ }
115+ else -> return nextTask
116+ }
117+ }
118+ }
119+
120+ override fun toString () = name ? : " $dispatcher .softLimitedParallelism($initialParallelism )"
121+
122+ /* *
123+ * Every running Worker holds a permit
124+ */
125+ private inner class Worker (private var currentTask : Runnable ) : Runnable, ParallelismCompensation {
126+ override fun run () {
127+ var fairnessCounter = 0
128+ while (true ) {
129+ try {
130+ currentTask.run ()
131+ } catch (e: Throwable ) {
132+ handleCoroutineException(EmptyCoroutineContext , e)
133+ }
134+ currentTask = obtainTaskOrDeallocateWorker() ? : return
135+ // 16 is our out-of-thin-air constant to emulate fairness. Used in JS dispatchers as well
136+ if (++ fairnessCounter >= 16 && dispatcher.safeIsDispatchNeeded(this @SoftLimitedDispatcher)) {
137+ // Do "yield" to let other views execute their runnable as well
138+ // Note that we do not decrement 'runningWorkers' as we are still committed to our part of work
139+ dispatcher.safeDispatch(this @SoftLimitedDispatcher, this )
140+ return
141+ }
142+ }
143+ }
144+
145+ override fun increaseParallelismAndLimit () {
146+ val newTask = obtainTaskOrDeallocateWorker() // either increases the number of permits or we launch a new worker (which holds a permit)
147+ if (newTask != null ) {
148+ dispatcher.safeDispatch(this @SoftLimitedDispatcher, Worker (newTask))
149+ }
150+ (currentTask as ? ParallelismCompensation )?.increaseParallelismAndLimit()
151+ }
152+
153+ override fun decreaseParallelismLimit () {
154+ try {
155+ (currentTask as ? ParallelismCompensation )?.decreaseParallelismLimit()
156+ } finally {
157+ availablePermits.decrementAndGet()
158+ }
159+ }
160+ }
161+ }
0 commit comments