from collections import deque
|
|
|
|
class RoundRobinQueue(object):
|
|
"""A round robin queue implemented using multiple internal queues (typically,
|
|
FIFO queues). The internal queue must implement the following methods:
|
|
* push(obj)
|
|
* pop()
|
|
* close()
|
|
* __len__()
|
|
The constructor receives a qfactory argument, which is a callable used to
|
|
instantiate a new (internal) queue when a new key is allocated. The
|
|
qfactory function is called with the key number as first and only
|
|
argument.
|
|
start_keys is a sequence of domains to start with. If the queue was
|
|
previously closed leaving some domain buckets non-empty, those domains
|
|
should be passed in start_keys.
|
|
|
|
The queue maintains a fifo queue of keys. The key that went last is
|
|
poped first and the next queue for that key is then poped. This allows
|
|
for a round robin
|
|
"""
|
|
|
|
def __init__(self, qfactory, start_domains=()):
|
|
self.queues = {}
|
|
self.qfactory = qfactory
|
|
for key in start_domains:
|
|
self.queues[key] = self.qfactory(key)
|
|
|
|
self.key_queue = deque(start_domains)
|
|
|
|
def push(self, obj, key):
|
|
if key not in self.key_queue:
|
|
self.queues[key] = self.qfactory(key)
|
|
self.key_queue.appendleft(key) # it's new, might as well pop first
|
|
|
|
q = self.queues[key]
|
|
q.push(obj) # this may fail (eg. serialization error)
|
|
|
|
def pop(self):
|
|
m = None
|
|
# pop until we find a valid object, closing necessary queues
|
|
while m is None:
|
|
try:
|
|
key = self.key_queue.pop()
|
|
except IndexError:
|
|
return
|
|
|
|
q = self.queues[key]
|
|
m = q.pop()
|
|
|
|
if len(q) == 0:
|
|
del self.queues[key]
|
|
q.close()
|
|
else:
|
|
self.key_queue.appendleft(key)
|
|
|
|
if m:
|
|
return m
|
|
|
|
def close(self):
|
|
active = []
|
|
for k, q in self.queues.items():
|
|
if len(q):
|
|
active.append(k)
|
|
q.close()
|
|
return active
|
|
|
|
def __len__(self):
|
|
return sum(len(x) for x in self.queues.values()) if self.queues else 0
|