1 /* <lambda>null2 * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license. 3 */ 4 5 package kotlinx.coroutines.internal 6 7 import kotlinx.atomicfu.* 8 import kotlinx.coroutines.* 9 import kotlin.jvm.* 10 11 private typealias Core<E> = LockFreeTaskQueueCore<E> 12 13 /** 14 * Lock-free Multiply-Producer xxx-Consumer Queue for task scheduling purposes. 15 * 16 * **Note 1: This queue is NOT linearizable. It provides only quiescent consistency for its operations.** 17 * However, this guarantee is strong enough for task-scheduling purposes. 18 * In particular, the following execution is permitted for this queue, but is not permitted for a linearizable queue: 19 * 20 * ``` 21 * Thread 1: addLast(1) = true, removeFirstOrNull() = null 22 * Thread 2: addLast(2) = 2 // this operation is concurrent with both operations in the first thread 23 * ``` 24 * 25 * **Note 2: When this queue is used with multiple consumers (`singleConsumer == false`) this it is NOT lock-free.** 26 * In particular, consumer spins until producer finishes its operation in the case of near-empty queue. 27 * It is a very short window that could manifest itself rarely and only under specific load conditions, 28 * but it still deprives this algorithm of its lock-freedom. 29 */ 30 internal open class LockFreeTaskQueue<E : Any>( 31 singleConsumer: Boolean // true when there is only a single consumer (slightly faster & lock-free) 32 ) { 33 private val _cur = atomic(Core<E>(Core.INITIAL_CAPACITY, singleConsumer)) 34 35 // Note: it is not atomic w.r.t. remove operation (remove can transiently fail when isEmpty is false) 36 val isEmpty: Boolean get() = _cur.value.isEmpty 37 val size: Int get() = _cur.value.size 38 39 fun close() { 40 _cur.loop { cur -> 41 if (cur.close()) return // closed this copy 42 _cur.compareAndSet(cur, cur.next()) // move to next 43 } 44 } 45 46 fun addLast(element: E): Boolean { 47 _cur.loop { cur -> 48 when (cur.addLast(element)) { 49 Core.ADD_SUCCESS -> return true 50 Core.ADD_CLOSED -> return false 51 Core.ADD_FROZEN -> _cur.compareAndSet(cur, cur.next()) // move to next 52 } 53 } 54 } 55 56 @Suppress("UNCHECKED_CAST") 57 fun removeFirstOrNull(): E? { 58 _cur.loop { cur -> 59 val result = cur.removeFirstOrNull() 60 if (result !== Core.REMOVE_FROZEN) return result as E? 61 _cur.compareAndSet(cur, cur.next()) 62 } 63 } 64 65 // Used for validation in tests only 66 fun <R> map(transform: (E) -> R): List<R> = _cur.value.map(transform) 67 68 // Used for validation in tests only 69 fun isClosed(): Boolean = _cur.value.isClosed() 70 } 71 72 /** 73 * Lock-free Multiply-Producer xxx-Consumer Queue core. 74 * @see LockFreeTaskQueue 75 */ 76 internal class LockFreeTaskQueueCore<E : Any>( 77 private val capacity: Int, 78 private val singleConsumer: Boolean // true when there is only a single consumer (slightly faster) 79 ) { 80 private val mask = capacity - 1 81 private val _next = atomic<Core<E>?>(null) 82 private val _state = atomic(0L) 83 private val array = atomicArrayOfNulls<Any?>(capacity) 84 85 init { 86 check(mask <= MAX_CAPACITY_MASK) 87 check(capacity and mask == 0) 88 } 89 90 // Note: it is not atomic w.r.t. remove operation (remove can transiently fail when isEmpty is false) headnull91 val isEmpty: Boolean get() = _state.value.withState { head, tail -> head == tail } headnull92 val size: Int get() = _state.value.withState { head, tail -> (tail - head) and MAX_CAPACITY_MASK } 93 closenull94 fun close(): Boolean { 95 _state.update { state -> 96 if (state and CLOSED_MASK != 0L) return true // ok - already closed 97 if (state and FROZEN_MASK != 0L) return false // frozen -- try next 98 state or CLOSED_MASK // try set closed bit 99 } 100 return true 101 } 102 103 // ADD_CLOSED | ADD_FROZEN | ADD_SUCCESS addLastnull104 fun addLast(element: E): Int { 105 _state.loop { state -> 106 if (state and (FROZEN_MASK or CLOSED_MASK) != 0L) return state.addFailReason() // cannot add 107 state.withState { head, tail -> 108 val mask = this.mask // manually move instance field to local for performance 109 // If queue is Single-Consumer then there could be one element beyond head that we cannot overwrite, 110 // so we check for full queue with an extra margin of one element 111 if ((tail + 2) and mask == head and mask) return ADD_FROZEN // overfull, so do freeze & copy 112 // If queue is Multi-Consumer then the consumer could still have not cleared element 113 // despite the above check for one free slot. 114 if (!singleConsumer && array[tail and mask].value != null) { 115 // There are two options in this situation 116 // 1. Spin-wait until consumer clears the slot 117 // 2. Freeze & resize to avoid spinning 118 // We use heuristic here to avoid memory-overallocation 119 // Freeze & reallocate when queue is small or more than half of the queue is used 120 if (capacity < MIN_ADD_SPIN_CAPACITY || (tail - head) and MAX_CAPACITY_MASK > capacity shr 1) { 121 return ADD_FROZEN 122 } 123 // otherwise spin 124 return@loop 125 } 126 val newTail = (tail + 1) and MAX_CAPACITY_MASK 127 if (_state.compareAndSet(state, state.updateTail(newTail))) { 128 // successfully added 129 array[tail and mask].value = element 130 // could have been frozen & copied before this item was set -- correct it by filling placeholder 131 var cur = this 132 while(true) { 133 if (cur._state.value and FROZEN_MASK == 0L) break // all fine -- not frozen yet 134 cur = cur.next().fillPlaceholder(tail, element) ?: break 135 } 136 return ADD_SUCCESS // added successfully 137 } 138 } 139 } 140 } 141 fillPlaceholdernull142 private fun fillPlaceholder(index: Int, element: E): Core<E>? { 143 val old = array[index and mask].value 144 /* 145 * addLast actions: 146 * 1) Commit tail slot 147 * 2) Write element to array slot 148 * 3) Check for array copy 149 * 150 * If copy happened between 2 and 3 then the consumer might have consumed our element, 151 * then another producer might have written its placeholder in our slot, so we should 152 * perform *unique* check that current placeholder is our to avoid overwriting another producer placeholder 153 */ 154 if (old is Placeholder && old.index == index) { 155 array[index and mask].value = element 156 // we've corrected missing element, should check if that propagated to further copies, just in case 157 return this 158 } 159 // it is Ok, no need for further action 160 return null 161 } 162 163 // REMOVE_FROZEN | null (EMPTY) | E (SUCCESS) removeFirstOrNullnull164 fun removeFirstOrNull(): Any? { 165 _state.loop { state -> 166 if (state and FROZEN_MASK != 0L) return REMOVE_FROZEN // frozen -- cannot modify 167 state.withState { head, tail -> 168 if ((tail and mask) == (head and mask)) return null // empty 169 val element = array[head and mask].value 170 if (element == null) { 171 // If queue is Single-Consumer, then element == null only when add has not finished yet 172 if (singleConsumer) return null // consider it not added yet 173 // retry (spin) until consumer adds it 174 return@loop 175 } 176 // element == Placeholder can only be when add has not finished yet 177 if (element is Placeholder) return null // consider it not added yet 178 // we cannot put null into array here, because copying thread could replace it with Placeholder and that is a disaster 179 val newHead = (head + 1) and MAX_CAPACITY_MASK 180 if (_state.compareAndSet(state, state.updateHead(newHead))) { 181 // Array could have been copied by another thread and it is perfectly fine, since only elements 182 // between head and tail were copied and there are no extra steps we should take here 183 array[head and mask].value = null // now can safely put null (state was updated) 184 return element // successfully removed in fast-path 185 } 186 // Multi-Consumer queue must retry this loop on CAS failure (another consumer might have removed element) 187 if (!singleConsumer) return@loop 188 // Single-consumer queue goes to slow-path for remove in case of interference 189 var cur = this 190 while (true) { 191 @Suppress("UNUSED_VALUE") 192 cur = cur.removeSlowPath(head, newHead) ?: return element 193 } 194 } 195 } 196 } 197 removeSlowPathnull198 private fun removeSlowPath(oldHead: Int, newHead: Int): Core<E>? { 199 _state.loop { state -> 200 state.withState { head, _ -> 201 assert { head == oldHead } // "This queue can have only one consumer" 202 if (state and FROZEN_MASK != 0L) { 203 // state was already frozen, so removed element was copied to next 204 return next() // continue to correct head in next 205 } 206 if (_state.compareAndSet(state, state.updateHead(newHead))) { 207 array[head and mask].value = null // now can safely put null (state was updated) 208 return null 209 } 210 } 211 } 212 } 213 nextnull214 fun next(): LockFreeTaskQueueCore<E> = allocateOrGetNextCopy(markFrozen()) 215 216 private fun markFrozen(): Long = 217 _state.updateAndGet { state -> 218 if (state and FROZEN_MASK != 0L) return state // already marked 219 state or FROZEN_MASK 220 } 221 allocateOrGetNextCopynull222 private fun allocateOrGetNextCopy(state: Long): Core<E> { 223 _next.loop { next -> 224 if (next != null) return next // already allocated & copied 225 _next.compareAndSet(null, allocateNextCopy(state)) 226 } 227 } 228 allocateNextCopynull229 private fun allocateNextCopy(state: Long): Core<E> { 230 val next = LockFreeTaskQueueCore<E>(capacity * 2, singleConsumer) 231 state.withState { head, tail -> 232 var index = head 233 while (index and mask != tail and mask) { 234 // replace nulls with placeholders on copy 235 val value = array[index and mask].value ?: Placeholder(index) 236 next.array[index and next.mask].value = value 237 index++ 238 } 239 next._state.value = state wo FROZEN_MASK 240 } 241 return next 242 } 243 244 // Used for validation in tests only mapnull245 fun <R> map(transform: (E) -> R): List<R> { 246 val res = ArrayList<R>(capacity) 247 _state.value.withState { head, tail -> 248 var index = head 249 while (index and mask != tail and mask) { 250 // replace nulls with placeholders on copy 251 val element = array[index and mask].value 252 @Suppress("UNCHECKED_CAST") 253 if (element != null && element !is Placeholder) res.add(transform(element as E)) 254 index++ 255 } 256 } 257 return res 258 } 259 260 // Used for validation in tests only isClosednull261 fun isClosed(): Boolean = _state.value and CLOSED_MASK != 0L 262 263 264 // Instance of this class is placed into array when we have to copy array, but addLast is in progress -- 265 // it had already reserved a slot in the array (with null) and have not yet put its value there. 266 // Placeholder keeps the actual index (not masked) to distinguish placeholders on different wraparounds of array 267 // Internal because of inlining 268 internal class Placeholder(@JvmField val index: Int) 269 270 @Suppress("PrivatePropertyName", "MemberVisibilityCanBePrivate") 271 internal companion object { 272 const val INITIAL_CAPACITY = 8 273 274 const val CAPACITY_BITS = 30 275 const val MAX_CAPACITY_MASK = (1 shl CAPACITY_BITS) - 1 276 const val HEAD_SHIFT = 0 277 const val HEAD_MASK = MAX_CAPACITY_MASK.toLong() shl HEAD_SHIFT 278 const val TAIL_SHIFT = HEAD_SHIFT + CAPACITY_BITS 279 const val TAIL_MASK = MAX_CAPACITY_MASK.toLong() shl TAIL_SHIFT 280 281 const val FROZEN_SHIFT = TAIL_SHIFT + CAPACITY_BITS 282 const val FROZEN_MASK = 1L shl FROZEN_SHIFT 283 const val CLOSED_SHIFT = FROZEN_SHIFT + 1 284 const val CLOSED_MASK = 1L shl CLOSED_SHIFT 285 286 const val MIN_ADD_SPIN_CAPACITY = 1024 287 288 @JvmField val REMOVE_FROZEN = Symbol("REMOVE_FROZEN") 289 290 const val ADD_SUCCESS = 0 291 const val ADD_FROZEN = 1 292 const val ADD_CLOSED = 2 293 294 infix fun Long.wo(other: Long) = this and other.inv() 295 fun Long.updateHead(newHead: Int) = (this wo HEAD_MASK) or (newHead.toLong() shl HEAD_SHIFT) 296 fun Long.updateTail(newTail: Int) = (this wo TAIL_MASK) or (newTail.toLong() shl TAIL_SHIFT) 297 298 inline fun <T> Long.withState(block: (head: Int, tail: Int) -> T): T { 299 val head = ((this and HEAD_MASK) shr HEAD_SHIFT).toInt() 300 val tail = ((this and TAIL_MASK) shr TAIL_SHIFT).toInt() 301 return block(head, tail) 302 } 303 304 // FROZEN | CLOSED 305 fun Long.addFailReason(): Int = if (this and CLOSED_MASK != 0L) ADD_CLOSED else ADD_FROZEN 306 } 307 } 308