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)
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 _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
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
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 _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()
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}')
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()
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}')
async def _handle_wrong_cmd(self):
Log.debug('[Server] Wrong cmd')
def die(self) -> None:
if not self._alive:
return
self._alive = False
Log.debug(f'Subscription died {self}')