Minimize the API surface of LockFreeLinkedList

dkhalanskyjb · dkhalanskyjb · commit 33fdfa8537de · 2024-04-09T13:01:51.000+02:00
With this change, `JobSupport` uses only a small and well-defined
portion of the functionality `LockFreeLinkedList` provides, which
makes it easier to replace the list implementation.
diff --git a/kotlinx-coroutines-core/common/src/JobSupport.kt b/kotlinx-coroutines-core/common/src/JobSupport.kt
@@ -559,11 +559,7 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
 
     private fun promoteSingleToNodeList(state: JobNode) {
         // try to promote it to list (SINGLE+ state)
-        state.addOneIfEmpty(NodeList())
-        // it must be in SINGLE+ state or state has changed (node could have need removed from state)
-        val list = state.nextNode // either our NodeList or somebody else won the race, updated state
-        // just attempt converting it to list if state is still the same, then we'll continue lock-free loop
-        _state.compareAndSet(state, list)
+        _state.compareAndSet(state, state.attachToList(NodeList()))
     }
 
     public final override suspend fun join() {
@@ -951,37 +947,33 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
             handler = ChildCompletion(this, state, child, proposedUpdate)
         )
         if (handle !== NonDisposableHandle) return true // child is not complete and we've started waiting for it
-        val nextChild = child.nextChild() ?: return false
+        val nextChild = state.list.nextChild(startAfter = child) ?: return false
         return tryWaitForChild(state, nextChild, proposedUpdate)
     }
 
     // ## IMPORTANT INVARIANT: Only one thread can be concurrently invoking this method.
     private fun continueCompleting(state: Finishing, lastChild: ChildHandleNode, proposedUpdate: Any?) {
         assert { this.state === state } // consistency check -- it cannot change while we are waiting for children
         // figure out if we need to wait for next child
-        val waitChild = lastChild.nextChild()
+        val waitChild = state.list.nextChild(startAfter = lastChild)
         // try wait for next child
         if (waitChild != null && tryWaitForChild(state, waitChild, proposedUpdate)) return // waiting for next child
         // no more children to wait -- stop accepting children
         state.list.close(LIST_CHILD_PERMISSION)
         // did any children get added?
-        val waitChildAgain = lastChild.nextChild()
+        val waitChildAgain = state.list.nextChild(startAfter = lastChild)
         // try wait for next child
         if (waitChildAgain != null && tryWaitForChild(state, waitChildAgain, proposedUpdate)) return // waiting for next child
         // no more children, now we are sure; try to update the state
         val finalState = finalizeFinishingState(state, proposedUpdate)
         afterCompletion(finalState)
     }
 
-    private fun LockFreeLinkedListNode.nextChild(): ChildHandleNode? {
-        var cur = this
-        while (cur.isRemoved) cur = cur.prevNode // rollback to prev non-removed (or list head)
-        while (true) {
-            cur = cur.nextNode
-            if (cur.isRemoved) continue
-            if (cur is ChildHandleNode) return cur
-            if (cur is NodeList) return null // checked all -- no more children
+    private fun NodeList.nextChild(startAfter: LockFreeLinkedListNode? = null): ChildHandleNode? {
+        forEach(startAfter) {
+            if (it is ChildHandleNode) return it
         }
+        return null
     }
 
     public final override val children: Sequence<Job> get() = sequence {
@@ -1046,7 +1038,7 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
                  * as this child didn't make it before [notifyCancelling] and won't be notified that it should be
                  * cancelled.
                  *
-                 * And if the parent wasn't cancelled and the previous [LockFreeLinkedListNode.addLast] failed because
+                 * And if the parent wasn't cancelled and the previous [LockFreeLinkedListHead.addLast] failed because
                  * the job is in its final state already, we won't be able to attach anyway, so we must just invoke
                  * the handler and return.
                  */
diff --git a/kotlinx-coroutines-core/common/src/internal/LockFreeLinkedList.common.kt b/kotlinx-coroutines-core/common/src/internal/LockFreeLinkedList.common.kt
@@ -4,22 +4,38 @@ package kotlinx.coroutines.internal
 
 /** @suppress **This is unstable API and it is subject to change.** */
 public expect open class LockFreeLinkedListNode() {
-    public val isRemoved: Boolean
-    public val nextNode: LockFreeLinkedListNode
-    public val prevNode: LockFreeLinkedListNode
-    public fun addLast(node: LockFreeLinkedListNode, permissionsBitmask: Int): Boolean
-    public fun addOneIfEmpty(node: LockFreeLinkedListNode): Boolean
-    public open fun remove(): Boolean
+    /**
+     * Try putting `node` into a list.
+     *
+     * Returns:
+     * - The new head of the list if the operation succeeded.
+     * - The head of the list if someone else concurrently added this node to the list,
+     *   but no other modifications to the list were made.
+     * - Some garbage if the list was already edited.
+     */
+    public fun attachToList(node: LockFreeLinkedListHead): LockFreeLinkedListNode
+
+    /**
+     * Remove this node from the list.
+     */
+    public open fun remove()
+}
+
+/** @suppress **This is unstable API and it is subject to change.** */
+public expect open class LockFreeLinkedListHead() {
+    /**
+     * Iterates over all non-removed elements in this list, skipping every node until (and including) [startAfter].
+     */
+    public inline fun forEach(startAfter: LockFreeLinkedListNode? = null, block: (LockFreeLinkedListNode) -> Unit)
 
     /**
      * Closes the list for anything that requests the permission [forbiddenElementsBit].
      * Only a single permission can be forbidden at a time, but this isn't checked.
      */
     public fun close(forbiddenElementsBit: Int)
-}
 
-/** @suppress **This is unstable API and it is subject to change.** */
-public expect open class LockFreeLinkedListHead() : LockFreeLinkedListNode {
-    public inline fun forEach(block: (LockFreeLinkedListNode) -> Unit)
-    public final override fun remove(): Nothing
+    /**
+     * Adds the [node] to the end of the list if every bit in [permissionsBitmask] is still allowed in the list.
+     */
+    public fun addLast(node: LockFreeLinkedListNode, permissionsBitmask: Int): Boolean
 }
diff --git a/kotlinx-coroutines-core/concurrent/src/internal/LockFreeLinkedList.kt b/kotlinx-coroutines-core/concurrent/src/internal/LockFreeLinkedList.kt
@@ -37,21 +37,21 @@ public actual open class LockFreeLinkedListNode {
     private fun removed(): Removed =
         _removedRef.value ?: Removed(this).also { _removedRef.lazySet(it) }
 
-    public actual open val isRemoved: Boolean get() = next is Removed
+    public open val isRemoved: Boolean get() = next is Removed
 
     // LINEARIZABLE. Returns Node | Removed
     public val next: Any get() = _next.value
 
     // LINEARIZABLE. Returns next non-removed Node
-    public actual val nextNode: Node get() =
+    val nextNode: Node get() =
         next.let { (it as? Removed)?.ref ?: it as Node } // unwraps the `next` node
 
     // LINEARIZABLE WHEN THIS NODE IS NOT REMOVED:
     // Returns prev non-removed Node, makes sure prev is correct (prev.next === this)
     // NOTE: if this node is removed, then returns non-removed previous node without applying
     // prev.next correction, which does not provide linearizable backwards iteration, but can be used to
     // resume forward iteration when current node was removed.
-    public actual val prevNode: Node
+    public val prevNode: Node
         get() = correctPrev() ?: findPrevNonRemoved(_prev.value)
 
     private tailrec fun findPrevNonRemoved(current: Node): Node {
@@ -61,16 +61,16 @@ public actual open class LockFreeLinkedListNode {
 
     // ------ addOneIfEmpty ------
 
-    public actual fun addOneIfEmpty(node: Node): Boolean {
-        node._prev.lazySet(this)
-        node._next.lazySet(this)
+    public actual fun attachToList(node: LockFreeLinkedListHead): Node {
+        (node as Node)._prev.lazySet(this)
+        (node as Node)._next.lazySet(this)
         while (true) {
             val next = next
-            if (next !== this) return false // this is not an empty list!
+            if (next !== this) return nextNode // this is not an empty list!
             if (_next.compareAndSet(this, node)) {
                 // added successfully (linearized add) -- fixup the list
-                node.finishAdd(this)
-                return true
+                (node as Node).finishAdd(this)
+                return node
             }
         }
     }
@@ -80,7 +80,7 @@ public actual open class LockFreeLinkedListNode {
     /**
      * Adds last item to this list. Returns `false` if the list is closed.
      */
-    public actual fun addLast(node: Node, permissionsBitmask: Int): Boolean {
+    fun addLast(node: Node, permissionsBitmask: Int): Boolean {
         while (true) { // lock-free loop on prev.next
             val currentPrev = prevNode
             return when {
@@ -93,13 +93,6 @@ public actual open class LockFreeLinkedListNode {
         }
     }
 
-    /**
-     * Forbids adding new items to this list.
-     */
-    public actual fun close(forbiddenElementsBit: Int) {
-        addLast(ListClosed(forbiddenElementsBit), forbiddenElementsBit)
-    }
-
     /**
      * Given:
      * ```
@@ -142,8 +135,9 @@ public actual open class LockFreeLinkedListNode {
      * **Note**: Invocation of this operation does not guarantee that remove was actually complete if result was `false`.
      * In particular, invoking [nextNode].[prevNode] might still return this node even though it is "already removed".
      */
-    public actual open fun remove(): Boolean =
-        removeOrNext() == null
+    public actual open fun remove() {
+        removeOrNext()
+    }
 
     // returns null if removed successfully or next node if this node is already removed
     @PublishedApi
@@ -271,19 +265,30 @@ public actual open class LockFreeLinkedListHead : LockFreeLinkedListNode() {
     /**
      * Iterates over all elements in this list of a specified type.
      */
-    public actual inline fun forEach(block: (Node) -> Unit) {
-        var cur: Node = next as Node
-        while (cur != this) {
-            block(cur)
+    public actual inline fun forEach(startAfter: LockFreeLinkedListNode?, block: (Node) -> Unit) {
+        var cur: Node = startAfter ?: this
+        while (cur.isRemoved) cur = cur.prevNode // rollback to prev non-removed (or list head)
+        cur = cur.nextNode
+        while (cur !== this) {
+            if (!cur.isRemoved) {
+                block(cur)
+            }
             cur = cur.nextNode
         }
     }
 
     // just a defensive programming -- makes sure that list head sentinel is never removed
-    public actual final override fun remove(): Nothing = error("head cannot be removed")
+    public final override fun remove(): Nothing = error("head cannot be removed")
 
     // optimization: because head is never removed, we don't have to read _next.value to check these:
     override val isRemoved: Boolean get() = false
+
+    /**
+     * Forbids adding new items to this list.
+     */
+    public actual fun close(forbiddenElementsBit: Int) {
+        addLast(ListClosed(forbiddenElementsBit), forbiddenElementsBit)
+    }
 }
 
 private class ListClosed(val forbiddenElementsBitmask: Int): LockFreeLinkedListNode()
diff --git a/kotlinx-coroutines-core/jsAndWasmShared/src/internal/LinkedList.kt b/kotlinx-coroutines-core/jsAndWasmShared/src/internal/LinkedList.kt
@@ -2,6 +2,8 @@
 
 package kotlinx.coroutines.internal
 
+import kotlinx.coroutines.*
+
 private typealias Node = LockFreeLinkedListNode
 
 /** @suppress **This is unstable API and it is subject to change.** */
@@ -10,11 +12,7 @@ public actual open class LockFreeLinkedListNode {
     @PublishedApi internal var _prev = this
     @PublishedApi internal var _removed: Boolean = false
 
-    public actual inline val nextNode get() = _next
-    inline actual val prevNode get() = _prev
-    inline actual val isRemoved get() = _removed
-
-    public actual fun addLast(node: Node, permissionsBitmask: Int): Boolean = when (val prev = this._prev) {
+    fun addLast(node: Node, permissionsBitmask: Int): Boolean = when (val prev = this._prev) {
         is ListClosed ->
             prev.forbiddenElementsBitmask and permissionsBitmask == 0 && prev.addLast(node, permissionsBitmask)
         else -> {
@@ -26,30 +24,26 @@ public actual open class LockFreeLinkedListNode {
         }
     }
 
-    public actual fun close(forbiddenElementsBit: Int) {
-        addLast(ListClosed(forbiddenElementsBit), forbiddenElementsBit)
-    }
-
     /*
      * Remove that is invoked as a virtual function with a
      * potentially augmented behaviour.
      * I.g. `LockFreeLinkedListHead` throws, while `SendElementWithUndeliveredHandler`
      * invokes handler on remove
      */
-    public actual open fun remove(): Boolean {
-        if (_removed) return false
+    public actual open fun remove() {
+        if (_removed) return
         val prev = this._prev
         val next = this._next
         prev._next = next
         next._prev = prev
         _removed = true
-        return true
     }
 
-    public actual fun addOneIfEmpty(node: Node): Boolean {
-        if (_next !== this) return false
-        addLast(node, Int.MIN_VALUE)
-        return true
+    public actual fun attachToList(node: LockFreeLinkedListHead): Node {
+        if (_next !== this) return _next
+        val success = addLast(node, Int.MIN_VALUE)
+        assert { success }
+        return node
     }
 }
 
@@ -58,16 +52,24 @@ public actual open class LockFreeLinkedListHead : Node() {
     /**
      * Iterates over all elements in this list of a specified type.
      */
-    public actual inline fun forEach(block: (Node) -> Unit) {
-        var cur: Node = _next
-        while (cur != this) {
-            block(cur)
+    public actual inline fun forEach(startAfter: LockFreeLinkedListNode?, block: (Node) -> Unit) {
+        var cur: Node = startAfter ?: this
+        while (cur._removed) cur = cur._prev // rollback to prev non-removed (or list head)
+        cur = cur._next
+        while (cur !== this) {
+            if (!cur._removed) {
+                block(cur)
+            }
             cur = cur._next
         }
     }
 
+    public actual fun close(forbiddenElementsBit: Int) {
+        addLast(ListClosed(forbiddenElementsBit), forbiddenElementsBit)
+    }
+
     // just a defensive programming -- makes sure that list head sentinel is never removed
-    public actual final override fun remove(): Nothing = throw UnsupportedOperationException()
+    public final override fun remove(): Nothing = throw UnsupportedOperationException()
 }
 
 private class ListClosed(val forbiddenElementsBitmask: Int): LockFreeLinkedListNode()
diff --git a/kotlinx-coroutines-core/jsAndWasmShared/test/internal/LinkedListTest.kt b/kotlinx-coroutines-core/jsAndWasmShared/test/internal/LinkedListTest.kt
@@ -20,14 +20,13 @@ class LinkedListTest {
         assertContents(list, 1, 2, 3)
         val n4 = IntNode(4).apply { list.addLast(this, Int.MAX_VALUE) }
         assertContents(list, 1, 2, 3, 4)
-        assertTrue(n1.remove())
+        n1.remove()
         assertContents(list, 2, 3, 4)
-        assertTrue(n3.remove())
+        n3.remove()
         assertContents(list, 2, 4)
-        assertTrue(n4.remove())
+        n4.remove()
         assertContents(list, 2)
-        assertTrue(n2.remove())
-        assertFalse(n2.remove())
+        n2.remove()
         assertContents(list)
     }
 
