kotlinx-coroutines-core/common/src/sync/Semaphore.kt - platform/external/kotlinx.coroutines - Gitiles

 /*
  * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */

 package kotlinx.coroutines.sync

 import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.internal.*
 import kotlin.contracts.*
 import kotlin.coroutines.*
 import kotlin.math.*
 import kotlin.native.concurrent.SharedImmutable

 /**
  * A counting semaphore for coroutines that logically maintains a number of available permits.
  * Each [acquire] takes a single permit or suspends until it is available.
  * Each [release] adds a permit, potentially releasing a suspended acquirer.
  * Semaphore is fair and maintains a FIFO order of acquirers.
  *
  * Semaphores are mostly used to limit the number of coroutines that have an access to particular resource.
  * Semaphore with `permits = 1` is essentially a [Mutex].
  **/
 public interface Semaphore {
     /**
      * Returns the current number of permits available in this semaphore.
      */
     public val availablePermits: Int

     /**
      * Acquires a permit from this semaphore, suspending until one is available.
      * All suspending acquirers are processed in first-in-first-out (FIFO) order.
      *
      * This suspending function is cancellable. If the [Job] of the current coroutine is cancelled or completed while this
      * function is suspended, this function immediately resumes with [CancellationException].
      * There is a **prompt cancellation guarantee**. If the job was cancelled while this function was
      * suspended, it will not resume successfully. See [suspendCancellableCoroutine] documentation for low-level details.
      * This function releases the semaphore if it was already acquired by this function before the [CancellationException]
      * was thrown.
      *
      * Note, that this function does not check for cancellation when it does not suspend.
      * Use [CoroutineScope.isActive] or [CoroutineScope.ensureActive] to periodically
      * check for cancellation in tight loops if needed.
      *
      * Use [tryAcquire] to try acquire a permit of this semaphore without suspension.
      */
     public suspend fun acquire()

     /**
      * Tries to acquire a permit from this semaphore without suspension.
      *
      * @return `true` if a permit was acquired, `false` otherwise.
      */
     public fun tryAcquire(): Boolean

     /**
      * Releases a permit, returning it into this semaphore. Resumes the first
      * suspending acquirer if there is one at the point of invocation.
      * Throws [IllegalStateException] if the number of [release] invocations is greater than the number of preceding [acquire].
      */
     public fun release()
 }

 /**
  * Creates new [Semaphore] instance.
  * @param permits the number of permits available in this semaphore.
  * @param acquiredPermits the number of already acquired permits,
  *        should be between `0` and `permits` (inclusively).
  */
 @Suppress("FunctionName")
 public fun Semaphore(permits: Int, acquiredPermits: Int = 0): Semaphore = SemaphoreImpl(permits, acquiredPermits)

 /**
  * Executes the given [action], acquiring a permit from this semaphore at the beginning
  * and releasing it after the [action] is completed.
  *
  * @return the return value of the [action].
  */
 @OptIn(ExperimentalContracts::class)
 public suspend inline fun <T> Semaphore.withPermit(action: () -> T): T {
     contract {
         callsInPlace(action, InvocationKind.EXACTLY_ONCE)
     }

     acquire()
     try {
         return action()
     } finally {
         release()
     }
 }

 private class SemaphoreImpl(private val permits: Int, acquiredPermits: Int) : Semaphore {
     /*
        The queue of waiting acquirers is essentially an infinite array based on the list of segments
        (see `SemaphoreSegment`); each segment contains a fixed number of slots. To determine a slot for each enqueue
        and dequeue operation, we increment the corresponding counter at the beginning of the operation
        and use the value before the increment as a slot number. This way, each enqueue-dequeue pair
        works with an individual cell. We use the corresponding segment pointers to find the required ones.

        Here is a state machine for cells. Note that only one `acquire` and at most one `release` operation
        can deal with each cell, and that `release` uses `getAndSet(PERMIT)` to perform transitions for performance reasons
        so that the state `PERMIT` represents different logical states.

          +------+ `acquire` suspends   +------+   `release` tries    +--------+                    // if `cont.tryResume(..)` succeeds, then
          | NULL | -------------------> | cont | -------------------> | PERMIT | (cont RETRIEVED)   // the corresponding `acquire` operation gets
          +------+                      +------+   to resume `cont`   +--------+                    // a permit and the `release` one completes.
             |                             |
             |                             | `acquire` request is cancelled and the continuation is
             | `release` comes             | replaced with a special `CANCEL` token to avoid memory leaks
             | to the slot before          V
             | `acquire` and puts    +-----------+   `release` has    +--------+
             | a permit into the     | CANCELLED | -----------------> | PERMIT | (RElEASE FAILED)
             | slot, waiting for     +-----------+        failed      +--------+
             | `acquire` after
             | that.
             |
             |           `acquire` gets   +-------+
             |        +-----------------> | TAKEN | (ELIMINATION HAPPENED)
             V        |    the permit     +-------+
         +--------+   |
         | PERMIT | -<
         +--------+  |
                     |  `release` has waited a bounded time,   +--------+
                     +---------------------------------------> | BROKEN | (BOTH RELEASE AND ACQUIRE FAILED)
                            but `acquire` has not come         +--------+
     */

     private val head: AtomicRef<SemaphoreSegment>
     private val deqIdx = atomic(0L)
     private val tail: AtomicRef<SemaphoreSegment>
     private val enqIdx = atomic(0L)

     init {
         require(permits > 0) { "Semaphore should have at least 1 permit, but had $permits" }
         require(acquiredPermits in 0..permits) { "The number of acquired permits should be in 0..$permits" }
         val s = SemaphoreSegment(0, null, 2)
         head = atomic(s)
         tail = atomic(s)
     }

     /**
      * This counter indicates a number of available permits if it is non-negative,
      * or the size with minus sign otherwise. Note, that 32-bit counter is enough here
      * since the maximal number of available permits is [permits] which is [Int],
      * and the maximum number of waiting acquirers cannot be greater than 2^31 in any
      * real application.
      */
     private val _availablePermits = atomic(permits - acquiredPermits)
     override val availablePermits: Int get() = max(_availablePermits.value, 0)

     private val onCancellationRelease = { _: Throwable -> release() }

     override fun tryAcquire(): Boolean {
         _availablePermits.loop { p ->
             if (p <= 0) return false
             if (_availablePermits.compareAndSet(p, p - 1)) return true
         }
     }

     override suspend fun acquire() {
         val p = _availablePermits.getAndDecrement()
         if (p > 0) return // permit acquired
         // While it looks better when the following function is inlined,
         // it is important to make `suspend` function invocations in a way
         // so that the tail-call optimization can be applied.
         acquireSlowPath()
     }

     private suspend fun acquireSlowPath() = suspendCancellableCoroutineReusable<Unit> sc@ { cont ->
         while (true) {
             if (addAcquireToQueue(cont)) return@sc
             val p = _availablePermits.getAndDecrement()
             if (p > 0) { // permit acquired
                 cont.resume(Unit)
                 return@sc
             }
         }
     }

     override fun release() {
         while (true) {
             val p = _availablePermits.getAndUpdate { cur ->
                 check(cur < permits) { "The number of released permits cannot be greater than $permits" }
                 cur + 1
             }
             if (p >= 0) return
             if (tryResumeNextFromQueue()) return
         }
     }

     /**
      * Returns `false` if the received permit cannot be used and the calling operation should restart.
      */
     private fun addAcquireToQueue(cont: CancellableContinuation<Unit>): Boolean {
         val curTail = this.tail.value
         val enqIdx = enqIdx.getAndIncrement()
         val segment = this.tail.findSegmentAndMoveForward(id = enqIdx / SEGMENT_SIZE, startFrom = curTail,
             createNewSegment = ::createSegment).segment // cannot be closed
         val i = (enqIdx % SEGMENT_SIZE).toInt()
         // the regular (fast) path -- if the cell is empty, try to install continuation
         if (segment.cas(i, null, cont)) { // installed continuation successfully
             cont.invokeOnCancellation(CancelSemaphoreAcquisitionHandler(segment, i).asHandler)
             return true
         }
         // On CAS failure -- the cell must be either PERMIT or BROKEN
         // If the cell already has PERMIT from tryResumeNextFromQueue, try to grab it
         if (segment.cas(i, PERMIT, TAKEN)) { // took permit thus eliminating acquire/release pair
             // The following resume must always succeed, since continuation was not published yet and we don't have
             // to pass onCancellationRelease handle, since the coroutine did not suspend yet and cannot be cancelled
             cont.resume(Unit)
             return true
         }
         assert { segment.get(i) === BROKEN } // it must be broken in this case, no other way around it
         return false // broken cell, need to retry on a different cell
     }

     @Suppress("UNCHECKED_CAST")
     private fun tryResumeNextFromQueue(): Boolean {
         val curHead = this.head.value
         val deqIdx = deqIdx.getAndIncrement()
         val id = deqIdx / SEGMENT_SIZE
         val segment = this.head.findSegmentAndMoveForward(id, startFrom = curHead,
             createNewSegment = ::createSegment).segment // cannot be closed
         segment.cleanPrev()
         if (segment.id > id) return false
         val i = (deqIdx % SEGMENT_SIZE).toInt()
         val cellState = segment.getAndSet(i, PERMIT) // set PERMIT and retrieve the prev cell state
         when {
             cellState === null -> {
                 // Acquire has not touched this cell yet, wait until it comes for a bounded time
                 // The cell state can only transition from PERMIT to TAKEN by addAcquireToQueue
                 repeat(MAX_SPIN_CYCLES) {
                     if (segment.get(i) === TAKEN) return true
                 }
                 // Try to break the slot in order not to wait
                 return !segment.cas(i, PERMIT, BROKEN)
             }
             cellState === CANCELLED -> return false // the acquire was already cancelled
             else -> return (cellState as CancellableContinuation<Unit>).tryResumeAcquire()
         }
     }

     private fun CancellableContinuation<Unit>.tryResumeAcquire(): Boolean {
         val token = tryResume(Unit, null, onCancellationRelease) ?: return false
         completeResume(token)
         return true
     }
 }

 private class CancelSemaphoreAcquisitionHandler(
     private val segment: SemaphoreSegment,
     private val index: Int
 ) : CancelHandler() {
     override fun invoke(cause: Throwable?) {
         segment.cancel(index)
     }

     override fun toString() = "CancelSemaphoreAcquisitionHandler[$segment, $index]"
 }

 private fun createSegment(id: Long, prev: SemaphoreSegment?) = SemaphoreSegment(id, prev, 0)

 private class SemaphoreSegment(id: Long, prev: SemaphoreSegment?, pointers: Int) : Segment<SemaphoreSegment>(id, prev, pointers) {
     val acquirers = atomicArrayOfNulls<Any?>(SEGMENT_SIZE)
     override val maxSlots: Int get() = SEGMENT_SIZE

     @Suppress("NOTHING_TO_INLINE")
     inline fun get(index: Int): Any? = acquirers[index].value

     @Suppress("NOTHING_TO_INLINE")
     inline fun set(index: Int, value: Any?) {
         acquirers[index].value = value
     }

     @Suppress("NOTHING_TO_INLINE")
     inline fun cas(index: Int, expected: Any?, value: Any?): Boolean = acquirers[index].compareAndSet(expected, value)

     @Suppress("NOTHING_TO_INLINE")
     inline fun getAndSet(index: Int, value: Any?) = acquirers[index].getAndSet(value)

     // Cleans the acquirer slot located by the specified index
     // and removes this segment physically if all slots are cleaned.
     fun cancel(index: Int) {
         // Clean the slot
         set(index, CANCELLED)
         // Remove this segment if needed
         onSlotCleaned()
     }

     override fun toString() = "SemaphoreSegment[id=$id, hashCode=${hashCode()}]"
 }
 @SharedImmutable
 private val MAX_SPIN_CYCLES = systemProp("kotlinx.coroutines.semaphore.maxSpinCycles", 100)
 @SharedImmutable
 private val PERMIT = Symbol("PERMIT")
 @SharedImmutable
 private val TAKEN = Symbol("TAKEN")
 @SharedImmutable
 private val BROKEN = Symbol("BROKEN")
 @SharedImmutable
 private val CANCELLED = Symbol("CANCELLED")
 @SharedImmutable
 private val SEGMENT_SIZE = systemProp("kotlinx.coroutines.semaphore.segmentSize", 16)
	/*
	* Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
	*/

	package kotlinx.coroutines.sync

	import kotlinx.atomicfu.*
	import kotlinx.coroutines.*
	import kotlinx.coroutines.internal.*
	import kotlin.contracts.*
	import kotlin.coroutines.*
	import kotlin.math.*
	import kotlin.native.concurrent.SharedImmutable

	/**
	* A counting semaphore for coroutines that logically maintains a number of available permits.
	* Each [acquire] takes a single permit or suspends until it is available.
	* Each [release] adds a permit, potentially releasing a suspended acquirer.
	* Semaphore is fair and maintains a FIFO order of acquirers.
	*
	* Semaphores are mostly used to limit the number of coroutines that have an access to particular resource.
	* Semaphore with `permits = 1` is essentially a [Mutex].
	**/
	public interface Semaphore {
	/**
	* Returns the current number of permits available in this semaphore.
	*/
	public val availablePermits: Int

	/**
	* Acquires a permit from this semaphore, suspending until one is available.
	* All suspending acquirers are processed in first-in-first-out (FIFO) order.
	*
	* This suspending function is cancellable. If the [Job] of the current coroutine is cancelled or completed while this
	* function is suspended, this function immediately resumes with [CancellationException].
	* There is a prompt cancellation guarantee. If the job was cancelled while this function was
	* suspended, it will not resume successfully. See [suspendCancellableCoroutine] documentation for low-level details.
	* This function releases the semaphore if it was already acquired by this function before the [CancellationException]
	* was thrown.
	*
	* Note, that this function does not check for cancellation when it does not suspend.
	* Use [CoroutineScope.isActive] or [CoroutineScope.ensureActive] to periodically
	* check for cancellation in tight loops if needed.
	*
	* Use [tryAcquire] to try acquire a permit of this semaphore without suspension.
	*/
	public suspend fun acquire()

	/**
	* Tries to acquire a permit from this semaphore without suspension.
	*
	* @return `true` if a permit was acquired, `false` otherwise.
	*/
	public fun tryAcquire(): Boolean

	/**
	* Releases a permit, returning it into this semaphore. Resumes the first
	* suspending acquirer if there is one at the point of invocation.
	* Throws [IllegalStateException] if the number of [release] invocations is greater than the number of preceding [acquire].
	*/
	public fun release()
	}

	/**
	* Creates new [Semaphore] instance.
	* @param permits the number of permits available in this semaphore.
	* @param acquiredPermits the number of already acquired permits,
	* should be between `0` and `permits` (inclusively).
	*/
	@Suppress("FunctionName")
	public fun Semaphore(permits: Int, acquiredPermits: Int = 0): Semaphore = SemaphoreImpl(permits, acquiredPermits)

	/**
	* Executes the given [action], acquiring a permit from this semaphore at the beginning
	* and releasing it after the [action] is completed.
	*
	* @return the return value of the [action].
	*/
	@OptIn(ExperimentalContracts::class)
	public suspend inline fun <T> Semaphore.withPermit(action: () -> T): T {
	contract {
	callsInPlace(action, InvocationKind.EXACTLY_ONCE)
	}

	acquire()
	try {
	return action()
	} finally {
	release()
	}
	}

	private class SemaphoreImpl(private val permits: Int, acquiredPermits: Int) : Semaphore {
	/*
	The queue of waiting acquirers is essentially an infinite array based on the list of segments
	(see `SemaphoreSegment`); each segment contains a fixed number of slots. To determine a slot for each enqueue
	and dequeue operation, we increment the corresponding counter at the beginning of the operation
	and use the value before the increment as a slot number. This way, each enqueue-dequeue pair
	works with an individual cell. We use the corresponding segment pointers to find the required ones.

	Here is a state machine for cells. Note that only one `acquire` and at most one `release` operation
	can deal with each cell, and that `release` uses `getAndSet(PERMIT)` to perform transitions for performance reasons
	so that the state `PERMIT` represents different logical states.

	+------+ `acquire` suspends +------+ `release` tries +--------+ // if `cont.tryResume(..)` succeeds, then
	\| NULL \| -------------------> \| cont \| -------------------> \| PERMIT \| (cont RETRIEVED) // the corresponding `acquire` operation gets
	+------+ +------+ to resume `cont` +--------+ // a permit and the `release` one completes.
	\| \|
	\| \| `acquire` request is cancelled and the continuation is
	\| `release` comes \| replaced with a special `CANCEL` token to avoid memory leaks
	\| to the slot before V
	\| `acquire` and puts +-----------+ `release` has +--------+
	\| a permit into the \| CANCELLED \| -----------------> \| PERMIT \| (RElEASE FAILED)
	\| slot, waiting for +-----------+ failed +--------+
	\| `acquire` after
	\| that.
	\|
	\| `acquire` gets +-------+
	\| +-----------------> \| TAKEN \| (ELIMINATION HAPPENED)
	V \| the permit +-------+
	+--------+ \|
	\| PERMIT \| -<
	+--------+ \|
	\| `release` has waited a bounded time, +--------+
	+---------------------------------------> \| BROKEN \| (BOTH RELEASE AND ACQUIRE FAILED)
	but `acquire` has not come +--------+
	*/

	private val head: AtomicRef<SemaphoreSegment>
	private val deqIdx = atomic(0L)
	private val tail: AtomicRef<SemaphoreSegment>
	private val enqIdx = atomic(0L)

	init {
	require(permits > 0) { "Semaphore should have at least 1 permit, but had $permits" }
	require(acquiredPermits in 0..permits) { "The number of acquired permits should be in 0..$permits" }
	val s = SemaphoreSegment(0, null, 2)
	head = atomic(s)
	tail = atomic(s)
	}

	/**
	* This counter indicates a number of available permits if it is non-negative,
	* or the size with minus sign otherwise. Note, that 32-bit counter is enough here
	* since the maximal number of available permits is [permits] which is [Int],
	* and the maximum number of waiting acquirers cannot be greater than 2^31 in any
	* real application.
	*/
	private val _availablePermits = atomic(permits - acquiredPermits)
	override val availablePermits: Int get() = max(_availablePermits.value, 0)

	private val onCancellationRelease = { _: Throwable -> release() }

	override fun tryAcquire(): Boolean {
	_availablePermits.loop { p ->
	if (p <= 0) return false
	if (_availablePermits.compareAndSet(p, p - 1)) return true
	}
	}

	override suspend fun acquire() {
	val p = _availablePermits.getAndDecrement()
	if (p > 0) return // permit acquired
	// While it looks better when the following function is inlined,
	// it is important to make `suspend` function invocations in a way
	// so that the tail-call optimization can be applied.
	acquireSlowPath()
	}

	private suspend fun acquireSlowPath() = suspendCancellableCoroutineReusable<Unit> sc@ { cont ->
	while (true) {
	if (addAcquireToQueue(cont)) return@sc
	val p = _availablePermits.getAndDecrement()
	if (p > 0) { // permit acquired
	cont.resume(Unit)
	return@sc
	}
	}
	}

	override fun release() {
	while (true) {
	val p = _availablePermits.getAndUpdate { cur ->
	check(cur < permits) { "The number of released permits cannot be greater than $permits" }
	cur + 1
	}
	if (p >= 0) return
	if (tryResumeNextFromQueue()) return
	}
	}

	/**
	* Returns `false` if the received permit cannot be used and the calling operation should restart.
	*/
	private fun addAcquireToQueue(cont: CancellableContinuation<Unit>): Boolean {
	val curTail = this.tail.value
	val enqIdx = enqIdx.getAndIncrement()
	val segment = this.tail.findSegmentAndMoveForward(id = enqIdx / SEGMENT_SIZE, startFrom = curTail,
	createNewSegment = ::createSegment).segment // cannot be closed
	val i = (enqIdx % SEGMENT_SIZE).toInt()
	// the regular (fast) path -- if the cell is empty, try to install continuation
	if (segment.cas(i, null, cont)) { // installed continuation successfully
	cont.invokeOnCancellation(CancelSemaphoreAcquisitionHandler(segment, i).asHandler)
	return true
	}
	// On CAS failure -- the cell must be either PERMIT or BROKEN
	// If the cell already has PERMIT from tryResumeNextFromQueue, try to grab it
	if (segment.cas(i, PERMIT, TAKEN)) { // took permit thus eliminating acquire/release pair
	// The following resume must always succeed, since continuation was not published yet and we don't have
	// to pass onCancellationRelease handle, since the coroutine did not suspend yet and cannot be cancelled
	cont.resume(Unit)
	return true
	}
	assert { segment.get(i) === BROKEN } // it must be broken in this case, no other way around it
	return false // broken cell, need to retry on a different cell
	}

	@Suppress("UNCHECKED_CAST")
	private fun tryResumeNextFromQueue(): Boolean {
	val curHead = this.head.value
	val deqIdx = deqIdx.getAndIncrement()
	val id = deqIdx / SEGMENT_SIZE
	val segment = this.head.findSegmentAndMoveForward(id, startFrom = curHead,
	createNewSegment = ::createSegment).segment // cannot be closed
	segment.cleanPrev()
	if (segment.id > id) return false
	val i = (deqIdx % SEGMENT_SIZE).toInt()
	val cellState = segment.getAndSet(i, PERMIT) // set PERMIT and retrieve the prev cell state
	when {
	cellState === null -> {
	// Acquire has not touched this cell yet, wait until it comes for a bounded time
	// The cell state can only transition from PERMIT to TAKEN by addAcquireToQueue
	repeat(MAX_SPIN_CYCLES) {
	if (segment.get(i) === TAKEN) return true
	}
	// Try to break the slot in order not to wait
	return !segment.cas(i, PERMIT, BROKEN)
	}
	cellState === CANCELLED -> return false // the acquire was already cancelled
	else -> return (cellState as CancellableContinuation<Unit>).tryResumeAcquire()
	}
	}

	private fun CancellableContinuation<Unit>.tryResumeAcquire(): Boolean {
	val token = tryResume(Unit, null, onCancellationRelease) ?: return false
	completeResume(token)
	return true
	}
	}

	private class CancelSemaphoreAcquisitionHandler(
	private val segment: SemaphoreSegment,
	private val index: Int
	) : CancelHandler() {
	override fun invoke(cause: Throwable?) {
	segment.cancel(index)
	}

	override fun toString() = "CancelSemaphoreAcquisitionHandler[$segment, $index]"
	}

	private fun createSegment(id: Long, prev: SemaphoreSegment?) = SemaphoreSegment(id, prev, 0)

	private class SemaphoreSegment(id: Long, prev: SemaphoreSegment?, pointers: Int) : Segment<SemaphoreSegment>(id, prev, pointers) {
	val acquirers = atomicArrayOfNulls<Any?>(SEGMENT_SIZE)
	override val maxSlots: Int get() = SEGMENT_SIZE

	@Suppress("NOTHING_TO_INLINE")
	inline fun get(index: Int): Any? = acquirers[index].value

	@Suppress("NOTHING_TO_INLINE")
	inline fun set(index: Int, value: Any?) {
	acquirers[index].value = value
	}

	@Suppress("NOTHING_TO_INLINE")
	inline fun cas(index: Int, expected: Any?, value: Any?): Boolean = acquirers[index].compareAndSet(expected, value)

	@Suppress("NOTHING_TO_INLINE")
	inline fun getAndSet(index: Int, value: Any?) = acquirers[index].getAndSet(value)

	// Cleans the acquirer slot located by the specified index
	// and removes this segment physically if all slots are cleaned.
	fun cancel(index: Int) {
	// Clean the slot
	set(index, CANCELLED)
	// Remove this segment if needed
	onSlotCleaned()
	}

	override fun toString() = "SemaphoreSegment[id=$id, hashCode=${hashCode()}]"
	}
	@SharedImmutable
	private val MAX_SPIN_CYCLES = systemProp("kotlinx.coroutines.semaphore.maxSpinCycles", 100)
	@SharedImmutable
	private val PERMIT = Symbol("PERMIT")
	@SharedImmutable
	private val TAKEN = Symbol("TAKEN")
	@SharedImmutable
	private val BROKEN = Symbol("BROKEN")
	@SharedImmutable
	private val CANCELLED = Symbol("CANCELLED")
	@SharedImmutable
	private val SEGMENT_SIZE = systemProp("kotlinx.coroutines.semaphore.segmentSize", 16)