2
\$\begingroup\$

Unbounded threadsafe suspending priority queue

  1. Items are ordered by an integer priority value passed into the enqueue() method
  2. dequeue() retrieves and removes the head of the queue or suspends until an item is available
  3. 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)
}
```
\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Browse other questions tagged or ask your own question.