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
