Skip to content

Commit

Permalink
Initial implementation of JobSupport on top of ConcurrentLinkedList
Browse files Browse the repository at this point in the history
  • Loading branch information
@dkhalanskyjb
dkhalanskyjb committed Apr 11, 2024
1 parent 33fdfa8 commit 1770738
Show file tree
Hide file tree
Showing 6 changed files with 224 additions and 545 deletions.
80 changes: 37 additions & 43 deletions kotlinx-coroutines-core/common/src/JobSupport.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import kotlinx.coroutines.internal.*
import kotlinx.coroutines.selects.*
import kotlin.coroutines.*
import kotlin.coroutines.intrinsics.*
import kotlin.experimental.*
import kotlin.js.*
import kotlin.jvm.*

Expand Down Expand Up @@ -319,8 +320,7 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
private fun notifyCancelling(list: NodeList, cause: Throwable) {
// first cancel our own children
onCancelling(cause)
list.close(LIST_CANCELLATION_PERMISSION)
notifyHandlers(list, cause) { it.onCancelling }
notifyHandlers(list, LIST_CANCELLATION_PERMISSION, cause) { it.onCancelling }
// then cancel parent
cancelParent(cause) // tentative cancellation -- does not matter if there is no parent
}
Expand Down Expand Up @@ -352,13 +352,12 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
}

private fun NodeList.notifyCompletion(cause: Throwable?) {
close(LIST_ON_COMPLETION_PERMISSION)
notifyHandlers(this, cause) { true }
notifyHandlers(this, LIST_ON_COMPLETION_PERMISSION, cause) { true }
}

private inline fun notifyHandlers(list: NodeList, cause: Throwable?, predicate: (JobNode) -> Boolean) {
private fun notifyHandlers(list: NodeList, permissionBitmask: Byte, cause: Throwable?, predicate: (JobNode) -> Boolean) {
var exception: Throwable? = null
list.forEach { node ->
list.forEach(forbidBitmask = permissionBitmask) { node ->
if (node is JobNode && predicate(node)) {
try {
node.invoke(cause)
Expand Down Expand Up @@ -925,57 +924,52 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
// process cancelling notification here -- it cancels all the children _before_ we start to wait them (sic!!!)
notifyRootCause?.let { notifyCancelling(list, it) }
// now wait for children
val child = list.nextChild()
if (child != null && tryWaitForChild(finishing, child, proposedUpdate))
return COMPLETING_WAITING_CHILDREN
list.close(LIST_CHILD_PERMISSION)
val anotherChild = list.nextChild()
if (anotherChild != null && tryWaitForChild(finishing, anotherChild, proposedUpdate))
return COMPLETING_WAITING_CHILDREN
if (shouldWaitForChildren(finishing, proposedUpdate)) return COMPLETING_WAITING_CHILDREN
// otherwise -- we have not children left (all were already cancelled?)
return finalizeFinishingState(finishing, proposedUpdate)
}

private val Any?.exceptionOrNull: Throwable?
get() = (this as? CompletedExceptionally)?.cause

// return false when there is no more incomplete children to wait
// ## IMPORTANT INVARIANT: Only one thread can be concurrently invoking this method.
private tailrec fun tryWaitForChild(state: Finishing, child: ChildHandleNode, proposedUpdate: Any?): Boolean {
val handle = child.childJob.invokeOnCompletion(
invokeImmediately = false,
handler = ChildCompletion(this, state, child, proposedUpdate)
)
if (handle !== NonDisposableHandle) return true // child is not complete and we've started waiting for it
val nextChild = state.list.nextChild(startAfter = child) ?: return false
return tryWaitForChild(state, nextChild, proposedUpdate)
private fun shouldWaitForChildren(state: Finishing, proposedUpdate: Any?, suggestedStart: ChildHandleNode? = null): Boolean {
val list = state.list
fun tryFindChildren(suggestedStart: ChildHandleNode?, closeList: Boolean): Boolean {
var startAfter: ChildHandleNode? = suggestedStart
while (true) {
val child = run {
list.forEach(forbidBitmask = if (closeList) LIST_CHILD_PERMISSION else 0, startAfter = startAfter) {
if (it is ChildHandleNode) return@run it
}
null
} ?: break
val handle = child.childJob.invokeOnCompletion(
invokeImmediately = false,
handler = ChildCompletion(this, state, child, proposedUpdate)
)
if (handle !== NonDisposableHandle) return true // child is not complete and we've started waiting for it
startAfter = child
}
return false
}
// Look for children that are currently in the list after the suggested start node.
if (tryFindChildren(suggestedStart = suggestedStart, closeList = false)) return true
// We didn't find anyone in the list after the suggested start node. Let's check the beginning now.
if (suggestedStart != null && tryFindChildren(suggestedStart = null, closeList = false)) return true
// Now we know that, at the moment this function started, there were no more children.
// We can close the list for the new children, and if we still don't find any, we can be sure there are none.
return tryFindChildren(suggestedStart = null, closeList = true)
}

// ## 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 = 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 = state.list.nextChild(startAfter = lastChild)
// try wait for next child
if (waitChildAgain != null && tryWaitForChild(state, waitChildAgain, proposedUpdate)) return // waiting for next child
if (shouldWaitForChildren(state, proposedUpdate, suggestedStart = lastChild)) return // waiting for the next child
// no more children, now we are sure; try to update the state
val finalState = finalizeFinishingState(state, proposedUpdate)
afterCompletion(finalState)
}

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 {
when (val state = this@JobSupport.state) {
is ChildHandleNode -> yield(state.childJob)
Expand Down Expand Up @@ -1389,9 +1383,9 @@ private val EMPTY_NEW = Empty(false)
private val EMPTY_ACTIVE = Empty(true)

// bit mask
private const val LIST_ON_COMPLETION_PERMISSION = 1
private const val LIST_CHILD_PERMISSION = 2
private const val LIST_CANCELLATION_PERMISSION = 4
private const val LIST_ON_COMPLETION_PERMISSION = 1.toByte()
private const val LIST_CHILD_PERMISSION = 2.toByte()
private const val LIST_CANCELLATION_PERMISSION = 4.toByte()

private class Empty(override val isActive: Boolean) : Incomplete {
override val list: NodeList? get() = null
Expand Down
198 changes: 187 additions & 11 deletions kotlinx-coroutines-core/common/src/internal/LockFreeLinkedList.common.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,40 +2,216 @@

package kotlinx.coroutines.internal

import kotlinx.atomicfu.*
import kotlinx.coroutines.*
import kotlin.coroutines.*
import kotlin.experimental.*
import kotlin.jvm.*

/** @suppress **This is unstable API and it is subject to change.** */
public expect open class LockFreeLinkedListNode() {
internal open class LockFreeLinkedListNode {
/**
* Try putting `node` into a list.
* Try putting this 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
fun attachToList(head: LockFreeLinkedListHead): LockFreeLinkedListHead {
val newAddress = head.addLastWithoutModifying(this, permissionsBitmask = 0)
assert { newAddress != null }
return if (_address.compareAndSet(null, newAddress)) {
head
} else {
_address.value!!.segment.head
}
}

/**
* Remove this node from the list.
*/
public open fun remove()
open fun remove() {
_address.value?.let {
val segment = it.segment
segment.clearSlot(it.index)
}
}

private val _address = atomic<Address?>(null)

val address: Address get() = _address.value!!

internal fun trySetAddress(address: Address) = this._address.compareAndSet(null, address)
}

/** @suppress **This is unstable API and it is subject to change.** */
public expect open class LockFreeLinkedListHead() {
internal open class LockFreeLinkedListHead {
private val head = LockFreeLinkedListSegment(
id = 0,
prev = null,
pointers = 2,
head = this,
)
private val tail = atomic(head)
private val nextElement = atomic(0L)

/**
* The list of bits that are forbidden from entering the list.
*
* TODO: we can store this in the extra bits in [head], there's enough space for that there, and it's never removed.
*/
private val forbiddenBits: AtomicInt = atomic(0)

/**
* 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)
inline fun forEach(
forbidBitmask: Byte = 0,
startAfter: LockFreeLinkedListNode? = null,
block: (LockFreeLinkedListNode) -> Unit
) {
forbiddenBits.update { it or forbidBitmask.toInt() }
val startAddress = startAfter?.address
var segment: LockFreeLinkedListSegment? = startAddress?.segment ?: head
var startIndex: Int = startAddress?.index?.let { it + 1 } ?: 0
while (segment != null) {
segment.forEach(forbidBitmask = forbidBitmask, startIndex = startIndex, block = block)
segment = segment.next
startIndex = 0
}
}

/**
* 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.
* Adds the [node] to the end of the list if every bit in [permissionsBitmask] is still allowed in the list,
* and then sets the [node]'s address to the new address.
*/
public fun close(forbiddenElementsBit: Int)
fun addLast(node: LockFreeLinkedListNode, permissionsBitmask: Byte): Boolean {
val address = addLastWithoutModifying(node, permissionsBitmask) ?: return false
val success = node.trySetAddress(address)
assert { success }
return true
}

/**
* Adds the [node] to the end of the list if every bit in [permissionsBitmask] is still allowed in the list.
* As opposed to [addLast], doesn't modify the [node]'s address.
*/
public fun addLast(node: LockFreeLinkedListNode, permissionsBitmask: Int): Boolean
fun addLastWithoutModifying(node: LockFreeLinkedListNode, permissionsBitmask: Byte): Address? {
/** First, avoid modifying the list at all if it was already closed for elements like ours. */
if (permissionsBitmask and forbiddenBits.value.toByte() != 0.toByte()) return null
/** Obtain the place from which the desired segment will certainly be reachable. */
val curTail = tail.value
/** Allocate a place for our element. */
val index = nextElement.getAndIncrement()
/** Find or create a segment where the node can be stored. */
val createNewSegment = ::createSegment // can't just pass the function, as the compiler crashes (KT-67332)
val segmentId = index / SEGMENT_SIZE
val segment = tail.findSegmentAndMoveForward(id = segmentId, curTail, createNewSegment).segment
assert { segment.id == segmentId }
val indexInSegment = (index % SEGMENT_SIZE).toInt()
/** Double-check that it's still not forbidden for the node to enter the list. */
if (permissionsBitmask and forbiddenBits.value.toByte() != 0.toByte()) return null
/** Now we know that the list was still not closed at some point *even after the segment* was created.
* Because [forbiddenBits] is set before [forEach] traverses the list, this means that [forEach] is guaranteed
* to observe the new segment and either break the cell where [node] wants to arrive or process the [node].
* In any case, we have linearizable behavior. */
return if (segment.tryAdd(node, permissionsBitmask = permissionsBitmask, indexInSegment = indexInSegment)) {
Address(segment, indexInSegment)
} else {
null
}
}
}

internal open class LockFreeLinkedListSegment(
id: Long,
prev: LockFreeLinkedListSegment?,
pointers: Int,
/** Used only during promoting of a single node to a list to ensure wait-freedom of the promotion operation.
* Without this, promotion can't be implemented without a (possibly bounded) spin loop: once the node is committed
* to be part of some list, the other threads can't do anything until that one thread sets the state to be the
* head of the list. */
@JvmField val head: LockFreeLinkedListHead,
) : Segment<LockFreeLinkedListSegment>(id = id, prev = prev, pointers = pointers)
{
/** Each cell is a [LockFreeLinkedListNode], a [BrokenForSomeElements], or `null`. */
private val cells = atomicArrayOfNulls<Any>(SEGMENT_SIZE)

override val numberOfSlots: Int get() = SEGMENT_SIZE

fun clearSlot(index: Int) {
cells[index].value = null
onSlotCleaned()
}

inline fun forEach(forbidBitmask: Byte, startIndex: Int, block: (LockFreeLinkedListNode) -> Unit) {
for (i in startIndex until SEGMENT_SIZE) {
val node = breakCellOrGetValue(forbidBitmask, i)
if (node != null) block(node)
}
}

private fun breakCellOrGetValue(forbidBitmask: Byte, index: Int): LockFreeLinkedListNode? {
while (true) {
val value = cells[index].value
if (value is BrokenForSomeElements?) {
val newForbiddenBits = value.forbiddenBits or forbidBitmask
if (newForbiddenBits == value.forbiddenBits
|| cells[index].compareAndSet(value, BrokenForSomeElements.fromBitmask(newForbiddenBits)))
return null
} else {
return value as LockFreeLinkedListNode
}
}
}

/**
* Adds the [node] to the array of cells if the slot wasn't broken.
*/
fun tryAdd(node: LockFreeLinkedListNode, permissionsBitmask: Byte, indexInSegment: Int): Boolean {
if (cells[indexInSegment].compareAndSet(null, node)) return true
cells[indexInSegment].loop { value ->
// This means that some elements are forbidden from entering the list.
value as BrokenForSomeElements
// Are *we* forbidden from entering the list?
if (value.forbiddenBits and permissionsBitmask != 0.toByte()) {
cells[indexInSegment].value = BrokenForSomeElements.FULLY_BROKEN
onSlotCleaned()
return false
}
// We aren't forbidden. Let's try entering it.
if (cells[indexInSegment].compareAndSet(value, node)) return true
}
}

override fun onCancellation(index: Int, cause: Throwable?, context: CoroutineContext) {
throw UnsupportedOperationException("Cancellation is not supported on LockFreeLinkedList")
}
}

internal class Address(@JvmField val segment: LockFreeLinkedListSegment, @JvmField val index: Int)

private fun createSegment(id: Long, prev: LockFreeLinkedListSegment): LockFreeLinkedListSegment =
LockFreeLinkedListSegment(
id = id,
prev = prev,
pointers = 0,
head = prev.head
)

private const val SEGMENT_SIZE = 8

@JvmInline
private value class BrokenForSomeElements private constructor(val forbiddenBits: Byte) {
companion object {
fun fromBitmask(forbiddenBits: Byte): BrokenForSomeElements? = when (forbiddenBits) {
0.toByte() -> null // no one is forbidden
else -> BrokenForSomeElements(forbiddenBits)
}

val FULLY_BROKEN = BrokenForSomeElements(255.toByte())
}
}

private val BrokenForSomeElements?.forbiddenBits get() = this?.forbiddenBits ?: 0

0 comments on commit 1770738

Please sign in to comment.