async def _make_pings_task(self):
""" Task that sends pings to all subscriptions in the queue.
All subs in this queue have timed out.
"""
while True:
await self._check_timeouts()
while not self._ping_pong_tasks.empty():
try:
subscriber = await self._ping_pong_tasks.get()
if subscriber.is_dead():
self._delete_subscription(subscriber)
else:
await subscriber.ping()
except (ConnectionResetError, BrokenPipeError):
Log.err(f'Connection unexpectedly closed {subscriber}')
self._delete_subscription(subscriber)
finally:
# We want to remove the task from the queue regardless
# if it fails or completes.
self._ping_pong_tasks.task_done()
# Go idle so other tasks can run.
await asyncio.sleep(TASK_DELAY_PING)
async def start(self, topic: str) -> bool:
"""
Starts the subscription.
param data_received is a callback that will be called
when new data arrvies at the subscribed topic.
"""
# Start connection and initialize a new Subscription.
await self.open()
await self.send(protocol.Commands.SUBSCRIBE)
# Wait for ACK, before sending the topic.
packet = await self.read()
if not await protocol.async_cmd_ok(
packet, protocol.Commands.SUBSCRIBE_ACK, self._writer):
return
# Send the topic we want to subscribe to.
await self.send(protocol.Commands.SUBSCRIBE,
data=topic.encode('utf-8'))
# Wait for OK/NOT OK.
sub_ack = await self.read()
if not await protocol.async_cmd_ok(sub_ack,
protocol.Commands.SUBSCRIBE_OK):
Log.err(f'Failed to subscribe to {topic}. Got no ACK.')
return False
# Enter ping-pong state where we just wait for data published to
# the chosen topic and give the data to the callback function provided.
await self._enter_ping_pong()
async def _handle_publish(self) -> Publication:
# Log.debug('Server: PUB')
# Send: Publish ACK
await protocol.send_packet(self._writer, protocol.Commands.PUBLISH_ACK)
# Receive: Publish + Data ('topic' | 'message')
packet = await protocol.read_packet(self._reader)
if not await protocol.async_cmd_ok(packet, protocol.Commands.PUBLISH,
self._writer):
return
data = packet.data.decode('utf-8')
# Ensure publish is OK according to the format required.
if not protocol.publish_ok(data):
Log.debug(f'[Server] Pub -> Publish "{data}" is incorrect format.')
await protocol.send_packet(self._writer,
protocol.Commands.INCORRECT_FORMAT)
return None
# Publish OK
await protocol.send_packet(self._writer, protocol.Commands.PUBLISH_OK)
topic, message = protocol.get_topic_and_msg(data)
publication = Publication(topic, message)
fd = self._writer.get_extra_info('socket').fileno()
Log.info(f'[Server] New Pub: "{fd}: {publication}"')
return publication
def __init__(self, ip: str, port: int):
self._subscriptions = SubscriptionContainer()
self._publications = asyncio.Queue()
self._ping_pong_tasks = asyncio.PriorityQueue()
self._ip = ip
self._port = port
Log.init(server=True)
async def subscribe(self,
topic: str,
data_received: callable = None) -> None:
""" Throws ConnectionRefusedError. """
sub = AsyncSubscribe(self._ip, self._port, data_received)
Log.info(f'[Client] Subscribing to "{topic}"')
await sub.start(topic)
async def start(self) -> None:
""" Starts the server. This method runs forever. """
server = await asyncio.start_server(self._request_handler,
self._ip, self._port)
ip, port = server.sockets[0].getsockname()
Log.info(f'{NAME} Server started at {ip} on port {port}')
async with server:
await asyncio.gather(self._make_pings_task(),
self._make_publications_task(),
server.serve_forever())
def _set_identifier(self, topic: str) -> None:
""" Sets the identification of the subscription.
This consists of:
1. Topic
2. File descripter number from reader/writer stream.
"""
self.topic = topic
try:
self.fd = self._writer.get_extra_info('socket').fileno()
except AttributeError:
# Streams are incorrect
Log.err(f'Incorrect streams to subscription to {self.topic}')
self.fd = None
def delete(self, subscription: Subscription) -> bool:
topic = subscription.topic
del_ok = False
try:
self._subscriptions[topic].remove(subscription)
subscription.die()
del_ok = True
# Remove topic if there's no subscribers left.
if len(self._subscriptions[topic]) == 0:
self._subscriptions[topic].remove(topic)
except KeyError:
Log.debug(f'Failed to find sub on topic {topic}')
return del_ok
async def _make_publications(self, publication: Publication) -> None:
""" Sends the publication to all the subscribers of the topic. """
subs = self._subscriptions.get(publication.topic)
for sub in subs.copy():
try:
Log.debug(f'[Server] Publishing: {self._publications.qsize()}')
pub_ok = await sub.new_data(publication.message)
if not pub_ok:
self._delete_subscription(sub)
except RuntimeError:
# This error is caused: RuntimeError: read() called while
# another coroutine is already waiting for incoming data.
# Should not do any harm, so therefore ignored.
pass
async def _request_handler(self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> None:
""" Handles a TCP request. """
request = Request(reader, writer)
result: RequestResult = await request.respond()
if result == Subscription:
await self._add_subscription(result.data)
elif result == Publication:
# Add new publication to queue
await self._publications.put(result.data)
elif result is None:
# This should not occur!
Log.debug('ALERT: Result is None! ')
async def ping(self) -> None:
""" Pings the subscriber and waits for a PONG back.
If the subscriber doesn't pong back, the subscription is closed.
"""
await protocol.send_packet(self._writer, protocol.Commands.PING)
Log.debug(f'Ping {self}')
pong = await protocol.read_packet(self._reader)
if await protocol.async_cmd_ok(pong, protocol.Commands.PONG):
# If PONG, reset timer.
self._time = 0
else:
Log.err(f'Bad ping! {self._alive} -> {self._state}')
# If no PONG, advance to next state, and potentially close.
alive = self._next_state()
if not alive:
self.die()
async def _handle_subscribe(self) -> Subscription:
Log.debug('[Server] New sub requested')
# Subscribe ACK
await protocol.send_packet(self._writer,
protocol.Commands.SUBSCRIBE_ACK)
# Subscribe 'topic'
packet = await protocol.read_packet(self._reader)
if not await protocol.async_cmd_ok(packet, protocol.Commands.SUBSCRIBE,
self._writer):
return None
topic = packet.data.decode('utf-8')
# Ensure topic is OK according to the format required.
if not protocol.topic_ok(topic):
Log.debug(f'[Server] Sub -> Topic "{topic}" is incorrect format.')
return None
# Subscribe OK
await protocol.send_packet(self._writer,
protocol.Commands.SUBSCRIBE_OK)
sub = Subscription(topic, self._reader, self._writer)
Log.info(f'[Server] New sub -> {sub}')
return sub
async def _enter_ping_pong(self) -> None:
self._alive = True
while self._alive:
packet = await self.read()
if packet.cmd == protocol.Commands.NEW_DATA:
# New data published!
await self._ack_new_data()
# Unpack the data and send it to callback.
data = packet.data.decode('utf-8')
if self._data_received is not None:
self._data_received(data)
else:
Log.err(f'No callback for new data provided!\n{data}')
else:
# Expecting a ping
if not await protocol.async_cmd_ok(packet,
protocol.Commands.PING):
Log.err('Failed to get ping command. Exiting.')
return
# Send a PONG back.
await self._pong()
Log.debug('[Client] Pong')
# If provided, call pong callback.
if self._pong_callback is not None:
self._pong_callback()
await asyncio.sleep(SLEEP_DELAY)
def get() -> dict:
"""
Returns the configurations.
If can't find the configuration file, default parameters are used
and a new one is created on disk.
"""
success = False
# Only need to read the config once (assumes no live-changes).
if Config._config is None:
try:
with open(CONFIG_PATH, 'rb') as f:
Config._config = json.load(f)
success = True
except FileNotFoundError:
Log.err(f'Failed to find config at {CONFIG_PATH}')
except json.decoder.JSONDecodeError as e:
Log.err(f'Failed to parse JSON {e}')
if not success:
Config._config = Config.__default
Config._save_config(Config._config)
return Config._config
async def publish(self, topic: str, message: str) -> bool:
"""
Returns if the publish is succesful or not.
Throws ConnectionRefusedError.
"""
publish = AsyncPublish(self._ip, self._port)
try:
pub_ok = await publish.start(topic, message)
except ConnectionRefusedError:
Log.err(f'[Client ]Failed to connect to server {self._ip} '
f'on port {self._port}')
return
if pub_ok:
Log.info(f'[Client] Published "{message}" to topic "{topic}"')
else:
Log.info(f'[Client] Failed to publish "{message}" to'
f' topic "{topic}"')
return pub_ok
def open(self):
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._sock.connect((self._ip, self._port))
self._writer = self._sock.makefile('wb')
self._reader = self._sock.makefile('rb')
Log.debug(f'Connecting to {self._ip}:{self._port}')
def _save_config(config: dict) -> None:
Log.info(f'Saving configuration file at {CONFIG_PATH}')
with open(CONFIG_PATH, 'w') as f:
json.dump(config, f)
async def _add_subscription(self, subscription: Subscription) -> None:
self._subscriptions.add(subscription)
Log.debug(f'Total subscribers: {self._subscriptions.get_all()}')
await subscription.start_idle()
def __init__(self, ip: str, port: int):
self._ip = ip
self._port = port
Log.init(server=False)
def die(self) -> None:
if not self._alive:
return
self._alive = False
Log.debug(f'Subscription died {self}')
async def _handle_wrong_cmd(self):
Log.debug('[Server] Wrong cmd')
def setUp(self):
self.ip = Config.get()['client']['ip']
self.port = Config.get()['client']['port']
self.client = Client(self.ip, self.port)
Log.disable()
async def open(self):
self._reader, self._writer = await asyncio.open_connection(
self._ip, self._port)
Log.debug(f'Connecting to {self._ip}:{self._port}')
def setUp(self):
self.container = SubscriptionContainer()
Log.disable()
