Unbounded threadsafe suspending priority queue
- Items are ordered by an integer priority value passed into the enqueue() method
- dequeue() retrieves and removes the head of the queue or suspends until an item is available
- enqueue(item, priority) inserts an item into the queue. suspends while aquiring a lock
Instead of using low level kotlin coroutine primitives to suspend when dequeue() is called on an empty queue, I create a new CompletableDeferred, add it to a queue of outstanding promises, then await it. When the next item is enqueued, a CompletableDeferred is popped from the queue and completed. The close() method cancels all outstanding promises
The class is used internally, so I didn't bother with the standard Queue interface.
Am I missing anything here that might cause a problem?
package com.something.something
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import java.util.*
import java.util.concurrent.atomic.AtomicBoolean
class AsyncQueueClosedError : RuntimeException("Queue is closed")
class AsyncPriorityQueue<T> : AutoCloseable {
val closed get() = isClosed.get()
val complete get() = isComplete.get() && items.size == 0
private val isClosed = AtomicBoolean(false)
private val isComplete = AtomicBoolean(false)
private val items = PriorityQueue<Prioritized<T>> { o1, o2 -> o1.priority.compareTo(o2.priority) }
private val lock = Mutex()
private val promises = ArrayDeque<CompletableDeferred<T>>()
override fun close() {
if (!isClosed.compareAndSet(false, true)) {
isComplete.set(true)
runBlocking {
lock.withLock {
for (p in promises) {
p.completeExceptionally(AsyncQueueClosedError())
}
promises.clear()
}
}
}
}
fun completeAdding() {
if (isComplete.compareAndSet(false, true)) {
if (items.size == 0) {
close()
}
}
}
private val key = object : Any() {}
suspend fun dequeue(): T {
lock.lock(key)
return try {
val item = items.poll()
if (item == null) {
if (isComplete.get()) {
close()
throw AsyncQueueClosedError()
}
val promise = CompletableDeferred<T>()
promises.add(promise)
lock.unlock(key)
promise.await()
} else {
lock.unlock(key)
item.item
}
} finally {
if (lock.holdsLock(key)) lock.unlock(key)
}
}
suspend fun enqueue(item: T, priority: Int) {
if (isComplete.get()) throw AsyncQueueClosedError()
lock.withLock {
promises.poll()
?.complete(item)
?: items.add(Prioritized(item, priority))
}
}
suspend fun size() = lock.withLock { items.size }
private class Prioritized<T> constructor(val item: T, val priority: Int)
}
```