core/kotlinx-coroutines-core/src/internal/LockFreeMPSCQueue.kt - platform/external/kotlinx.coroutines - Gitiles

 /*
  * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */

 package kotlinx.coroutines.experimental.internal

 import kotlinx.atomicfu.*
 import java.util.concurrent.atomic.*

 private typealias Core<E> = LockFreeMPSCQueueCore<E>

 /**
  * Lock-free Multiply-Producer Single-Consumer Queue.
  * *Note: This queue is NOT linearizable. It provides only quiescent consistency for its operations.*
  *
  * In particular, the following execution is permitted for this queue, but is not permitted for a linearizable queue:
  *
  * ```
  * Thread 1: addLast(1) = true, removeFirstOrNull() = null
  * Thread 2: addLast(2) = 2 // this operation is concurrent with both operations in the first thread
  * ```
  *
  * @suppress **This is unstable API and it is subject to change.**
  */
 class LockFreeMPSCQueue<E : Any> {
     private val _cur = atomic(Core<E>(Core.INITIAL_CAPACITY))

     // Note: it is not atomic w.r.t. remove operation (remove can transiently fail when isEmpty is false)
     val isEmpty: Boolean get() = _cur.value.isEmpty

     fun close() {
         _cur.loop { cur ->
             if (cur.close()) return // closed this copy
             _cur.compareAndSet(cur, cur.next()) // move to next
         }
     }

     fun addLast(element: E): Boolean {
         _cur.loop { cur ->
             when (cur.addLast(element)) {
                 Core.ADD_SUCCESS -> return true
                 Core.ADD_CLOSED -> return false
                 Core.ADD_FROZEN -> _cur.compareAndSet(cur, cur.next()) // move to next
             }
         }
     }

     @Suppress("UNCHECKED_CAST")
     fun removeFirstOrNull(): E? {
         _cur.loop { cur ->
             val result = cur.removeFirstOrNull()
             if (result !== Core.REMOVE_FROZEN) return result as E?
             _cur.compareAndSet(cur, cur.next())
         }
     }
 }

 /**
  * Lock-free Multiply-Producer Single-Consumer Queue core.
  * *Note: This queue is NOT linearizable. It provides only quiescent consistency for its operations.*
  *
  * @see LockFreeMPSCQueue
  * @suppress **This is unstable API and it is subject to change.**
  */
 internal class LockFreeMPSCQueueCore<E : Any>(private val capacity: Int) {
     private val mask = capacity - 1
     private val _next = atomic<Core<E>?>(null)
     private val _state = atomic(0L)
     private val array = AtomicReferenceArray<Any?>(capacity);

     init {
         check(mask <= MAX_CAPACITY_MASK)
         check(capacity and mask == 0)
     }

     // Note: it is not atomic w.r.t. remove operation (remove can transiently fail when isEmpty is false)
     val isEmpty: Boolean get() = _state.value.withState { head, tail -> head == tail }

     fun close(): Boolean {
         _state.update { state ->
             if (state and CLOSED_MASK != 0L) return true // ok - already closed
             if (state and FROZEN_MASK != 0L) return false // frozen -- try next
             state or CLOSED_MASK // try set closed bit
         }
         return true
     }

     // ADD_CLOSED | ADD_FROZEN | ADD_SUCCESS
     fun addLast(element: E): Int {
         _state.loop { state ->
             if (state and (FROZEN_MASK or CLOSED_MASK) != 0L) return state.addFailReason() // cannot add
             state.withState { head, tail ->
                 // there could be one REMOVE element beyond head that we cannot stump up,
                 // so we check for full queue with an extra margin of one element
                 if ((tail + 2) and mask == head and mask) return ADD_FROZEN // overfull, so do freeze & copy
                 val newTail = (tail + 1) and MAX_CAPACITY_MASK
                 if (_state.compareAndSet(state, state.updateTail(newTail))) {
                     // successfully added
                     array[tail and mask] = element
                     // could have been frozen & copied before this item was set -- correct it by filling placeholder
                     var cur = this
                     while(true) {
                         if (cur._state.value and FROZEN_MASK == 0L) break // all fine -- not frozen yet
                         cur = cur.next().fillPlaceholder(tail, element) ?: break
                     }
                     return ADD_SUCCESS // added successfully
                 }
             }
         }
     }

     private fun fillPlaceholder(index: Int, element: E): Core<E>? {
         val old = array.get(index and mask)
         /*
          * addLast actions:
          * 1) Commit tail slot
          * 2) Write element to array slot
          * 3) Check for array copy
          *
          * If copy happened between 2 and 3 then the consumer might have consumed our element,
          * then another producer might have written its placeholder in our slot, so we should
          * perform *unique* check that current placeholder is our to avoid overwriting another producer placeholder
          */
         if (old is Placeholder && old.index == index) {
             array.set(index and mask, element)
             // we've corrected missing element, should check if that propagated to further copies, just in case
             return this
         }
         // it is Ok, no need for further action
         return null
     }

     // SINGLE CONSUMER
     // REMOVE_FROZEN | null (EMPTY) | E (SUCCESS)
     fun removeFirstOrNull(): Any? {
         _state.loop { state ->
             if (state and FROZEN_MASK != 0L) return REMOVE_FROZEN // frozen -- cannot modify
             state.withState { head, tail ->
                 if ((tail and mask) == (head and mask)) return null // empty
                 // because queue is Single Consumer, then element == null|Placeholder can only be when add has not finished yet
                 val element = array[head and mask] ?: return null
                 if (element is Placeholder) return null // same story -- consider it not added yet
                 // we cannot put null into array here, because copying thread could replace it with Placeholder and that is a disaster
                 val newHead = (head + 1) and MAX_CAPACITY_MASK
                 if (_state.compareAndSet(state, state.updateHead(newHead))) {
                     array[head and mask] = null // now can safely put null (state was updated)
                     return element // successfully removed in fast-path
                 }
                 // Slow-path for remove in case of interference
                 var cur = this
                 while (true) {
                     cur = cur.removeSlowPath(head, newHead) ?: return element
                 }
             }
         }
     }

     private fun removeSlowPath(oldHead: Int, newHead: Int): Core<E>? {
         _state.loop { state ->
             state.withState { head, _ ->
                 check(head == oldHead) { "This queue can have only one consumer" }
                 if (state and FROZEN_MASK != 0L) {
                     // state was already frozen, so removed element was copied to next
                     return next() // continue to correct head in next
                 }
                 if (_state.compareAndSet(state, state.updateHead(newHead))) {
                     array[head and mask] = null // now can safely put null (state was updated)
                     return null
                 }
             }
         }
     }

     fun next(): LockFreeMPSCQueueCore<E> = allocateOrGetNextCopy(markFrozen())

     private fun markFrozen(): Long =
         _state.updateAndGet { state ->
             if (state and FROZEN_MASK != 0L) return state // already marked
             state or FROZEN_MASK
         }

     private fun allocateOrGetNextCopy(state: Long): Core<E> {
         _next.loop { next ->
             if (next != null) return next // already allocated & copied
             _next.compareAndSet(null, allocateNextCopy(state))
         }
     }

     private fun allocateNextCopy(state: Long): Core<E> {
         val next = LockFreeMPSCQueueCore<E>(capacity * 2)
         state.withState { head, tail ->
             var index = head
             while (index and mask != tail and mask) {
                 // replace nulls with placeholders on copy
                 next.array[index and next.mask] = array[index and mask] ?: Placeholder(index)
                 index++
             }
             next._state.value = state wo FROZEN_MASK
         }
         return next
     }

     // Instance of this class is placed into array when we have to copy array, but addLast is in progress --
     // it had already reserved a slot in the array (with null) and have not yet put its value there.
     // Placeholder keeps the actual index (not masked) to distinguish placeholders on different wraparounds of array
     private class Placeholder(@JvmField val index: Int)

     @Suppress("PrivatePropertyName")
     internal companion object {
         internal const val INITIAL_CAPACITY = 8

         private const val CAPACITY_BITS = 30
         private const val MAX_CAPACITY_MASK = (1 shl CAPACITY_BITS) - 1
         private const val HEAD_SHIFT = 0
         private const val HEAD_MASK = MAX_CAPACITY_MASK.toLong() shl HEAD_SHIFT
         private const val TAIL_SHIFT = HEAD_SHIFT + CAPACITY_BITS
         private const val TAIL_MASK = MAX_CAPACITY_MASK.toLong() shl TAIL_SHIFT

         private const val FROZEN_SHIFT = TAIL_SHIFT + CAPACITY_BITS
         private const val FROZEN_MASK = 1L shl FROZEN_SHIFT
         private const val CLOSED_SHIFT = FROZEN_SHIFT + 1
         private const val CLOSED_MASK = 1L shl CLOSED_SHIFT

         @JvmField internal val REMOVE_FROZEN = Symbol("REMOVE_FROZEN")

         internal const val ADD_SUCCESS = 0
         internal const val ADD_FROZEN = 1
         internal const val ADD_CLOSED = 2

         private infix fun Long.wo(other: Long) = this and other.inv()
         private fun Long.updateHead(newHead: Int) = (this wo HEAD_MASK) or (newHead.toLong() shl HEAD_SHIFT)
         private fun Long.updateTail(newTail: Int) = (this wo TAIL_MASK) or (newTail.toLong() shl TAIL_SHIFT)

         private inline fun <T> Long.withState(block: (head: Int, tail: Int) -> T): T {
             val head = ((this and HEAD_MASK) shr HEAD_SHIFT).toInt()
             val tail = ((this and TAIL_MASK) shr TAIL_SHIFT).toInt()
             return block(head, tail)
         }

         // FROZEN | CLOSED
         private fun Long.addFailReason(): Int = if (this and CLOSED_MASK != 0L) ADD_CLOSED else ADD_FROZEN
     }
 }
	/*
	* Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
	*/

	package kotlinx.coroutines.experimental.internal

	import kotlinx.atomicfu.*
	import java.util.concurrent.atomic.*

	private typealias Core<E> = LockFreeMPSCQueueCore<E>

	/**
	* Lock-free Multiply-Producer Single-Consumer Queue.
	* Note: This queue is NOT linearizable. It provides only quiescent consistency for its operations.
	*
	* In particular, the following execution is permitted for this queue, but is not permitted for a linearizable queue:
	*
	* ```
	* Thread 1: addLast(1) = true, removeFirstOrNull() = null
	* Thread 2: addLast(2) = 2 // this operation is concurrent with both operations in the first thread
	* ```
	*
	* @suppress This is unstable API and it is subject to change.
	*/
	class LockFreeMPSCQueue<E : Any> {
	private val _cur = atomic(Core<E>(Core.INITIAL_CAPACITY))

	// Note: it is not atomic w.r.t. remove operation (remove can transiently fail when isEmpty is false)
	val isEmpty: Boolean get() = _cur.value.isEmpty

	fun close() {
	_cur.loop { cur ->
	if (cur.close()) return // closed this copy
	_cur.compareAndSet(cur, cur.next()) // move to next
	}
	}

	fun addLast(element: E): Boolean {
	_cur.loop { cur ->
	when (cur.addLast(element)) {
	Core.ADD_SUCCESS -> return true
	Core.ADD_CLOSED -> return false
	Core.ADD_FROZEN -> _cur.compareAndSet(cur, cur.next()) // move to next
	}
	}
	}

	@Suppress("UNCHECKED_CAST")
	fun removeFirstOrNull(): E? {
	_cur.loop { cur ->
	val result = cur.removeFirstOrNull()
	if (result !== Core.REMOVE_FROZEN) return result as E?
	_cur.compareAndSet(cur, cur.next())
	}
	}
	}

	/**
	* Lock-free Multiply-Producer Single-Consumer Queue core.
	* Note: This queue is NOT linearizable. It provides only quiescent consistency for its operations.
	*
	* @see LockFreeMPSCQueue
	* @suppress This is unstable API and it is subject to change.
	*/
	internal class LockFreeMPSCQueueCore<E : Any>(private val capacity: Int) {
	private val mask = capacity - 1
	private val _next = atomic<Core<E>?>(null)
	private val _state = atomic(0L)
	private val array = AtomicReferenceArray<Any?>(capacity);

	init {
	check(mask <= MAX_CAPACITY_MASK)
	check(capacity and mask == 0)
	}

	// Note: it is not atomic w.r.t. remove operation (remove can transiently fail when isEmpty is false)
	val isEmpty: Boolean get() = _state.value.withState { head, tail -> head == tail }

	fun close(): Boolean {
	_state.update { state ->
	if (state and CLOSED_MASK != 0L) return true // ok - already closed
	if (state and FROZEN_MASK != 0L) return false // frozen -- try next
	state or CLOSED_MASK // try set closed bit
	}
	return true
	}

	// ADD_CLOSED \| ADD_FROZEN \| ADD_SUCCESS
	fun addLast(element: E): Int {
	_state.loop { state ->
	if (state and (FROZEN_MASK or CLOSED_MASK) != 0L) return state.addFailReason() // cannot add
	state.withState { head, tail ->
	// there could be one REMOVE element beyond head that we cannot stump up,
	// so we check for full queue with an extra margin of one element
	if ((tail + 2) and mask == head and mask) return ADD_FROZEN // overfull, so do freeze & copy
	val newTail = (tail + 1) and MAX_CAPACITY_MASK
	if (_state.compareAndSet(state, state.updateTail(newTail))) {
	// successfully added
	array[tail and mask] = element
	// could have been frozen & copied before this item was set -- correct it by filling placeholder
	var cur = this
	while(true) {
	if (cur._state.value and FROZEN_MASK == 0L) break // all fine -- not frozen yet
	cur = cur.next().fillPlaceholder(tail, element) ?: break
	}
	return ADD_SUCCESS // added successfully
	}
	}
	}
	}

	private fun fillPlaceholder(index: Int, element: E): Core<E>? {
	val old = array.get(index and mask)
	/*
	* addLast actions:
	* 1) Commit tail slot
	* 2) Write element to array slot
	* 3) Check for array copy
	*
	* If copy happened between 2 and 3 then the consumer might have consumed our element,
	* then another producer might have written its placeholder in our slot, so we should
	* perform unique check that current placeholder is our to avoid overwriting another producer placeholder
	*/
	if (old is Placeholder && old.index == index) {
	array.set(index and mask, element)
	// we've corrected missing element, should check if that propagated to further copies, just in case
	return this
	}
	// it is Ok, no need for further action
	return null
	}

	// SINGLE CONSUMER
	// REMOVE_FROZEN \| null (EMPTY) \| E (SUCCESS)
	fun removeFirstOrNull(): Any? {
	_state.loop { state ->
	if (state and FROZEN_MASK != 0L) return REMOVE_FROZEN // frozen -- cannot modify
	state.withState { head, tail ->
	if ((tail and mask) == (head and mask)) return null // empty
	// because queue is Single Consumer, then element == null\|Placeholder can only be when add has not finished yet
	val element = array[head and mask] ?: return null
	if (element is Placeholder) return null // same story -- consider it not added yet
	// we cannot put null into array here, because copying thread could replace it with Placeholder and that is a disaster
	val newHead = (head + 1) and MAX_CAPACITY_MASK
	if (_state.compareAndSet(state, state.updateHead(newHead))) {
	array[head and mask] = null // now can safely put null (state was updated)
	return element // successfully removed in fast-path
	}
	// Slow-path for remove in case of interference
	var cur = this
	while (true) {
	cur = cur.removeSlowPath(head, newHead) ?: return element
	}
	}
	}
	}

	private fun removeSlowPath(oldHead: Int, newHead: Int): Core<E>? {
	_state.loop { state ->
	state.withState { head, _ ->
	check(head == oldHead) { "This queue can have only one consumer" }
	if (state and FROZEN_MASK != 0L) {
	// state was already frozen, so removed element was copied to next
	return next() // continue to correct head in next
	}
	if (_state.compareAndSet(state, state.updateHead(newHead))) {
	array[head and mask] = null // now can safely put null (state was updated)
	return null
	}
	}
	}
	}

	fun next(): LockFreeMPSCQueueCore<E> = allocateOrGetNextCopy(markFrozen())

	private fun markFrozen(): Long =
	_state.updateAndGet { state ->
	if (state and FROZEN_MASK != 0L) return state // already marked
	state or FROZEN_MASK
	}

	private fun allocateOrGetNextCopy(state: Long): Core<E> {
	_next.loop { next ->
	if (next != null) return next // already allocated & copied
	_next.compareAndSet(null, allocateNextCopy(state))
	}
	}

	private fun allocateNextCopy(state: Long): Core<E> {
	val next = LockFreeMPSCQueueCore<E>(capacity * 2)
	state.withState { head, tail ->
	var index = head
	while (index and mask != tail and mask) {
	// replace nulls with placeholders on copy
	next.array[index and next.mask] = array[index and mask] ?: Placeholder(index)
	index++
	}
	next._state.value = state wo FROZEN_MASK
	}
	return next
	}

	// Instance of this class is placed into array when we have to copy array, but addLast is in progress --
	// it had already reserved a slot in the array (with null) and have not yet put its value there.
	// Placeholder keeps the actual index (not masked) to distinguish placeholders on different wraparounds of array
	private class Placeholder(@JvmField val index: Int)

	@Suppress("PrivatePropertyName")
	internal companion object {
	internal const val INITIAL_CAPACITY = 8

	private const val CAPACITY_BITS = 30
	private const val MAX_CAPACITY_MASK = (1 shl CAPACITY_BITS) - 1
	private const val HEAD_SHIFT = 0
	private const val HEAD_MASK = MAX_CAPACITY_MASK.toLong() shl HEAD_SHIFT
	private const val TAIL_SHIFT = HEAD_SHIFT + CAPACITY_BITS
	private const val TAIL_MASK = MAX_CAPACITY_MASK.toLong() shl TAIL_SHIFT

	private const val FROZEN_SHIFT = TAIL_SHIFT + CAPACITY_BITS
	private const val FROZEN_MASK = 1L shl FROZEN_SHIFT
	private const val CLOSED_SHIFT = FROZEN_SHIFT + 1
	private const val CLOSED_MASK = 1L shl CLOSED_SHIFT

	@JvmField internal val REMOVE_FROZEN = Symbol("REMOVE_FROZEN")

	internal const val ADD_SUCCESS = 0
	internal const val ADD_FROZEN = 1
	internal const val ADD_CLOSED = 2

	private infix fun Long.wo(other: Long) = this and other.inv()
	private fun Long.updateHead(newHead: Int) = (this wo HEAD_MASK) or (newHead.toLong() shl HEAD_SHIFT)
	private fun Long.updateTail(newTail: Int) = (this wo TAIL_MASK) or (newTail.toLong() shl TAIL_SHIFT)

	private inline fun <T> Long.withState(block: (head: Int, tail: Int) -> T): T {
	val head = ((this and HEAD_MASK) shr HEAD_SHIFT).toInt()
	val tail = ((this and TAIL_MASK) shr TAIL_SHIFT).toInt()
	return block(head, tail)
	}

	// FROZEN \| CLOSED
	private fun Long.addFailReason(): Int = if (this and CLOSED_MASK != 0L) ADD_CLOSED else ADD_FROZEN
	}
	}