@@ -84,12 +84,41 @@ public suspend inline fun <T> Semaphore.withPermit(action: () -> T): T {
84
84
}
85
85
86
86
private class SemaphoreImpl (private val permits : Int , acquiredPermits : Int ) : Semaphore {
87
+ /*
88
+ The queue of waiting acquirers is essentially an infinite array based on the list of segments
89
+ (see `SemaphoreSegment`); each segment contains a fixed number of slots. To determine a slot for each enqueue
90
+ and dequeue operation, we increment the corresponding counter at the beginning of the operation
91
+ and use the value before the increment as a slot number. This way, each enqueue-dequeue pair
92
+ works with an individual cell. We use the corresponding segment pointers to find the required ones.
93
+
94
+ Here is a state machine for cells. Note that only one `acquire` and at most one `release` operation
95
+ can deal with each cell, and that `release` uses `getAndSet(PERMIT)` to perform transitions for performance reasons
96
+ so that the state `PERMIT` represents different logical states.
97
+
98
+ +------+ `acquire` suspends +------+ `release` tries +--------+ // if `cont.tryResume(..)` succeeds, then
99
+ | NULL | -------------------> | cont | -------------------> | PERMIT | (cont RETRIEVED) // the corresponding `acquire` operation gets
100
+ +------+ +------+ to resume `cont` +--------+ // a permit and the `release` one completes.
101
+ | |
102
+ | | `acquire` request is cancelled and the continuation is
103
+ | `release` comes | replaced with a special `CANCEL` token to avoid memory leaks
104
+ | to the slot before V
105
+ | `acquire` and puts +-----------+ `release` has +--------+
106
+ | a permit into the | CANCELLED | -----------------> | PERMIT | (RElEASE FAILED)
107
+ | slot, waiting for +-----------+ failed +--------+
108
+ | `acquire` after
109
+ | that.
110
+ |
111
+ | `acquire` gets +-------+
112
+ | +-----------------> | TAKEN | (ELIMINATION HAPPENED)
113
+ V | the permit +-------+
114
+ +--------+ |
115
+ | PERMIT | -<
116
+ +--------+ |
117
+ | `release` has waited a bounded time, +--------+
118
+ +---------------------------------------> | BROKEN | (BOTH RELEASE AND ACQUIRE FAILED)
119
+ but `acquire` has not come +--------+
120
+ */
87
121
88
- // The queue of waiting acquirers is essentially an infinite array based on the list of segments
89
- // (see `SemaphoreSegment`); each segment contains a fixed number of slots. To determine a slot for each enqueue
90
- // and dequeue operation, we increment the corresponding counter at the beginning of the operation
91
- // and use the value before the increment as a slot number. This way, each enqueue-dequeue pair
92
- // works with an individual cell.We use the corresponding segment pointer to find the required ones.
93
122
private val head: AtomicRef <SemaphoreSegment >
94
123
private val deqIdx = atomic(0L )
95
124
private val tail: AtomicRef <SemaphoreSegment >
@@ -123,74 +152,100 @@ private class SemaphoreImpl(private val permits: Int, acquiredPermits: Int) : Se
123
152
override suspend fun acquire () {
124
153
val p = _availablePermits .getAndDecrement()
125
154
if (p > 0 ) return // permit acquired
126
- addToQueueAndSuspend()
155
+ // While it looks better when the following function is inlined,
156
+ // it is important to make `suspend` function invocations in a way
157
+ // so that the tail-call optimization can be applied.
158
+ acquireSlowPath()
127
159
}
128
160
129
- override fun release () {
130
- val p = incPermits()
131
- if (p >= 0 ) return // no waiters
132
- resumeNextFromQueue()
161
+ private suspend fun acquireSlowPath () = suspendAtomicCancellableCoroutineReusable<Unit > sc@ { cont ->
162
+ while (true ) {
163
+ if (addAcquireToQueue(cont)) return @sc
164
+ val p = _availablePermits .getAndDecrement()
165
+ if (p > 0 ) { // permit acquired
166
+ cont.resume(Unit )
167
+ return @sc
168
+ }
169
+ }
133
170
}
134
171
135
- fun incPermits () = _availablePermits .getAndUpdate { cur ->
136
- check(cur < permits) { " The number of released permits cannot be greater than $permits " }
137
- cur + 1
172
+ override fun release () {
173
+ while (true ) {
174
+ val p = _availablePermits .getAndUpdate { cur ->
175
+ check(cur < permits) { " The number of released permits cannot be greater than $permits " }
176
+ cur + 1
177
+ }
178
+ if (p >= 0 ) return
179
+ if (tryResumeNextFromQueue()) return
180
+ }
138
181
}
139
182
140
- private suspend fun addToQueueAndSuspend () = suspendAtomicCancellableCoroutineReusable<Unit > sc@{ cont ->
183
+ /* *
184
+ * Returns `false` if the received permit cannot be used and the calling operation should restart.
185
+ */
186
+ private fun addAcquireToQueue (cont : CancellableContinuation <Unit >): Boolean {
141
187
val curTail = this .tail.value
142
188
val enqIdx = enqIdx.getAndIncrement()
143
189
val segment = this .tail.findSegmentAndMoveForward(id = enqIdx / SEGMENT_SIZE , startFrom = curTail,
144
- createNewSegment = ::createSegment).run { segment } // cannot be closed
190
+ createNewSegment = ::createSegment).segment // cannot be closed
145
191
val i = (enqIdx % SEGMENT_SIZE ).toInt()
146
- if (segment.get(i) == = RESUMED || ! segment.cas(i, null , cont)) {
147
- // already resumed
192
+ // the regular (fast) path -- if the cell is empty, try to install continuation
193
+ if (segment.cas(i, null , cont)) { // installed continuation successfully
194
+ cont.invokeOnCancellation(CancelSemaphoreAcquisitionHandler (segment, i).asHandler)
195
+ return true
196
+ }
197
+ // On CAS failure -- the cell must be either PERMIT or BROKEN
198
+ // If the cell already has PERMIT from tryResumeNextFromQueue, try to grab it
199
+ if (segment.cas(i, PERMIT , TAKEN )) { // took permit thus eliminating acquire/release pair
148
200
cont.resume(Unit )
149
- return @sc
201
+ return true
150
202
}
151
- cont.invokeOnCancellation(CancelSemaphoreAcquisitionHandler (this , segment, i).asHandler)
203
+ assert { segment.get(i) == = BROKEN } // it must be broken in this case, no other way around it
204
+ return false // broken cell, need to retry on a different cell
152
205
}
153
206
154
207
@Suppress(" UNCHECKED_CAST" )
155
- internal fun resumeNextFromQueue () {
156
- try_again@ while (true ) {
157
- val curHead = this .head.value
158
- val deqIdx = deqIdx.getAndIncrement()
159
- val id = deqIdx / SEGMENT_SIZE
160
- val segment = this .head.findSegmentAndMoveForward(id, startFrom = curHead,
161
- createNewSegment = ::createSegment).run { segment } // cannot be closed
162
- segment.cleanPrev()
163
- if (segment.id > id) {
164
- this .deqIdx.updateIfLower(segment.id * SEGMENT_SIZE )
165
- continue @try_again
208
+ private fun tryResumeNextFromQueue (): Boolean {
209
+ val curHead = this .head.value
210
+ val deqIdx = deqIdx.getAndIncrement()
211
+ val id = deqIdx / SEGMENT_SIZE
212
+ val segment = this .head.findSegmentAndMoveForward(id, startFrom = curHead,
213
+ createNewSegment = ::createSegment).segment // cannot be closed
214
+ segment.cleanPrev()
215
+ if (segment.id > id) return false
216
+ val i = (deqIdx % SEGMENT_SIZE ).toInt()
217
+ val cellState = segment.getAndSet(i, PERMIT ) // set PERMIT and retrieve the prev cell state
218
+ when {
219
+ cellState == = null -> {
220
+ // Acquire has not touched this cell yet, wait until it comes for a bounded time
221
+ // The cell state can only transition from PERMIT to TAKEN by addAcquireToQueue
222
+ repeat(MAX_SPIN_CYCLES ) {
223
+ if (segment.get(i) == = TAKEN ) return true
224
+ }
225
+ // Try to break the slot in order not to wait
226
+ return ! segment.cas(i, PERMIT , BROKEN )
166
227
}
167
- val i = (deqIdx % SEGMENT_SIZE ).toInt()
168
- val cont = segment.getAndSet(i, RESUMED )
169
- if (cont == = null ) return // just resumed
170
- if (cont == = CANCELLED ) continue @try_again
171
- (cont as CancellableContinuation <Unit >).resume(Unit )
172
- return
228
+ cellState == = CANCELLED -> return false // the acquire was already cancelled
229
+ else -> return (cellState as CancellableContinuation <Unit >).tryResume()
173
230
}
174
231
}
175
232
}
176
233
177
- private inline fun AtomicLong.updateIfLower (value : Long ): Unit = loop { cur ->
178
- if (cur >= value || compareAndSet(cur, value)) return
234
+ private fun CancellableContinuation<Unit>.tryResume (): Boolean {
235
+ val token = tryResume(Unit ) ? : return false
236
+ completeResume(token)
237
+ return true
179
238
}
180
239
181
240
private class CancelSemaphoreAcquisitionHandler (
182
- private val semaphore : SemaphoreImpl ,
183
241
private val segment : SemaphoreSegment ,
184
242
private val index : Int
185
243
) : CancelHandler() {
186
244
override fun invoke (cause : Throwable ? ) {
187
- val p = semaphore.incPermits()
188
- if (p >= 0 ) return
189
- if (segment.cancel(index)) return
190
- semaphore.resumeNextFromQueue()
245
+ segment.cancel(index)
191
246
}
192
247
193
- override fun toString () = " CancelSemaphoreAcquisitionHandler[$semaphore , $ segment , $index ]"
248
+ override fun toString () = " CancelSemaphoreAcquisitionHandler[$segment , $index ]"
194
249
}
195
250
196
251
private fun createSegment (id : Long , prev : SemaphoreSegment ? ) = SemaphoreSegment (id, prev, 0 )
@@ -202,6 +257,11 @@ private class SemaphoreSegment(id: Long, prev: SemaphoreSegment?, pointers: Int)
202
257
@Suppress(" NOTHING_TO_INLINE" )
203
258
inline fun get (index : Int ): Any? = acquirers[index].value
204
259
260
+ @Suppress(" NOTHING_TO_INLINE" )
261
+ inline fun set (index : Int , value : Any? ) {
262
+ acquirers[index].value = value
263
+ }
264
+
205
265
@Suppress(" NOTHING_TO_INLINE" )
206
266
inline fun cas (index : Int , expected : Any? , value : Any? ): Boolean = acquirers[index].compareAndSet(expected, value)
207
267
@@ -210,19 +270,23 @@ private class SemaphoreSegment(id: Long, prev: SemaphoreSegment?, pointers: Int)
210
270
211
271
// Cleans the acquirer slot located by the specified index
212
272
// and removes this segment physically if all slots are cleaned.
213
- fun cancel (index : Int ): Boolean {
214
- // Try to cancel the slot
215
- val cancelled = getAndSet (index, CANCELLED ) != = RESUMED
273
+ fun cancel (index : Int ) {
274
+ // Clean the slot
275
+ set (index, CANCELLED )
216
276
// Remove this segment if needed
217
277
onSlotCleaned()
218
- return cancelled
219
278
}
220
279
221
280
override fun toString () = " SemaphoreSegment[id=$id , hashCode=${hashCode()} ]"
222
281
}
223
-
224
282
@SharedImmutable
225
- private val RESUMED = Symbol (" RESUMED" )
283
+ private val MAX_SPIN_CYCLES = systemProp(" kotlinx.coroutines.semaphore.maxSpinCycles" , 100 )
284
+ @SharedImmutable
285
+ private val PERMIT = Symbol (" PERMIT" )
286
+ @SharedImmutable
287
+ private val TAKEN = Symbol (" TAKEN" )
288
+ @SharedImmutable
289
+ private val BROKEN = Symbol (" BROKEN" )
226
290
@SharedImmutable
227
291
private val CANCELLED = Symbol (" CANCELLED" )
228
292
@SharedImmutable
0 commit comments