本文分享自天翼云开发者社区《python 实现消费者优先级队列》，作者：Frost 关键字 条件变量，信号量，消费者优先级，公平性，堆队列算法
需求背景 常见的是消息队列支持为消息指定优先级，但支持为消费者指定优先级的却很少见，作者在网上检索一般能查到 rabbitMQ 的消费者优先级相关资料。并没有找到其它语言的资料。
而 python 标准库里所有队列都是公平的，并没有使用非公平的参数，因此可能不能满足有些场景的需求。
什么是公平与非公平呢，这个一般是指互斥锁的特征，互斥锁的多个尝试取锁的线程其实很类似队列的多个消费者，以 waiter 统称。
假设有 A, B, C, D 四个 waiter，他们按照字母顺序依次调用 acquire()/get()，
那么等到有线程释放锁或队列放入了一条消息，会按照先来后到的顺序，唤醒对应的 waiter，也就是这里的 A，同理，按照排队顺序，B -> C -> D 将是后续的唤醒顺序，其实简单讲就是 FIFO。
一般来说 FIFO 策略具有普适性，可以避免有的消费者被饿死，但某些场景我们希望给队列的消费者赋予优先级，每次优先唤醒仍在等待消费的优先级最高的消费者。 下面会给出 pure python 的实现。
实现原理 先阅读 python 自带的 SimpleQueue 源码 (pure python 版本，位于 Lib\queue.py）。
1 '''Simple, unbounded FIFO queue. 2 3 This pure Python implementation is not reentrant. 4 ''' 5 # Note: while this pure Python version provides fairness 6 # (by using a threading.Semaphore which is itself fair, being based 7 # on threading.Condition), fairness is not part of the API contract. 8 # This allows the C version to use a different implementation. 9 10 def __init__(self): 11 self._queue = deque() 12 self._count = threading.Semaphore(0) 13 14 def put(self, item, block=True, timeout=None): 15 '''Put the item on the queue. 16 17 The optional 'block' and 'timeout' arguments are ignored, as this method 18 never blocks. They are provided for compatibility with the Queue class. 19 ''' 20 self._queue.append(item) 21 self._count.release() 22 23 def get(self, block=True, timeout=None): 24 '''Remove and return an item from the queue. 25 26 If optional args 'block' is true and 'timeout' is None (the default), 27 block if necessary until an item is available. If 'timeout' is 28 a non-negative number, it blocks at most 'timeout' seconds and raises 29 the Empty exception if no item was available within that time. 30 Otherwise ('block' is false), return an item if one is immediately 31 available, else raise the Empty exception ('timeout' is ignored 32 in that case). 33 ''' 34 if timeout is not None and timeout < 0: 35 raise ValueError("'timeout' must be a non-negative number") 36 if not self._count.acquire(block, timeout): 37 raise Empty 38 return self._queue.popleft() 39 40 def put_nowait(self, item): 41 '''Put an item into the queue without blocking. 42 43 This is exactly equivalent to `put(item, block=False)` and is only provided 44 for compatibility with the Queue class. 45 ''' 46 return self.put(item, block=False) 47 48 def get_nowait(self): 49 '''Remove and return an item from the queue without blocking. 50 51 Only get an item if one is immediately available. Otherwise 52 raise the Empty exception. 53 ''' 54 return self.get(block=False) 55 56 def empty(self): 57 '''Return True if the queue is empty, False otherwise (not reliable!).''' 58 return len(self._queue) == 0 59 60 def qsize(self): 61 '''Return the approximate size of the queue (not reliable!).''' 62 return len(self._queue) 63 64 __class_getitem__ = classmethod(types.GenericAlias)
docstring 里面说明，这个队列是保证了公平性，因为其使用的信号量实现是公平的。
符合直觉的是，我们在 get 方法以及信号量的 acquire 方法增加一个优先级数值的参数，那么再来看信号量的实现，看看能不能做到这一点，
1 """This class implements semaphore objects. 2 3 Semaphores manage a counter representing the number of release() calls minus 4 the number of acquire() calls, plus an initial value. The acquire() method 5 blocks if necessary until it can return without making the counter 6 negative. If not given, value defaults to 1. 7 8 """ 9 10 # After Tim Peters' semaphore class, but not quite the same (no maximum) 11 12 def __init__(self, value=1): 13 if value < 0: 14 raise ValueError("semaphore initial value must be >= 0") 15 self._cond = Condition(Lock()) 16 self._value = value 17 18 def acquire(self, blocking=True, timeout=None): 19 """Acquire a semaphore, decrementing the internal counter by one. 20 21 When invoked without arguments: if the internal counter is larger than 22 zero on entry, decrement it by one and return immediately. If it is zero 23 on entry, block, waiting until some other thread has called release() to 24 make it larger than zero. This is done with proper interlocking so that 25 if multiple acquire() calls are blocked, release() will wake exactly one 26 of them up. The implementation may pick one at random, so the order in 27 which blocked threads are awakened should not be relied on. There is no 28 return value in this case. 29 30 When invoked with blocking set to true, do the same thing as when called 31 without arguments, and return true. 32 33 When invoked with blocking set to false, do not block. If a call without 34 an argument would block, return false immediately; otherwise, do the 35 same thing as when called without arguments, and return true. 36 37 When invoked with a timeout other than None, it will block for at 38 most timeout seconds. If acquire does not complete successfully in 39 that interval, return false. Return true otherwise. 40 41 """ 42 if not blocking and timeout is not None: 43 raise ValueError("can't specify timeout for non-blocking acquire") 44 rc = False 45 endtime = None 46 with self._cond: 47 while self._value == 0: 48 if not blocking: 49 break 50 if timeout is not None: 51 if endtime is None: 52 endtime = _time() + timeout 53 else: 54 timeout = endtime - _time() 55 if timeout <= 0: 56 break 57 self._cond.wait(timeout) 58 else: 59 self._value -= 1 60 rc = True 61 return rc 62 63 __enter__ = acquire 64 65 def release(self, n=1): 66 """Release a semaphore, incrementing the internal counter by one or more. 67 68 When the counter is zero on entry and another thread is waiting for it 69 to become larger than zero again, wake up that thread. 70 71 """ 72 if n < 1: 73 raise ValueError('n must be one or more') 74 with self._cond: 75 self._value += n 76 for i in range(n): 77 self._cond.notify() 78 79 def __exit__(self, t, v, tb): 80 self.release()
可以看到信号量其实使用条件变量实现的，分析可知只能在条件变量的 wait 方法增加优先级数值。下面只展示不完整的，只有关键部分的源码
1 """Class that implements a condition variable. 2 3 A condition variable allows one or more threads to wait until they are 4 notified by another thread. 5 6 If the lock argument is given and not None, it must be a Lock or RLock 7 object, and it is used as the underlying lock. Otherwise, a new RLock object 8 is created and used as the underlying lock. 9 10 """ 11 12 def __init__(self, lock=None): 13 if lock is None: 14 lock = RLock() 15 self._lock = lock 16 # Export the lock's acquire() and release() methods 17 self.acquire = lock.acquire 18 self.release = lock.release 19 # If the lock defines _release_save() and/or _acquire_restore(), 20 # these override the default implementations (which just call 21 # release() and acquire() on the lock). Ditto for _is_owned(). 22 try: 23 self._release_save = lock._release_save 24 except AttributeError: 25 pass 26 try: 27 self._acquire_restore = lock._acquire_restore 28 except AttributeError: 29 pass 30 try: 31 self._is_owned = lock._is_owned 32 except AttributeError: 33 pass 34 self._waiters = _deque() 35 36 def _at_fork_reinit(self):... 37 def __enter__(self):... 38 def __exit__(self, *args):... 39 def __repr__(self):... 40 def _release_save(self):... 41 def _acquire_restore(self, x):... 42 def _is_owned(self):... 43 def wait(self, timeout=None): 44 """Wait until notified or until a timeout occurs. 45 46 If the calling thread has not acquired the lock when this method is 47 called, a RuntimeError is raised. 48 49 This method releases the underlying lock, and then blocks until it is 50 awakened by a notify() or notify_all() call for the same condition 51 variable in another thread, or until the optional timeout occurs. Once 52 awakened or timed out, it re-acquires the lock and returns. 53 54 When the timeout argument is present and not None, it should be a 55 floating point number specifying a timeout for the operation in seconds 56 (or fractions thereof). 57 58 When the underlying lock is an RLock, it is not released using its 59 release() method, since this may not actually unlock the lock when it 60 was acquired multiple times recursively. Instead, an internal interface 61 of the RLock class is used, which really unlocks it even when it has 62 been recursively acquired several times. Another internal interface is 63 then used to restore the recursion level when the lock is reacquired. 64 65 """ 66 if not self._is_owned(): 67 raise RuntimeError("cannot wait on un-acquired lock") 68 waiter = _allocate_lock() 69 waiter.acquire() 70 self._waiters.append(waiter) 71 saved_state = self._release_save() 72 gotit = False 73 try: # restore state no matter what (e.g., KeyboardInterrupt) 74 if timeout is None: 75 waiter.acquire() 76 gotit = True 77 else: 78 if timeout > 0: 79 gotit = waiter.acquire(True, timeout) 80 else: 81 gotit = waiter.acquire(False) 82 return gotit 83 finally: 84 self._acquire_restore(saved_state) 85 if not gotit: 86 try: 87 self._waiters.remove(waiter) 88 except ValueError: 89 pass 90 91 def wait_for(self, predicate, timeout=None):... 92 93 def notify(self, n=1): 94 """Wake up one or more threads waiting on this condition, if any. 95 96 If the calling thread has not acquired the lock when this method is 97 called, a RuntimeError is raised. 98 99 This method wakes up at most n of the threads waiting for the condition 100 variable; it is a no-op if no threads are waiting. 101 102 """ 103 if not self._is_owned(): 104 raise RuntimeError("cannot notify on un-acquired lock") 105 waiters = self._waiters 106 while waiters and n > 0: 107 waiter = waiters[0] 108 try: 109 waiter.release() 110 except RuntimeError: 111 # gh-92530: The previous call of notify() released the lock, 112 # but was interrupted before removing it from the queue. 113 # It can happen if a signal handler raises an exception, 114 # like CTRL+C which raises KeyboardInterrupt. 115 pass 116 else: 117 n -= 1 118 try: 119 waiters.remove(waiter) 120 except ValueError: 121 pass 122 123 def notify_all(self):... 124 125 def notifyAll(self):...
容易看出，关键部分在于 notify 方法，这里决定了通知哪一个 waiter，每个 waiter 对应队列的一个消费者。代码第 106 行的 waiters[0] 完美说明了 FIFO 原则，那么我们只需要改动这个 self._waiters 这个双向队列内 waiter 插入/获取的策略就行了。
可以使用标准库内 heapq 这个包来实现，也就是使用堆队列算法。另外 wait 方法要增加一个 priority 参数。具体实现如下：
1import itertools 2from dataclasses import dataclass, field 3from typing import Any, Tuple 4 5@dataclass(order=True) 6class PrioritizedWaiter: 7 priority: Tuple 8 waiter: Any=field(compare=False) 9 10class Condition: 11 """Class that implements a condition variable. 12 13 A condition variable allows one or more threads to wait until they are 14 notified by another thread. 15 16 If the lock argument is given and not None, it must be a Lock or RLock 17 object, and it is used as the underlying lock. Otherwise, a new RLock object 18 is created and used as the underlying lock. 19 20 """ 21 22 def __init__(self, lock=None): 23 if lock is None: 24 lock = RLock() 25 self._lock = lock 26 # Export the lock's acquire() and release() methods 27 self.acquire = lock.acquire 28 self.release = lock.release 29 # If the lock defines _release_save() and/or _acquire_restore(), 30 # these override the default implementations (which just call 31 # release() and acquire() on the lock). Ditto for _is_owned(). 32 try: 33 self._release_save = lock._release_save 34 except AttributeError: 35 pass 36 try: 37 self._acquire_restore = lock._acquire_restore 38 except AttributeError: 39 pass 40 try: 41 self._is_owned = lock._is_owned 42 except AttributeError: 43 pass 44 self._waiters = [] 45 self._counter = itertools.count() 46 47 def _at_fork_reinit(self):... 48 def __enter__(self):... 49 def __exit__(self, *args):... 50 def __repr__(self):... 51 def _release_save(self):... 52 def _acquire_restore(self, x):... 53 def _is_owned(self):... 54 def wait(self, timeout=None, priority=10): 55 """Wait until notified or until a timeout occurs. 56 57 If the calling thread has not acquired the lock when this method is 58 called, a RuntimeError is raised. 59 60 This method releases the underlying lock, and then blocks until it is 61 awakened by a notify() or notify_all() call for the same condition 62 variable in another thread, or until the optional timeout occurs. Once 63 awakened or timed out, it re-acquires the lock and returns. 64 65 When the timeout argument is present and not None, it should be a 66 floating point number specifying a timeout for the operation in seconds 67 (or fractions thereof). 68 69 When the underlying lock is an RLock, it is not released using its 70 release() method, since this may not actually unlock the lock when it 71 was acquired multiple times recursively. Instead, an internal interface 72 of the RLock class is used, which really unlocks it even when it has 73 been recursively acquired several times. Another internal interface is 74 then used to restore the recursion level when the lock is reacquired. 75 76 """ 77 if not self._is_owned(): 78 raise RuntimeError("cannot wait on un-acquired lock") 79 waiter = _allocate_lock() 80 waiter.acquire() 81 item = PrioritizedWaiter((priority, next(self._counter), waiter) 82 heapq.heappush(self._waiters, item) 83 saved_state = self._release_save() 84 gotit = False 85 try: # restore state no matter what (e.g., KeyboardInterrupt) 86 if timeout is None: 87 waiter.acquire() 88 gotit = True 89 else: 90 if timeout > 0: 91 gotit = waiter.acquire(True, timeout) 92 else: 93 gotit = waiter.acquire(False) 94 return gotit 95 finally: 96 self._acquire_restore(saved_state) 97 if not gotit: 98 try: 99 self._waiters.remove(item) 100 except ValueError: 101 pass 102 else: 103 heapq.heapify(self._waiters) 104 105 def wait_for(self, predicate, timeout=None):... 106 107 def notify(self, n=1): 108 """Wake up one or more threads waiting on this condition, if any. 109 110 If the calling thread has not acquired the lock when this method is 111 called, a RuntimeError is raised. 112 113 This method wakes up at most n of the threads waiting for the condition 114 variable; it is a no-op if no threads are waiting. 115 116 """ 117 if not self._is_owned(): 118 raise RuntimeError("cannot notify on un-acquired lock") 119 waiters = self._waiters 120 while waiters and n > 0: 121 waiter = waiters[0].waiter 122 try: 123 waiter.release() 124 except RuntimeError: 125 # gh-92530: The previous call of notify() released the lock, 126 # but was interrupted before removing it from the queue. 127 # It can happen if a signal handler raises an exception, 128 # like CTRL+C which raises KeyboardInterrupt. 129 pass 130 else: 131 n -= 1 132 try: 133 heapq.heappop(waiters) 134 except ValueError: 135 pass 136 137 def notify_all(self):... 138 def notifyAll(self):...
由此，最关键的实现完成了，接下来只需要给 _PySimpleQueue.get 方法也增加 priority 参数，并传入 Semaphore.acquire 方法。 Semaphore.acquire 方法增加 priority 参数，并传入给 Condition.wait 方法，就完成啦，限于篇幅这里就不全写下来了。
另外这里虽然加入了 priority 参数，但完全不使用这个参数时，其行为和原始版本时没有区别的，即依然符合 FIFO 策略。