async def frequency_limiter(sem: Semaphore) -> None: """ Function that must be used only inside router module to wait some time before answer to the client and let the background task to start :param sem: semaphore to store iot or user data """ try: await wait_for(sem.acquire(), 1) sem.release() except TimeoutError: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, headers={"Retry-After": str(randrange(1, 29) * random() + 1)}, )
class PriorityQueue: def __init__(self): self.queue = [] self.items = Semaphore(value=0) async def push(self, data, priority=0): self.queue.append((priority, json.dumps(data))) self.queue.sort() self.items.release() async def pop(self, timeout: int = 1) -> Any: try: await wait_for(self.items.acquire(), timeout) return json.loads(self.queue.pop(-1)[1]) except TimeoutError: return None async def pop_ready(self) -> Any: if self.items.locked(): return None await self.items.acquire() return json.loads(self.queue.pop(-1)[1]) async def score(self, data): data = json.dumps(data) for priority, item in self.queue: if data == item: return priority return None async def rank(self, data): data = json.dumps(data) for index, (_, item) in enumerate(self.queue): if data == item: return len(self.queue) - index - 1 return None async def clear(self): self.queue = [] self.items = Semaphore(value=0) async def length(self): return len(self.queue)
class AsyncioSubscriptionManager(SubscriptionManager): def __init__(self, pubnub_instance): subscription_manager = self self._message_worker = None self._message_queue = Queue() self._subscription_lock = Semaphore(1) self._subscribe_loop_task = None self._heartbeat_periodic_callback = None self._reconnection_manager = AsyncioReconnectionManager(pubnub_instance) super(AsyncioSubscriptionManager, self).__init__(pubnub_instance) self._start_worker() class AsyncioReconnectionCallback(ReconnectionCallback): def on_reconnect(self): subscription_manager.reconnect() pn_status = PNStatus() pn_status.category = PNStatusCategory.PNReconnectedCategory pn_status.error = False subscription_manager._subscription_status_announced = True subscription_manager._listener_manager.announce_status(pn_status) self._reconnection_listener = AsyncioReconnectionCallback() self._reconnection_manager.set_reconnection_listener(self._reconnection_listener) def _set_consumer_event(self): if not self._message_worker.cancelled(): self._message_worker.cancel() def _message_queue_put(self, message): self._message_queue.put_nowait(message) def _start_worker(self): consumer = AsyncioSubscribeMessageWorker(self._pubnub, self._listener_manager, self._message_queue, None) self._message_worker = asyncio.ensure_future(consumer.run(), loop=self._pubnub.event_loop) def reconnect(self): # TODO: method is synchronized in Java self._should_stop = False self._subscribe_loop_task = asyncio.ensure_future(self._start_subscribe_loop()) self._register_heartbeat_timer() def disconnect(self): # TODO: method is synchronized in Java self._should_stop = True self._stop_heartbeat_timer() self._stop_subscribe_loop() def stop(self): super(AsyncioSubscriptionManager, self).stop() self._reconnection_manager.stop_polling() if self._subscribe_loop_task is not None and not self._subscribe_loop_task.cancelled(): self._subscribe_loop_task.cancel() @asyncio.coroutine def _start_subscribe_loop(self): self._stop_subscribe_loop() yield from self._subscription_lock.acquire() combined_channels = self._subscription_state.prepare_channel_list(True) combined_groups = self._subscription_state.prepare_channel_group_list(True) if len(combined_channels) == 0 and len(combined_groups) == 0: self._subscription_lock.release() return self._subscribe_request_task = asyncio.ensure_future(Subscribe(self._pubnub) .channels(combined_channels) .channel_groups(combined_groups) .timetoken(self._timetoken).region(self._region) .filter_expression(self._pubnub.config.filter_expression) .future()) e = yield from self._subscribe_request_task if self._subscribe_request_task.cancelled(): self._subscription_lock.release() return if e.is_error(): if e.status is not None and e.status.category == PNStatusCategory.PNCancelledCategory: self._subscription_lock.release() return if e.status is not None and e.status.category == PNStatusCategory.PNTimeoutCategory: self._pubnub.event_loop.call_soon(self._start_subscribe_loop) self._subscription_lock.release() return logger.error("Exception in subscribe loop: %s" % str(e)) if e.status is not None and e.status.category == PNStatusCategory.PNAccessDeniedCategory: e.status.operation = PNOperationType.PNUnsubscribeOperation # TODO: raise error self._listener_manager.announce_status(e.status) self._reconnection_manager.start_polling() self._subscription_lock.release() self.disconnect() return else: self._handle_endpoint_call(e.result, e.status) self._subscription_lock.release() self._subscribe_loop_task = asyncio.ensure_future(self._start_subscribe_loop()) self._subscription_lock.release() def _stop_subscribe_loop(self): if self._subscribe_request_task is not None and not self._subscribe_request_task.cancelled(): self._subscribe_request_task.cancel() def _stop_heartbeat_timer(self): if self._heartbeat_periodic_callback is not None: self._heartbeat_periodic_callback.stop() def _register_heartbeat_timer(self): super(AsyncioSubscriptionManager, self)._register_heartbeat_timer() self._heartbeat_periodic_callback = AsyncioPeriodicCallback( self._perform_heartbeat_loop, self._pubnub.config.heartbeat_interval * 1000, self._pubnub.event_loop) if not self._should_stop: self._heartbeat_periodic_callback.start() @asyncio.coroutine def _perform_heartbeat_loop(self): if self._heartbeat_call is not None: # TODO: cancel call pass cancellation_event = Event() state_payload = self._subscription_state.state_payload() presence_channels = self._subscription_state.prepare_channel_list(False) presence_groups = self._subscription_state.prepare_channel_group_list(False) if len(presence_channels) == 0 and len(presence_groups) == 0: return try: heartbeat_call = (Heartbeat(self._pubnub) .channels(presence_channels) .channel_groups(presence_groups) .state(state_payload) .cancellation_event(cancellation_event) .future()) envelope = yield from heartbeat_call heartbeat_verbosity = self._pubnub.config.heartbeat_notification_options if envelope.status.is_error: if heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL or \ heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL: self._listener_manager.announce_status(envelope.status) else: if heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL: self._listener_manager.announce_status(envelope.status) except PubNubAsyncioException as e: pass # TODO: check correctness # if e.status is not None and e.status.category == PNStatusCategory.PNTimeoutCategory: # self._start_subscribe_loop() # else: # self._listener_manager.announce_status(e.status) finally: cancellation_event.set() def _send_leave(self, unsubscribe_operation): asyncio.ensure_future(self._send_leave_helper(unsubscribe_operation)) @asyncio.coroutine def _send_leave_helper(self, unsubscribe_operation): envelope = yield from Leave(self._pubnub) \ .channels(unsubscribe_operation.channels) \ .channel_groups(unsubscribe_operation.channel_groups).future() self._listener_manager.announce_status(envelope.status)
async def _drain_semaphore(semaphore: asyncio.Semaphore): while not semaphore.locked(): try: await asyncio.wait_for(semaphore.acquire(), 0.1) except asyncio.TimeoutError: break
class scheduler(threading.Thread): def __init__(self, worker_q, ws, map_jobs, n_reducer=30, url=""): ''' worker_q: we put availabe workers into this priority queue. ws: a websocket connection object. map_jobs: a list storing all map tasks (more precisely: which part of the input file) n_reducer: Number of reducer. ''' super(scheduler, self).__init__() self.url_list = [] self.url = url self.mutex = Semaphore() self.map_jobs = map_jobs self.n_reducer = n_reducer self.stoprequest = threading.Event() self.worker_q = worker_q self.ws = ws self.dead_worker = set() self.mapCount = 0 self.reduceCount = n_reducer # a dict used to track status of all map/reduce jobs, # entities in map_jobs are keys, # a list of worker(if not finished) or None object(if finished) # is the corresponding value self.map_status = {} self.reduce_status = {} self.tid_map = {} for job in map_jobs: self.map_status[job] = [] for i in range(n_reducer): self.reduce_status[i] = [] def removeWorker(self, uid): self.dead_worker.add(uid) def jobFinished(self, tid, url, type): ''' Nth slice of the map job is finished, mark them as done in map_status ''' if type == "m": if self.mapCount == 0: return self.mutex.acquire() n = self.tid_map[tid] if self.map_status[n] == None: self.mutex.release() return for worker in self.map_status[n]: self.worker_q.put(worker) self.url_list.append(url) self.map_status[n] = None self.mapCount -= 1 self.map_jobs.remove(n) self.mutex.release() return if type == "r": self.mutex.acquire() n = self.tid_map[tid] if self.reduce_status[n] == None: self.mutex.release() return for worker in self.reduce_status[n]: self.worker_q.put(worker) self.reduce_status[n] = None self.reduceCount -= 1 self.mutex.release() return def schedule_reduce(self): print("start reduce") counter = 0 while (True): new_worker = self.worker_q.get(True) if new_worker[1] in self.dead_worker: self.dead_worker.remove(new_worker[1]) continue self.mutex.acquire() if self.reduceCount == 0: return counter %= self.n_reducer while (self.reduce_status[counter] == None): counter += 1 self.reduce_status[counter].append(new_worker) self.mutex.release() tid = (int(time.time() * 1000) + new_worker[0]) self.tid_map[tid] = counter data = { 'type': 'r', 'uid': new_worker[1], 'tid': tid, 'slice': counter, 'url': self.url_list } self.ws.send(msg_generator(1, "", "task", data)) counter += 1 def schedule_map(self): ''' Map tasks are scheduled in this function, it will return iff all map tasks are finished. ''' self.mapCount = len(self.map_jobs) counter = 0 while (True): new_worker = self.worker_q.get(True) if new_worker[1] in self.dead_worker: self.dead_worker.remove(new_worker[1]) continue self.mutex.acquire() if self.mapCount == 0: self.worker_q.put(new_worker) return counter %= self.mapCount job = self.map_jobs[counter] counter += 1 self.map_status[job].append(new_worker) self.mutex.release() tid = (int(time.time() * 1000) + new_worker[0]) #tid = counter self.tid_map[tid] = job data = { 'type': 'm', 'uid': new_worker[1], 'tid': tid, 'slice': job, 'url': [self.url] } self.ws.send(msg_generator(1, "", "task", data)) def run(self): while not self.stoprequest.isSet(): self.schedule_map() print("map done!") self.schedule_reduce() print("reduce done!")
class AsyncioSubscriptionManager(SubscriptionManager): def __init__(self, pubnub_instance): subscription_manager = self self._message_worker = None self._message_queue = Queue() self._subscription_lock = Semaphore(1) self._subscribe_loop_task = None self._heartbeat_periodic_callback = None self._reconnection_manager = AsyncioReconnectionManager(pubnub_instance) super(AsyncioSubscriptionManager, self).__init__(pubnub_instance) self._start_worker() class AsyncioReconnectionCallback(ReconnectionCallback): def on_reconnect(self): subscription_manager.reconnect() pn_status = PNStatus() pn_status.category = PNStatusCategory.PNReconnectedCategory pn_status.error = False subscription_manager._subscription_status_announced = True subscription_manager._listener_manager.announce_status(pn_status) self._reconnection_listener = AsyncioReconnectionCallback() self._reconnection_manager.set_reconnection_listener(self._reconnection_listener) def _set_consumer_event(self): if not self._message_worker.cancelled(): self._message_worker.cancel() def _message_queue_put(self, message): self._message_queue.put_nowait(message) def _start_worker(self): consumer = AsyncioSubscribeMessageWorker(self._pubnub, self._listener_manager, self._message_queue, None) self._message_worker = asyncio.ensure_future(consumer.run(), loop=self._pubnub.event_loop) def reconnect(self): # TODO: method is synchronized in Java self._should_stop = False self._subscribe_loop_task = asyncio.ensure_future(self._start_subscribe_loop()) self._register_heartbeat_timer() def disconnect(self): # TODO: method is synchronized in Java self._should_stop = True self._stop_heartbeat_timer() self._stop_subscribe_loop() def stop(self): super(AsyncioSubscriptionManager, self).stop() self._reconnection_manager.stop_polling() if self._subscribe_loop_task is not None and not self._subscribe_loop_task.cancelled(): self._subscribe_loop_task.cancel() @asyncio.coroutine def _start_subscribe_loop(self): self._stop_subscribe_loop() yield from self._subscription_lock.acquire() combined_channels = self._subscription_state.prepare_channel_list(True) combined_groups = self._subscription_state.prepare_channel_group_list(True) if len(combined_channels) == 0 and len(combined_groups) == 0: self._subscription_lock.release() return self._subscribe_request_task = asyncio.ensure_future(Subscribe(self._pubnub) .channels(combined_channels) .channel_groups(combined_groups) .timetoken(self._timetoken).region(self._region) .filter_expression(self._pubnub.config.filter_expression) .future()) e = yield from self._subscribe_request_task if self._subscribe_request_task.cancelled(): self._subscription_lock.release() return if e.is_error(): if e.status is not None and e.status.category == PNStatusCategory.PNCancelledCategory: self._subscription_lock.release() return if e.status is not None and e.status.category == PNStatusCategory.PNTimeoutCategory: self._pubnub.event_loop.call_soon(self._start_subscribe_loop) self._subscription_lock.release() return logger.error("Exception in subscribe loop: %s" % str(e)) if e.status is not None and e.status.category == PNStatusCategory.PNAccessDeniedCategory: e.status.operation = PNOperationType.PNUnsubscribeOperation # TODO: raise error self._listener_manager.announce_status(e.status) self._reconnection_manager.start_polling() self._subscription_lock.release() self.disconnect() return else: self._handle_endpoint_call(e.result, e.status) self._subscription_lock.release() self._subscribe_loop_task = asyncio.ensure_future(self._start_subscribe_loop()) self._subscription_lock.release() def _stop_subscribe_loop(self): if self._subscribe_request_task is not None and not self._subscribe_request_task.cancelled(): self._subscribe_request_task.cancel() def _stop_heartbeat_timer(self): if self._heartbeat_periodic_callback is not None: self._heartbeat_periodic_callback.stop() def _register_heartbeat_timer(self): super(AsyncioSubscriptionManager, self)._register_heartbeat_timer() self._heartbeat_periodic_callback = AsyncioPeriodicCallback( self._perform_heartbeat_loop, self._pubnub.config.heartbeat_interval * 1000, self._pubnub.event_loop) if not self._should_stop: self._heartbeat_periodic_callback.start() @asyncio.coroutine def _perform_heartbeat_loop(self): if self._heartbeat_call is not None: # TODO: cancel call pass cancellation_event = Event() state_payload = self._subscription_state.state_payload() presence_channels = self._subscription_state.prepare_channel_list(False) presence_groups = self._subscription_state.prepare_channel_group_list(False) if len(presence_channels) == 0 and len(presence_groups) == 0: return try: heartbeat_call = (Heartbeat(self._pubnub) .channels(presence_channels) .channel_groups(presence_groups) .state(state_payload) .cancellation_event(cancellation_event) .future()) envelope = yield from heartbeat_call heartbeat_verbosity = self._pubnub.config.heartbeat_notification_options if envelope.status.is_error: if heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL or \ heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL: self._listener_manager.announce_stateus(envelope.status) else: if heartbeat_verbosity == PNHeartbeatNotificationOptions.ALL: self._listener_manager.announce_stateus(envelope.status) except PubNubAsyncioException as e: pass # TODO: check correctness # if e.status is not None and e.status.category == PNStatusCategory.PNTimeoutCategory: # self._start_subscribe_loop() # else: # self._listener_manager.announce_status(e.status) finally: cancellation_event.set() def _send_leave(self, unsubscribe_operation): asyncio.ensure_future(self._send_leave_helper(unsubscribe_operation)) @asyncio.coroutine def _send_leave_helper(self, unsubscribe_operation): envelope = yield from Leave(self._pubnub) \ .channels(unsubscribe_operation.channels) \ .channel_groups(unsubscribe_operation.channel_groups).future() self._listener_manager.announce_status(envelope.status)
class ConnectionPool: def __init__(self, addr, configuration, maxconnections): self._addr = addr self._configuration = configuration if not maxconnections: raise Exception('Please set a maximum limit of connections in the pool') self._limit = Semaphore(maxconnections) self._available = deque() self._inuse = set() self._fast = None self._lock = Lock() async def take(self, is_slow, timeout=None): if not is_slow: if self._fast is None or self._fast.closed: async with self._lock: if self._fast is None or self._fast.closed: self._fast = await Connection.create(self._addr, self._configuration) return self._fast else: try: while True: conn = self._available.popleft() if not conn.closed: break self._limit.release() except IndexError: if timeout is None: await self._limit.acquire() else: await wait_for(self._limit.acquire(), timeout) try: conn = await Connection.create(self._addr, self._configuration) conn.set_release_cb(self.release) except: self._limit.release() raise self._inuse.add(conn) return conn def release(self, conn): # This is also a protection against double .release call on a single .take call self._inuse.remove(conn) if not conn.closed: self._available.append(conn) else: self._limit.release() async def aclose(self): if self._fast: await self._fast.aclose() self._fast = None if self._available is not None and self._inuse is not None: tmp = list(self._available) + list(self._inuse) for conn in tmp: await conn.aclose() self._addr = None self._configuration = None self._limit = None self._available = None self._inuse = None self._lock = None