Add documentation + small fixes

Author: ndkoval

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

@@ -3,35 +3,62 @@ package kotlinx.coroutines.internal
 import kotlinx.atomicfu.AtomicRef
 import kotlinx.atomicfu.atomic
 
-internal abstract class SegmentQueue<S: Segment<S>>(createFirstSegment: Boolean = true) {
+/**
+ * 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?>
-    protected val head: S? get() = _head.value
+    /**
+     * 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 tail: S? get() = _tail.value
+    protected val last: S? get() = _tail.value
 
     init {
-        val initialSegment = if (createFirstSegment) newSegment(0) else null
+        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 = head!!
+                startFrom = first!!
             }
         }
         if (startFrom.id > id) return null
-        // This method goes through `next` references and
-        // adds new segments if needed, similarly to the `push` in
-        // the Michael-Scott queue algorithm.
+        // 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
@@ -54,20 +81,23 @@ internal abstract class SegmentQueue<S: Segment<S>>(createFirstSegment: Boolean
         return cur
     }
 
-    protected fun getSegmentAndMoveHeadForward(startFrom: S?, id: Long): S? {
+    /**
+     * 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
+     * Updates [_head] to the specified segment
      * if its `id` is greater.
      */
     private fun moveHeadForward(new: S) {
         while (true) {
-            val cur = head!!
+            val cur = first!!
             if (cur.id > new.id) return
             if (_head.compareAndSet(cur, new)) {
                 new.prev.value = null
@@ -77,71 +107,62 @@ internal abstract class SegmentQueue<S: Segment<S>>(createFirstSegment: Boolean
     }
 
     /**
-     * Updates [tail] to the specified segment
+     * Updates [_tail] to the specified segment
      * if its `id` is greater.
      */
     private fun moveTailForward(new: S) {
         while (true) {
-            val cur = tail
+            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 segments, updates
-    // similarly to the Michael-Scott queue algorithm.
-    val next = atomic<S?>(null) // null (not set) | Segment | CLOSED
-    // Pointer to the previous non-empty segment (can be null!),
-    // updates lazily (see `remove()` function).
+    // 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 node from the waiting queue and cleans all references to it.
+     * 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 "}
-        val next = this.next.value ?: return // tail can't be removed
-        // Find the first non-removed node (tail is always non-removed)
+        // 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()
-
-//        while (next.removed) {
-//            next = next.next.value ?: return
-//        }
-//        // Find the first non-removed `prev` and remove this node
-//        var prev = prev.value
-//        while (true) {
-//            if (prev === null) {
-//                next.prev.value = null
-//                return
-//            }
-//            if (prev.removed) {
-//                prev = prev.prev.value
-//                continue
-//            }
-//            next.movePrevToLeft(prev)
-//            prev.movePrevNextToRight(next)
-//            if (next.removed || !prev.removed) return
-//            prev = prev.prev.value
-//        }
     }
 
     /**
-     * Update [Segment.next] pointer to the specified one if
-     * the `id` of the specified segment is greater.
+     * Updates [next] pointer to the specified segment if
+     * the [id] of the specified segment is greater.
      */
     private fun movePrevNextToRight(next: S) {
         while (true) {
@@ -152,8 +173,8 @@ internal abstract class Segment<S: Segment<S>>(val id: Long, prev: S?) {
     }
 
     /**
-     * Update [Segment.prev] pointer to the specified segment if
-     * its `id` is lower.
+     * Updates [prev] pointer to the specified segment if
+     * the [id] of the specified segment is lower.
      */
     private fun movePrevToLeft(prev: S) {
         while (true) {

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

@@ -3,7 +3,6 @@ package kotlinx.coroutines.sync
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.internal.*
-import kotlinx.coroutines.selects.select
 import kotlin.coroutines.resume
 import kotlin.jvm.JvmField
 import kotlin.math.max
@@ -74,24 +73,30 @@ public suspend inline fun <T> Semaphore.withSemaphore(action: () -> T): T {
 }
 
 private class SemaphoreImpl(@JvmField val maxPermits: Int, acquiredPermits: Int)
-    : Semaphore, SegmentQueue<SemaphoreSegment>(createFirstSegment = false)
+    : Semaphore, SegmentQueue<SemaphoreSegment>(createFirstSegmentLazily = true)
 {
     init {
         require(maxPermits > 0) { "Semaphore should have at least 1 permit"}
-        require(acquiredPermits in 0..maxPermits) { "TODO" }
+        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.toInt(), 0)
-
-    // Queue of waiting `acquire` requests. The queue is represented via enqueue and dequeue
-    // indices (`enqIdx` and `deqIdx` respectively) and infinite array. Each enqueue and dequeue
-    // operations increment the corresponding counter and goes to the corresponding unique cell.
-    // Due to the fact that we do not need to really send anything through this queue, we can make
-    // this queue wait-free. In order to do this,
+    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)
 
@@ -110,69 +115,69 @@ private class SemaphoreImpl(@JvmField val maxPermits: Int, acquiredPermits: Int)
 
     override fun release() {
         val p = _availablePermits.getAndUpdate { cur ->
-            check(cur < maxPermits) { "Cannot TODO" }
+            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 tail = this.tail
+        val last = this.last
         val enqIdx = enqIdx.getAndIncrement()
-        val segment = getSegment(tail, enqIdx / SEGMENT_SIZE)
+        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.clean(i)
+                segment.cancel(i)
                 release()
             }
         }.asHandler)
     }
 
     private fun resumeNextFromQueue() {
-        val head = this.head
+        val first = this.first
         val deqIdx = deqIdx.getAndIncrement()
-        val segment = getSegmentAndMoveHeadForward(head, deqIdx / SEGMENT_SIZE) ?: return
+        val segment = getSegmentAndMoveFirst(first, deqIdx / SEGMENT_SIZE) ?: return
         val i = (deqIdx % SEGMENT_SIZE).toInt()
         val cont = segment[i].getAndUpdate {
-            if (it === CLEANED) it else RESUMED
+            if (it === CANCELLED) CANCELLED else RESUMED
         }
-        if (cont === CLEANED) return
+        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) {
-    // == Waiters Array ==
-    private val waiters = atomicArrayOfNulls<Any?>(SEGMENT_SIZE)
+    private val acquirers = atomicArrayOfNulls<Any?>(SEGMENT_SIZE)
 
-    operator fun get(index: Int): AtomicRef<Any?> = waiters[index]
+    operator fun get(index: Int): AtomicRef<Any?> = acquirers[index]
 
-    private val cleaned = atomic(0)
-    override val removed get() = cleaned.value == SEGMENT_SIZE
+    private val cancelledSlots = atomic(0)
+    override val removed get() = cancelledSlots.value == SEGMENT_SIZE
 
-    /**
-     * Cleans the waiter located by the specified index in this segment.
-     */
-    fun clean(index: Int) {
-        // Clean the specified waiter and
-        // check if all node items are cleaned.
-        waiters[index].value = CLEANED
+    // 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 (cleaned.incrementAndGet() == SEGMENT_SIZE)
+        if (cancelledSlots.incrementAndGet() == SEGMENT_SIZE)
             remove()
     }
 }
 
 @SharedImmutable
 private val RESUMED = Symbol("RESUMED")
 @SharedImmutable
-private val CLEANED = Symbol("CLEANED")
+private val CANCELLED = Symbol("CANCELLED")
 @SharedImmutable
 private val SEGMENT_SIZE = systemProp("kotlinx.coroutines.semaphore.segmentSize", 32)