Enable two way communication

Author: Marco Paolini

pushpull/amqp/client.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 
 from .gateway import Exchanger
@@ -6,8 +7,21 @@
 logger = logging.getLogger(__name__)
 
 
-async def challenge(fd):
-    async with Exchanger('1') as (sender, _):
-        async for line in FdLineReader(fd):
-            logger.debug('sending message %r', line)
-            await sender.send(line)
+async def challenge(name, fd_in, fd_out):
+    async with Exchanger(name, Exchanger.ROLE_APP) as (amqp_sender, amqp_receiver):
+        sender = send_from_fd_to_amqp(fd_in, amqp_sender)
+        receiver = send_from_amqp_to_fd(amqp_receiver, fd_out)
+        _, pending = await asyncio.wait([sender, receiver], return_when=asyncio.FIRST_COMPLETED)
+        for task in pending:
+            task.cancel()
+
+
+async def send_from_fd_to_amqp(fd, sender):
+    async for line in FdLineReader(fd):
+        logger.debug('sending message %r', line)
+        await sender.send(line)
+
+
+async def send_from_amqp_to_fd(receiver, fd):
+    async for msg in receiver:
+        fd.write(msg)

pushpull/amqp/gateway.py

@@ -9,19 +9,31 @@
 
 class Exchanger:
 
-    def __init__(self, name, **connection_params):
-        self.exchange_name = name
+    ROLE_WS = 1
+    ROLE_APP = 2
+
+    def __init__(self, name, role, **connection_params):
         self._conn_params = connection_params
+        if role not in [self.ROLE_WS, self.ROLE_APP]:
+            raise ValueError('bad role {}'.format(role))
+        self.role = role
+        self.q_names = [x.format(name) for x in ('{}.from_ws', '{}.to_ws')]
+        if role == self.ROLE_APP:
+            self.q_names.reverse()
 
     async def __aenter__(self):
         logger.debug('connecting sender and receiver')
         self._conn = await asynqp.connect(**self._conn_params)
         self._chan = await self._conn.open_channel()
-        exchange = await self._chan.declare_exchange(self.exchange_name, 'direct')
-        queue = await self._chan.declare_queue(self.exchange_name)
-        await queue.bind(exchange, self.exchange_name)
+        queues, exchanges = [], []
+        for name in self.q_names:
+            exchange = await self._chan.declare_exchange(name, 'direct')
+            exchanges.append(exchange)
+            queue = await self._chan.declare_queue(name)
+            await queue.bind(exchange, name)
+            queues.append(queue)
         logger.debug('connected sender and receiver')
-        return Sender(exchange, self.exchange_name), Receiver(queue)
+        return Sender(exchanges[0], self.q_names[0]), Receiver(queues[1])
 
     async def __aexit__(self, exc_type, exc_value, traceback):
         logger.debug('closing connection and channel %r %r', exc_type, exc_value)
@@ -56,6 +68,8 @@ def __call__(self, amqp_message):
             self._fifo.put_nowait(amqp_message.body.decode(amqp_message.content_encoding))
         except asyncio.QueueFull:
             logger.warning('queue full')
+        else:
+            amqp_message.ack()
 
     def on_error(self, exc):
         logger.error('error received: %r', exc)

pushpull/cli/client.py

@@ -15,16 +15,18 @@ def client():
 
 @click.command()
 @click.argument('url')
-def challenge_websocket(url):
+@click.argument('name')
+def challenge_websocket(url, name):
     logging.basicConfig(level=logging.DEBUG)
-    click.echo(asyncio.get_event_loop().run_until_complete(websocket_client.challenge(url, sys.stdin, sys.stdout)))
+    click.echo(asyncio.get_event_loop().run_until_complete(websocket_client.challenge(url, name, sys.stdin, sys.stdout)))
 client.add_command(challenge_websocket)
 
 
 @click.command()
-def challenge_amqp():
+@click.argument('name')
+def challenge_amqp(name):
     logging.basicConfig(level=logging.DEBUG)
-    click.echo(asyncio.get_event_loop().run_until_complete(amqp_client.challenge(sys.stdin)))
+    click.echo(asyncio.get_event_loop().run_until_complete(amqp_client.challenge(name, sys.stdin, sys.stdout)))
 client.add_command(challenge_amqp)
 
 

pushpull/websocket/client.py

@@ -9,30 +9,30 @@
 logger = logging.getLogger(__name__)
 
 
-async def challenge(url, fd_in, fd_out):
+async def challenge(url, name, fd_in, fd_out):
     with aiohttp.ClientSession() as session:
-        async with session.ws_connect(url) as ws:
+        async with session.ws_connect('{}?name={}'.format(url, name)) as ws:
             logger.debug('opening websocket')
-            await send_from_fd_to_ws(fd_in, ws)
-            return
-            sender = asyncio.ensure_future(send_from_fd_to_ws(fd_in, ws))
-            receiver = asyncio.ensure_future(send_from_ws_to_fd(fd_out, ws))
-            asyncio.gather(sender, receiver)
+            sender = send_from_fd_to_ws(fd_in, ws)
+            receiver = send_from_ws_to_fd(ws, fd_out)
+            _, pending = await asyncio.wait([sender, receiver], return_when=asyncio.FIRST_COMPLETED)
+            for task in pending:
+                task.cancel()
             logger.debug('closing websocket')
             # await ws.close()
 
 
 async def send_from_fd_to_ws(fd, ws):
     async for line in FdLineReader(fd):
-        logger.debug('sending message %r', line)
+        logger.debug('sending line from fd to ws %r', line)
         ws.send_str(line)
         await ws._writer.writer.drain()
 
 
 async def send_from_ws_to_fd(ws, fd):
     async for msg in ws:
         if msg.tp == aiohttp.MsgType.text:
-            logger.debug('got data: %s', msg.data)
+            logger.debug('sending data from ws to fd: %s', msg.data)
             fd.write(msg.data)
         elif msg.tp == aiohttp.MsgType.error:
             logger.error('ws connection closed with exception %s', ws.exception())

pushpull/websocket/gateway.py

@@ -14,35 +14,23 @@
 
 
 async def websocket_rabbitmq_gateway(request):
-
+    name = request.GET.get('name')
+    if not name:
+        return aiohttp.web.Response(body=b'name is required', status=400)
     ws = aiohttp.web.WebSocketResponse()
-    await ws.prepare(request)
     logger.debug('websocket connection open')
     try:
-        # await echo_websocket(ws)
-        async with Exchanger('2') as (amqp_sender, amqp_receiver):
+        await ws.prepare(request)
+        async with Exchanger(name, Exchanger.ROLE_WS) as (amqp_sender, amqp_receiver):
             send_coro = asyncio.ensure_future(send_from_amqp_to_websocket(amqp_receiver, ws))
             receive_coro = asyncio.ensure_future(send_from_websocket_to_amqp(ws, amqp_sender))
             await asyncio.gather(receive_coro, send_coro)
-    except Exception as exc:
-        logger.exception('fatal error')
     finally:
         logger.debug('websocket connection closing')
         await ws.close()
         return ws
 
 
-async def echo_websocket(ws):
-    """For testing only """
-    async for msg in ws:
-        if msg.tp == aiohttp.MsgType.text:
-            logger.debug('echoing data: %s', msg.data)
-            ws.send_str(msg.data)
-        elif msg.tp == aiohttp.MsgType.error:
-            logger.error('ws connection closed with exception %s', ws.exception())
-            return
-
-
 async def send_from_websocket_to_amqp(ws, sender):
     async for msg in ws:
         if msg.tp == aiohttp.MsgType.text:
@@ -56,4 +44,15 @@ async def send_from_websocket_to_amqp(ws, sender):
 async def send_from_amqp_to_websocket(receiver, ws):
     async for data in receiver:
         ws.send_str(data)
-        #await ws._writer.writer.drain()
+        # await ws._writer.writer.drain()
+
+
+async def echo_websocket(ws):
+    """For testing only """
+    async for msg in ws:
+        if msg.tp == aiohttp.MsgType.text:
+            logger.debug('echoing data: %s', msg.data)
+            ws.send_str(msg.data)
+        elif msg.tp == aiohttp.MsgType.error:
+            logger.error('ws connection closed with exception %s', ws.exception())
+            return