• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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