|
|
- 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
|
