def __init__(self, expiredIn): self.messageQueue = RBTree() self.subscribers = RBTree() self.expiredIn = expiredIn self.idGenerator = IdGenerator() task.LoopingCall(self.purgeMessageQueue).start(1, False) task.LoopingCall(self.purgeSubscribers).start(1, False)
class Channel: """ Channel provides pub/sub service for a single channel""" def __init__(self, expiredIn): self.messageQueue = RBTree() self.subscribers = RBTree() self.expiredIn = expiredIn self.idGenerator = IdGenerator() task.LoopingCall(self.purgeMessageQueue).start(1, False) task.LoopingCall(self.purgeSubscribers).start(1, False) def subscriber_size(self): count = 0 for s in self.subscribers.values(): count += len(s) return count def publish(self, data, isAsync = True): """ Publish the data to the message queue. After messages are published, interested subscribers will be notified. @params isAsync: specify if the subscribers should be notified asynchronously or synchronously. """ time = now() dataWithId = (self.idGenerator.generateId(), data) self.messageQueue.setdefault(time, []).append(dataWithId) self.notify(time, dataWithId, isAsync) def subscribe(self, onReceive, onTimeout, timeFrom = 0, timeTo = sys.maxint, minId = 0, timeoutSec = 0): """ subscribe messages within a specific time span. @params onReceive: if the interested messages are retrieved, onReceive will be invoked to notify the subscribers. @params onTimeout: if subscriber waits for more than `timeoutSec` seconds, onTimeout will be invoked. @params timeFrom: only retrieve messages after timestamp `timeFrom`; time is represented in unix time. @params timeTo: only retrieve messages before timestamp `timeTo`; time is represented in unix time. @params minId: this is HACK... """ messages = self._flatten(self.messageQueue[timeFrom: timeTo], minId) messages = list(messages) if len(messages) != 0 or timeoutSec == 0: onReceive(messages) return waitUntil = now() + timeoutSec subscription = (timeFrom, timeTo, onReceive, onTimeout) self.subscribers.setdefault(waitUntil, []).append(subscription) def notify(self, time, data, isAsync = True): # purge expired subscribers self.purgeSubscribers() for time, bucket in self.subscribers.items(): index = partition(bucket, lambda sub: time < sub[0] or time > sub[1]) data = [data] for i in range(index, len(bucket)): callback = bucket[i][2] if isAsync: reactor.callLater(0, callback, data) else: callback(data) del bucket[index:] if len(bucket) == 0: del self.subscribers[time] def purgeMessageQueue(self): expired = now() - self.expiredIn + 1 del self.messageQueue[:expired] def purgeSubscribers(self): expired = now() + 1 for data in self._flatten(self.subscribers[:expired], 0): data[3]() del self.subscribers[:expired] # --- Utilities --- def _flatten(self, buckets, minId): for bucket in buckets.values(): for item in bucket: if item[0] >= minId: yield item