def __init__(self, qkey=None): self.__lastput__ = -1 self.__lastget__ = -1 self.kqmap = {} self.__qkey__ = qkey or "anonymous" self.queues = Proxy(self._getqueues) self.immediate = PriorityQueue() self.mutex = threading.Lock() self.not_empty = threading.Condition(self.mutex)
class KGroupQueueAbstract(object): def __init__(self, qkey=None): self.__lastput__ = -1 self.__lastget__ = -1 self.kqmap = {} self.__qkey__ = qkey or "anonymous" self.queues = Proxy(self._getqueues) self.immediate = PriorityQueue() self.mutex = threading.Lock() self.not_empty = threading.Condition(self.mutex) def _nextput(self): queues = self._getqueues() if not queues: raise KQueueNotFound() cur_put = (self.__lastput__ + 1) % len(queues) kq = queues[cur_put] self.__lastput__ = cur_put return kq def _nextget(self): queues = [q for q in self._getqueues() if not q.empty()] if not queues: raise Empty() cur_get = (self.__lastget__ + 1) % len(queues) kq = self.kqmap.values()[cur_get] self.__lastget__ = cur_get return kq def _next_kq(self, next, notfound_exc, key=None): if key is None: kq = next() else: kq = self.find(key) if not kq: raise notfound_exc() if isqueue(kq): return kq else: return kq._next_kq(next=next, notfound_exc=notfound_exc) def _kq_toget(self, key=None): """ Find a kq to get """ return self._next_kq(key=key, next=self._nextget, notfound_exc=KQueueNotFound) def _kq_toput(self, key=None): """ Find a kq to put """ return self._next_kq(key=key, next=self._nextput, notfound_exc=Empty) def put(self, item, key=None): kq = self._kq_toput(key) kq.put_nowait(item) # wake up a greenlet with self.mutex: self.not_empty.notify() put_nowait = put def _get_immediate(self): return self.immediate.get(block=False) def _get(self): try: return 0, self._get_immediate() except Empty: pass remaining_times = [] queues = sorted(self._getqueues(), key=attrgetter("priority")) for q in queues: remaining = q.expected_time() if not remaining: try: self.immediate.put((q.priority, q.get_nowait())) except Empty: pass except RateLimitExceeded: remaining_times.append(q.expected_time()) else: remaining_times.append(remaining) try: return 0, self._get_immediate() except Empty: if not remaining_times: raise return min(remaining_times), None def get(self, block=True, timeout=None): tstart = time() get = self._get not_empty = self.not_empty with not_empty: while 1: try: remaining_time, priority_item = get() except Empty: if not block or (timeout and time() - tstart > timeout): raise not_empty.wait(timeout) continue if remaining_time: if not block or (timeout and time() - tstart > timeout): raise Empty() sleep(min(remaining_time, timeout or 1)) else: return priority_item[1] def get_nowait(self): return self.get(block=False) def touch(self, key, rate_limits=None, capacities=None, priority=0): """ Make sub groups and queue, the deepest queue will be returned """ if key: keys = key.split('/') qkey = keys[-1] groupkeys = keys[0:-1] que = self.find(qkey) or KQueue.make( qkey=qkey, rate_limit=rate_limits[-1] if rate_limits else None, capacity=capacities[-1] if capacities else None, priority=priority) if groupkeys: groupqueue = self.makegroup( '/'.join(groupkeys), rate_limits[0:-1] if rate_limits else None, capacities[0:-1] if capacities else None) else: groupqueue = self groupqueue.addque(qkey, que) return que def makegroup(self, key, rate_limits=None, capacities=None): """ Make sub groups, the deepest group will be returned """ if key: keys = key.split('/') qkey = keys[0] que = self.find(qkey) or KGroupQueue.make( qkey=qkey, rate_limit=rate_limits[0] if rate_limits else None, capacity=capacities[0] if capacities else None) if len(keys) > 1: subgroup = que.makegroup( '/'.join(keys[1:]), rate_limits[1:] if rate_limits else None, capacities[-1] if capacities else None) self.addque(qkey, que) return subgroup else: self.addque(qkey, que) return que def addque(self, qkey, queue, update=False): """ Directly add queue to the current queue group. If update is False, the queue will not be updated if it already exists in the kqmap. This is a low level interface, use it cautiously """ if update or (qkey not in self.kqmap): self.kqmap[qkey] = queue def find(self, key): """ find the group that matches the key """ if key: keys = key.split('/') cur_key = keys[0] if len(keys) == 1: return self.kqmap.get(cur_key) elif cur_key in self.kqmap: return self.kqmap[cur_key].find('/'.join(keys[1:])) def __repr__(self): output = "%s{" % self.__qkey__ for key, que in items(self.kqmap): output += "{0}".format(repr(que)) output += "}" return output def get_priority(self): try: return min(q.priority for q in self.queues) except ValueError: # No queues in the group return -1 def set_priority(self, priority): for q in self.queues: q.priority = priority priority = property(get_priority, set_priority) def _getqueues(self): """ Return all the queues of the group Don't call self.queues in the func """ queues = [] for key, que in items(self.kqmap): if isqueue(que): queues.append(que) else: queues.extend(que._getqueues()) return queues def qsize(self): return sum(q.qsize() for q in self.queues) def empty(self): return all(q.empty() for q in self.queues) def clear(self): for q in self.queues: q.clear()