Add Semaphore.

In addition, a `SegmentQueue` data structure, which emulates an infinite array with fast removing from the middle, is introduced for storing suspended acquirers in the semaphore algorithm.

Fixes #1088

Author: ndkoval

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt

@@ -981,6 +981,19 @@ public final class kotlinx/coroutines/sync/MutexKt {
 	public static synthetic fun withLock$default (Lkotlinx/coroutines/sync/Mutex;Ljava/lang/Object;Lkotlin/jvm/functions/Function0;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 }
 
+public abstract interface class kotlinx/coroutines/sync/Semaphore {
+	public abstract fun acquire (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public abstract fun getAvailablePermits ()I
+	public abstract fun release ()V
+	public abstract fun tryAcquire ()Z
+}
+
+public final class kotlinx/coroutines/sync/SemaphoreKt {
+	public static final fun Semaphore (II)Lkotlinx/coroutines/sync/Semaphore;
+	public static synthetic fun Semaphore$default (IIILjava/lang/Object;)Lkotlinx/coroutines/sync/Semaphore;
+	public static final fun withSemaphore (Lkotlinx/coroutines/sync/Semaphore;Lkotlin/jvm/functions/Function0;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+}
+
 public final class kotlinx/coroutines/test/TestCoroutineContext : kotlin/coroutines/CoroutineContext {
 	public fun <init> ()V
 	public fun <init> (Ljava/lang/String;)V

kotlinx-coroutines-core/common/src/internal/SegmentQueue.kt

@@ -0,0 +1,186 @@
+package kotlinx.coroutines.internal
+
+import kotlinx.atomicfu.AtomicRef
+import kotlinx.atomicfu.atomic
+
+/**
+ * Essentially, this segment queue is an infinite array of segments, which is represented as
+ * a Michael-Scott queue of them. All segments are instances of [Segment] interface and
+ * follow in natural order (see [Segment.id]) in the queue.
+ *
+ * In some data structures, like `Semaphore`, this queue is used for storing suspended continuations
+ * and is always empty in uncontended scenarios. Therefore, there is no need in creating
+ * the first segment in advance in this case. A special `createFirstSegmentLazily` is introduced
+ * to create segments lazily, on the first [getSegment] invocation; it is set to `false` by default.
+ */
+internal abstract class SegmentQueue<S: Segment<S>>(createFirstSegmentLazily: Boolean = false) {
+    private val _head: AtomicRef<S?>
+    /**
+     * Returns the first segment in the queue. All segments with lower [id]
+     */
+    protected val first: S? get() = _head.value
+
+    private val _tail: AtomicRef<S?>
+    protected val last: S? get() = _tail.value
+
+    init {
+        val initialSegment = if (createFirstSegmentLazily) null else newSegment(0)
+        _head = atomic(initialSegment)
+        _tail = atomic(initialSegment)
+    }
+
+    /**
+     * The implementation should create an instance of segment [S] with the specified id
+     * and initial reference to the previous one.
+     */
+    abstract fun newSegment(id: Long, prev: S? = null): S
+
+    /**
+     * Finds a segment with the specified [id] following by next references from the
+     * [startFrom] segment. The typical use-case is reading [last] or [first], doing some
+     * synchronization, and invoking [getSegment] or [getSegmentAndMoveFirst] correspondingly
+     * to find the required segment.
+     */
+    protected fun getSegment(startFrom: S?, id: Long): S? {
+        // Try to create the first segment if [startFrom] is null.
+        // This occurs if `createFirstSegmentLazily` was set to `true`.
+        var startFrom = startFrom
+        if (startFrom === null) {
+            val firstSegment = newSegment(0)
+            if (_head.compareAndSet(null, firstSegment))
+                startFrom = firstSegment
+            else {
+                startFrom = first!!
+            }
+        }
+        if (startFrom.id > id) return null
+        // Go through `next` references and add new segments if needed,
+        // similarly to the `push` in the Michael-Scott queue algorithm.
+        // The only difference is that `CAS failure` means that the
+        // required segment has already been added, so the algorithm just
+        // uses it. This way, only one segment with each id can be in the queue.
+        var cur: S = startFrom
+        while (cur.id < id) {
+            var curNext = cur.next.value
+            if (curNext == null) {
+                // Add a new segment.
+                val newTail = newSegment(cur.id + 1, cur)
+                curNext = if (cur.next.compareAndSet(null, newTail)) {
+                    if (cur.removed) {
+                        cur.remove()
+                    }
+                    moveTailForward(newTail)
+                    newTail
+                } else {
+                    cur.next.value!!
+                }
+            }
+            cur = curNext
+        }
+        if (cur.id != id) return null
+        return cur
+    }
+
+    /**
+     * Invokes [getSegment] and replaces [first] with the result if its [id] is greater.
+     */
+    protected fun getSegmentAndMoveFirst(startFrom: S?, id: Long): S? {
+        if (startFrom !== null && startFrom.id == id) return startFrom
+        val s = getSegment(startFrom, id) ?: return null
+        moveHeadForward(s)
+        return s
+    }
+
+    /**
+     * Updates [_head] to the specified segment
+     * if its `id` is greater.
+     */
+    private fun moveHeadForward(new: S) {
+        while (true) {
+            val cur = first!!
+            if (cur.id > new.id) return
+            if (_head.compareAndSet(cur, new)) {
+                new.prev.value = null
+                return
+            }
+        }
+    }
+
+    /**
+     * Updates [_tail] to the specified segment
+     * if its `id` is greater.
+     */
+    private fun moveTailForward(new: S) {
+        while (true) {
+            val cur = last
+            if (cur !== null && cur.id > new.id) return
+            if (_tail.compareAndSet(cur, new)) return
+        }
+    }
+}
+
+/**
+ * Each segment in [SegmentQueue] has a unique id and is created by [SegmentQueue.newSegment].
+ * Essentially, this is a node in the Michael-Scott queue algorithm, but with
+ * maintaining [prev] pointer for efficient [remove] implementation.
+ */
+internal abstract class Segment<S: Segment<S>>(val id: Long, prev: S?) {
+    // Pointer to the next segment, updates similarly to the Michael-Scott queue algorithm.
+    val next = atomic<S?>(null)
+    // Pointer to the previous segment, updates in [remove] function.
+    val prev = atomic<S?>(null)
+
+    /**
+     * Returns `true` if this segment is logically removed from the queue.
+     * The [remove] function should be called right after it becomes logically removed.
+     */
+    abstract val removed: Boolean
+
+    init {
+        this.prev.value = prev
+    }
+
+    /**
+     * Removes this segment physically from the segment queue. The segment should be
+     * logically removed (so [removed] returns `true`) at the point of invocation.
+     */
+    fun remove() {
+        check(removed) { " The segment should be logically removed at first "}
+        // Read `next` and `prev` pointers.
+        val next = this.next.value ?: return // tail cannot be removed
+        val prev = prev.value ?: return // head cannot be removed
+        // Link `next` and `prev`.
+        next.movePrevToLeft(prev)
+        prev.movePrevNextToRight(next)
+        // Check whether `prev` and `next` are still in the queue
+        // and help with removing them if needed.
+        if (prev.removed)
+            prev.remove()
+        if (next.removed)
+            next.remove()
+    }
+
+    /**
+     * Updates [next] pointer to the specified segment if
+     * the [id] of the specified segment is greater.
+     */
+    private fun movePrevNextToRight(next: S) {
+        while (true) {
+            val curNext = this.next.value as S
+            if (next.id <= curNext.id) return
+            if (this.next.compareAndSet(curNext, next)) return
+        }
+    }
+
+    /**
+     * Updates [prev] pointer to the specified segment if
+     * the [id] of the specified segment is lower.
+     */
+    private fun movePrevToLeft(prev: S) {
+        while (true) {
+            val curPrev = this.prev.value ?: return
+            if (curPrev.id <= prev.id) return
+            if (this.prev.compareAndSet(curPrev, prev)) return
+        }
+    }
+}

kotlinx-coroutines-core/common/src/sync/Semaphore.kt

@@ -0,0 +1,183 @@
+package kotlinx.coroutines.sync
+
+import kotlinx.atomicfu.*
+import kotlinx.coroutines.*
+import kotlinx.coroutines.internal.*
+import kotlin.coroutines.resume
+import kotlin.jvm.JvmField
+import kotlin.math.max
+
+/**
+ * A counting semaphore for coroutines. It maintains a number of available permits.
+ * Each [acquire] suspends if necessary until a permit is available, and then takes it.
+ * Each [release] adds a permit, potentially releasing a suspended acquirer.
+ *
+ * Semaphore with `maxPermits = 1` is essentially a [Mutex].
+ **/
+public interface Semaphore {
+    /**
+     * Returns the current number of available 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].
+     *
+     * *Cancellation of suspended lock invocation is atomic* -- when this function
+     * throws [CancellationException] it means that the mutex was not locked.
+     *
+     * Note, that this function does not check for cancellation when it is not suspended.
+     * Use [yield] or [CoroutineScope.isActive] 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.
+     */
+    public fun release()
+}
+
+/**
+ * Creates new [Semaphore] instance.
+ */
+@Suppress("FunctionName")
+public fun Semaphore(maxPermits: Int, acquiredPermits: Int = 0): Semaphore = SemaphoreImpl(maxPermits, acquiredPermits)
+
+/**
+ * Executes the given [action] with acquiring a permit from this semaphore at the beginning
+ * and releasing it after the [action] is completed.
+ *
+ * @return the return value of the [action].
+ */
+public suspend inline fun <T> Semaphore.withSemaphore(action: () -> T): T {
+    acquire()
+    try {
+        return action()
+    } finally {
+        release()
+    }
+}
+
+private class SemaphoreImpl(@JvmField val maxPermits: Int, acquiredPermits: Int)
+    : Semaphore, SegmentQueue<SemaphoreSegment>(createFirstSegmentLazily = true)
+{
+    init {
+        require(maxPermits > 0) { "Semaphore should have at least 1 permit"}
+        require(acquiredPermits in 0..maxPermits) { "The number of acquired permits should be ≥ 0 and ≤ maxPermits" }
+    }
+
+    override fun newSegment(id: Long, prev: SemaphoreSegment?)= SemaphoreSegment(id, prev)
+
+    /**
+     * 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 [maxPermits] which is [Int],
+     * and the maximum number of waiting acquirers cannot be greater than 2^31 in any
+     * real application.
+     */
+    private val _availablePermits = atomic(maxPermits)
+    override val availablePermits: Int get() = max(_availablePermits.value, 0)
+
+    // The queue of waiting acquirers is essentially an infinite array based on `SegmentQueue`;
+    // 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.
+    private val enqIdx = atomic(0L)
+    private val deqIdx = atomic(0L)
+
+    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
+        addToQueueAndSuspend()
+    }
+
+    override fun release() {
+        val p = _availablePermits.getAndUpdate { cur ->
+            check(cur < maxPermits) { "The number of acquired permits cannot be greater than maxPermits" }
+            cur + 1
+        }
+        if (p >= 0) return // no waiters
+        resumeNextFromQueue()
+    }
+
+    private suspend fun addToQueueAndSuspend() = suspendAtomicCancellableCoroutine<Unit> sc@ { cont ->
+        val last = this.last
+        val enqIdx = enqIdx.getAndIncrement()
+        val segment = getSegment(last, enqIdx / SEGMENT_SIZE)
+        val i = (enqIdx % SEGMENT_SIZE).toInt()
+        if (segment === null || segment[i].value === RESUMED || !segment[i].compareAndSet(null, cont)) {
+            // already resumed
+            cont.resume(Unit)
+            return@sc
+        }
+        cont.invokeOnCancellation(handler = object : CancelHandler() {
+            override fun invoke(cause: Throwable?) {
+                segment.cancel(i)
+                release()
+            }
+        }.asHandler)
+    }
+
+    private fun resumeNextFromQueue() {
+        val first = this.first
+        val deqIdx = deqIdx.getAndIncrement()
+        val segment = getSegmentAndMoveFirst(first, deqIdx / SEGMENT_SIZE) ?: return
+        val i = (deqIdx % SEGMENT_SIZE).toInt()
+        val cont = segment[i].getAndUpdate {
+            if (it === CANCELLED) CANCELLED else RESUMED
+        }
+        if (cont === null) return // just resumed
+        if (cont === CANCELLED) return // Cancelled continuation invokes `release`
+                                       // and resumes next suspended acquirer if needed.
+        cont as CancellableContinuation<Unit>
+        cont.resume(Unit)
+    }
+}
+
+private class SemaphoreSegment(id: Long, prev: SemaphoreSegment?): Segment<SemaphoreSegment>(id, prev) {
+    private val acquirers = atomicArrayOfNulls<Any?>(SEGMENT_SIZE)
+
+    operator fun get(index: Int): AtomicRef<Any?> = acquirers[index]
+
+    private val cancelledSlots = atomic(0)
+    override val removed get() = cancelledSlots.value == SEGMENT_SIZE
+
+    // 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 specified waiter
+        acquirers[index].value = CANCELLED
+        // Remove this segment if needed
+        if (cancelledSlots.incrementAndGet() == SEGMENT_SIZE)
+            remove()
+    }
+}
+
+@SharedImmutable
+private val RESUMED = Symbol("RESUMED")
+@SharedImmutable
+private val CANCELLED = Symbol("CANCELLED")
+@SharedImmutable
+private val SEGMENT_SIZE = systemProp("kotlinx.coroutines.semaphore.segmentSize", 32)