async def req(self, endpoint: str, model):
if self.ratelimit:
now = datetime.datetime.now()
delta = (now - self.last_req).total_seconds()
if delta < self.ratelimit:
await sleep(self.ratelimit - delta)
self.last_req = now
if not endpoint.startswith("/"):
endpoint = "/" + endpoint
headers = {
"X-Application-Origin": self.app_name,
"X-TOA-Key": self.key,
"Content-Type": "application/json"
}
async with atimeout(5) as _, self.session.get(
"https://theorangealliance.org/api" + endpoint,
headers=headers) as response:
# toa _still_ sometimes returns json data as text/html, making response.json() throw an exception
# _sigh_
data = json.loads(await response.text())
# toa never returns data in dicts, it's always lists
if isinstance(data, dict):
raise AioTOAError(
f"Request to {endpoint} failed with {response.status} {response.reason} (data={data})"
)
return to_model(data, model)
async def run(self):
"""
This can stop with a TimeoutError if a timeout was specified.
If a coroutine was not specified, the message queue is read
once, and the last message is returned.
"""
last_message = None
loop_count = 0
# print('Running {0!s}'.format(self))
while True:
loop_count = loop_count + 1
# print('Waiting for messages ({})'.format(loop_count))
async with atimeout(self.timeout, loop=self.channel.socket.loop):
last_message = await self.channel.receive()
if last_message.event != AbsintheEvent.subscription_data.value:
continue
if 'result' not in last_message.payload:
raise InvalidSubscriptionDataMessage(str(last_message))
if self.coroutine is None:
break
keep_going = await self.coroutine(last_message, **self.kwargs)
if not keep_going:
break
# print('Returning {}'.format(last_message))
return last_message
async def connect(self):
if self.connected:
return
try:
url_parts = [self.url]
if self.params:
qs = urllib.parse.urlencode(self.params)
if len(qs) > 0:
url_parts.append(qs)
url_with_params = '?'.join(url_parts)
logger.info('connect to {}'.format(url_with_params))
self.connected = True
async with atimeout(self._timeout_secs * 2, loop=self.loop):
self._socket = await websockets.connect(url_with_params,
ssl=self._ssl,
loop=self.loop)
self._coroutines = [
asyncio.ensure_future(self.recv(), loop=self.loop),
asyncio.ensure_future(self.heartbeat(), loop=self.loop),
asyncio.ensure_future(self.wait_for_disconnection(),
loop=self.loop)
]
except OSError as e:
if re.compile(r'.*Errno 61.*').match(str(e)):
raise ConnectionRefusedError('server is not responding')
raise e
except asyncio.TimeoutError:
raise TimeoutError('timeout in opening a websocket')
async def test_in_limit(event_loop):
rl = RateLimit(loop=event_loop)
rl.add_limit('get_server_time', 10)
async with atimeout(0.9, loop=event_loop):
for i in range(10):
await rl.wait('get_server_time')
async def wait(self, timeout=None):
gathered = asyncio.gather(*[job.job_task for job in self._in_execution],
loop=self.loop,
return_exceptions=True)
async with atimeout(timeout, loop=self.loop) as to:
await gathered
if to.expired:
raise asyncio.TimeoutError()
async def test_over_limit(event_loop):
rl = RateLimit(loop=event_loop)
rl.add_limit('get_server_time', 10)
with pytest.raises(asyncio.TimeoutError):
async with atimeout(1, loop=event_loop):
for i in range(11):
await rl.wait('get_server_time')
async def joins(topic):
async with phoenix.channel(topic) as ch:
while True:
try:
async with atimeout(0.5, loop=event_loop):
await ch.receive()
await cnt_q.put(1)
except asyncio.TimeoutError:
break
async def _pool_execute(
self,
pool: AbcPool,
args: Sequence,
kwargs: Dict,
*,
timeout: float = None,
) -> Any:
async with atimeout(timeout):
result = await pool.execute(*args, **kwargs)
return result
async def run_client():
with pytest.raises(ConnectionClosed):
try:
async with atimeout(1, loop=event_loop):
async with Phoenix(
'ws://127.0.0.1:{}'.format(unused_tcp_port),
loop=event_loop) as phoenix:
async with phoenix.channel('/channel') as ch:
await asyncio.sleep(10, loop=phoenix.loop)
except asyncio.TimeoutError:
pass
async def run_client():
async with Phoenix('ws://127.0.0.1:{}'.format(unused_tcp_port),
loop=event_loop,
heartbeat_secs=heartbeat_secs) as phoenix:
try:
async with atimeout(heartbeat_secs * (1.5 * num_heartbeat),
loop=phoenix.loop):
run_forever = asyncio.Future(loop=phoenix.loop)
await run_forever
except asyncio.TimeoutError:
pass
async def run_client():
cnt = 0
async with Phoenix('ws://127.0.0.1:{}'.format(unused_tcp_port),
loop=event_loop) as phoenix:
async with phoenix.channel('/channel') as ch:
while True:
try:
async with atimeout(0.5, loop=event_loop):
await ch.receive()
cnt += 1
except asyncio.TimeoutError:
break
assert cnt == notice_cnt
async def handler_ok(ws, _):
while True:
try:
async with atimeout(0.5, loop=ws.loop):
recv_msg = str_to_msg(await ws.recv())
send_msg = msg_to_str(recv_msg.ref, recv_msg.join_ref,
recv_msg.topic,
PHOENIX_EVENT['REPLY'], {
'status': 'ok',
'response': {}
})
await ws.send(send_msg)
except asyncio.TimeoutError:
break
async def test_lock_redis_diff_keys(redis_url: str):
lock1 = await create_lock(redis_url)
lock2 = await create_lock(redis_url)
locks = []
async def coro(lock, key: str, sleep: float, end_lock_cnt: int):
async with lock(key):
locks.append(key)
await asyncio.sleep(sleep)
locks.remove(key)
assert len(locks) == end_lock_cnt
with atimeout(10):
await asyncio.gather(coro(lock1, '1', 1, 1), coro(lock2, '2', 2, 0))
assert len(locks) == 0
async def test_lock_redis_same_keys(redis_url: str):
lock1 = await create_lock(redis_url)
lock2 = await create_lock(redis_url)
locks = []
async def coro(lock, key: str, sleep: float) -> None:
async with lock(key):
assert len(locks) == 0
locks.append(key)
await asyncio.sleep(sleep)
locks.remove(key)
assert len(locks) == 0
key = rndstr()
with atimeout(10):
assert len(locks) == 0
await asyncio.gather(coro(lock1, key, 1), coro(lock2, key, 2))
assert len(locks) == 0
async def _conn_execute(
self,
pool: AbcPool,
args: Sequence,
kwargs: Dict,
*,
timeout: float = None,
asking: bool = False,
) -> Any:
result: Any
async with pool.get() as conn:
try:
async with atimeout(timeout):
if asking:
# emulate command pipeline
results = await asyncio.gather(
conn.execute(b"ASKING"),
conn.execute(*args, **kwargs),
return_exceptions=True,
)
# raise first error
for result in results:
if isinstance(result, BaseException):
raise result
result = results[1]
else:
result = await conn.execute(*args, **kwargs)
return result
except asyncio.TimeoutError:
logger.warning(
"Execute command %s on %s is timed out. Closing connection",
args[0],
pool.address,
)
conn.close()
raise
async def connect(self):
if self.connected:
return
try:
logger.info("connect to %s.", furl(self.url).add(self.params).url)
self.connected = True
async with atimeout(self._timeout_secs * 2, loop=self.loop):
self._socket = await websockets.connect(furl(self.url).add(
self.params).url,
ssl=self._ssl,
loop=self.loop)
self._coroutines = [
asyncio.ensure_future(self.recv(), loop=self.loop),
asyncio.ensure_future(self.heartbeat(), loop=self.loop),
asyncio.ensure_future(self.wait_for_disconnection(),
loop=self.loop)
]
except OSError as e:
if re.compile(r'.*Errno 61.*').match(str(e)):
raise ConnectionRefusedError('Server is not responding.')
raise e
except asyncio.TimeoutError:
raise TimeoutError('Timeout in opening a websocket.')
async def push(self,
channel,
event,
payload,
timeout=None,
retry=3,
wait_response=True):
if not self.connected:
raise CommunicationError("sent a message before connection.")
if timeout is None:
timeout = self._timeout_secs
msg_response = None
ref = None
if retry <= 0:
retry = 0
retry += 1
for i in range(retry):
try:
async with atimeout(timeout, loop=self.loop):
ref = self.make_ref()
if PHOENIX_EVENT['JOIN'] == event:
channel.join_ref = ref
if payload is None:
payload = {}
msg = json.dumps(
OutMessage(
topic=channel.topic,
event=event,
payload=payload,
ref=ref,
join_ref=channel.join_ref,
).asdict())
logger.debug("> %s", msg)
await self._send(msg)
if wait_response:
msg_response = asyncio.Future(loop=self.loop)
if ref in self._waited_messages:
self._waited_messages[ref].cancel()
self._waited_messages[ref] = msg_response
try:
await msg_response
except asyncio.CancelledError:
raise asyncio.TimeoutError
except asyncio.TimeoutError:
retry_str = ''
if i > 0:
retry_str = ' retry {}/{}'.format(i, retry - 1)
logger.warning(
('{}:{} - failed to get a response.' + retry_str).format(
channel.topic, event))
if wait_response and ref is not None:
self._waited_messages.pop(ref, None)
else:
if msg_response is None:
return
return msg_response.result()
raise CommunicationError("failed to send a message.")
async def push(self,
channel,
event,
payload,
timeout=None,
retry=3,
wait_response=True):
if not self.connected:
raise CommunicationError(
'attempting to send a message before connection')
if timeout is None:
timeout = self._timeout_secs
msg_response = None
ref = None
if retry <= 0:
retry = 0
retry += 1
for i in range(retry):
try:
async with atimeout(timeout, loop=self.loop):
ref = self.make_ref()
if PhoenixEvent.join.value == event:
channel.join_ref = ref
if payload is None:
payload = {}
message_data = {
'topic': channel.topic,
'event': event,
'payload': payload,
'ref': ref,
'join_ref': channel.join_ref
}
msg = json.dumps(message_data)
logger.debug('> {}'.format(msg))
await self._send(msg)
if wait_response:
msg_response = asyncio.Future(loop=self.loop)
if ref in self._waited_messages:
self._waited_messages[ref].cancel()
self._waited_messages[ref] = msg_response
try:
await msg_response
except asyncio.CancelledError:
raise asyncio.TimeoutError
except asyncio.TimeoutError:
retry_str = ''
if i > 0:
retry_str = ' retry {}/{}'.format(i, retry - 1)
logger.warning(
('{}/{} - failed to get a response' + retry_str).format(
channel.topic, event))
if wait_response and ref is not None:
self._waited_messages.pop(ref, None)
else:
if msg_response is None:
return
return msg_response.result()
raise CommunicationError('failed to send a message')
async def execute_pubsub(self, *args, **kwargs) -> List[PubsubResponse]:
"""Execute Redis (p)subscribe/(p)unsubscribe commands."""
ctx = self._make_exec_context(args, kwargs)
# get first pattern and calculate slot
channel_name = ctx.cmd[1]
is_pattern = ctx.cmd_name in {"PSUBSCRIBE", "PUNSUBSCRIBE"}
is_unsubscribe = ctx.cmd_name in {"UNSUBSCRIBE", "PUNSUBSCRIBE"}
ctx.slot = key_slot(channel_name)
exec_fail_props: Optional[ExecuteFailProps] = None
result: List[List]
while ctx.attempt < ctx.max_attempts:
self._check_closed()
ctx.attempt += 1
exec_fail_props = None
state = await self._manager.get_state()
node_addr = self._pooler.get_pubsub_addr(channel_name, is_pattern)
# if unsuscribe command and no node found for pattern
# probably pubsub connection is already close
if is_unsubscribe and node_addr is None:
result = [[ctx.cmd[0], channel_name, 0]]
break
if node_addr is None:
try:
node_addr = state.random_slot_node(ctx.slot).addr
except UncoveredSlotError:
logger.warning("No any node found by slot %d", ctx.slot)
node_addr = state.random_node().addr
try:
pool = await self._pooler.ensure_pool(node_addr)
async with atimeout(self._attempt_timeout):
result = await pool.execute_pubsub(*ctx.cmd, **ctx.kwargs)
if is_unsubscribe:
self._pooler.remove_pubsub_channel(channel_name, is_pattern)
else:
self._pooler.add_pubsub_channel(node_addr, channel_name, is_pattern)
except asyncio.CancelledError:
raise
except Exception as e:
exec_fail_props = ExecuteFailProps(
node_addr=node_addr,
error=e,
)
if exec_fail_props:
await self._on_execute_fail(ctx, exec_fail_props)
continue
break
return [(cmd, name, count) for cmd, name, count in result]
async def wait(self, timeout=None):
self._waited = True
async with atimeout(timeout) as to:
await self.job_task
if to.expired:
raise asyncio.TimeoutError()
