fix: Massively improve logging context isolation.

src/meta/ipc/client.py

@@ -3,7 +3,7 @@ import asyncio
 import logging
 import pickle
 
-from ..logger import logging_context
+from ..logger import logging_context, log_wrap, set_logging_context
 
 
 logger = logging.getLogger(__name__)
@@ -99,68 +99,70 @@ class AppClient:
         # TODO
         ...
 
+    @log_wrap(action="Req")
     async def request(self, appid, payload: 'AppPayload', wait_for_reply=True):
-        with logging_context(action=f"Req {appid}"):
-            try:
-                if appid not in self.peers:
-                    raise ValueError(f"Peer '{appid}' not found.")
-                logger.debug(f"Sending request to app '{appid}' with payload {payload}")
+        set_logging_context(action=appid)
+        try:
+            if appid not in self.peers:
+                raise ValueError(f"Peer '{appid}' not found.")
+            logger.debug(f"Sending request to app '{appid}' with payload {payload}")
 
-                address = self.peers[appid]
-                reader, writer = await asyncio.open_connection(**address)
+            address = self.peers[appid]
+            reader, writer = await asyncio.open_connection(**address)
 
-                writer.write(payload.encoded())
-                await writer.drain()
-                writer.write_eof()
-                if wait_for_reply:
-                    result = await reader.read()
-                    writer.close()
-                    decoded = payload.route.decode(result)
-                    return decoded
-                else:
-                    return None
-            except Exception:
-                logging.exception(f"Failed to send request to {appid}'")
+            writer.write(payload.encoded())
+            await writer.drain()
+            writer.write_eof()
+            if wait_for_reply:
+                result = await reader.read()
+                writer.close()
+                decoded = payload.route.decode(result)
+                return decoded
+            else:
                 return None
+        except Exception:
+            logging.exception(f"Failed to send request to {appid}'")
+            return None
 
+    @log_wrap(action="Broadcast")
     async def requestall(self, payload, except_self=True, only_my_peers=True):
-        with logging_context(action="Broadcast"):
-            peerlist = self.my_peers if only_my_peers else self.peers
-            results = await asyncio.gather(
-                *(self.request(appid, payload) for appid in peerlist if (appid != self.appid or not except_self)),
-                return_exceptions=True
-            )
-            return dict(zip(self.peers.keys(), results))
+        peerlist = self.my_peers if only_my_peers else self.peers
+        results = await asyncio.gather(
+            *(self.request(appid, payload) for appid in peerlist if (appid != self.appid or not except_self)),
+            return_exceptions=True
+        )
+        return dict(zip(self.peers.keys(), results))
 
     async def handle_request(self, reader, writer):
-        with logging_context(action="SERV"):
-            data = await reader.read()
-            loaded = pickle.loads(data)
-            route, args, kwargs = loaded
+        set_logging_context(action="SERV")
+        data = await reader.read()
+        loaded = pickle.loads(data)
+        route, args, kwargs = loaded
 
-        with logging_context(action=f"SERV {route}"):
-            logger.debug(f"AppClient {self.appid} handling request on route '{route}' with args {args} and kwargs {kwargs}")
+        set_logging_context(action=route)
 
-            if route in self.routes:
-                try:
-                    await self.routes[route].run((reader, writer), args, kwargs)
-                except Exception:
-                    logger.exception(f"Fatal exception during route '{route}'. This should never happen!")
-            else:
-                logger.warning(f"Appclient '{self.appid}' recieved unknown route {route}. Ignoring.")
-            writer.write_eof()
+        logger.debug(f"AppClient {self.appid} handling request on route '{route}' with args {args} and kwargs {kwargs}")
 
+        if route in self.routes:
+            try:
+                await self.routes[route].run((reader, writer), args, kwargs)
+            except Exception:
+                logger.exception(f"Fatal exception during route '{route}'. This should never happen!")
+        else:
+            logger.warning(f"Appclient '{self.appid}' recieved unknown route {route}. Ignoring.")
+        writer.write_eof()
+
+    @log_wrap(stack=("ShardTalk",))
     async def connect(self):
         """
         Start the local peer server.
         Connect to the address server.
         """
-        with logging_context(stack=['ShardTalk']):
-            # Start the client server
-            self._listener = await asyncio.start_server(self.handle_request, **self.address, start_serving=True)
+        # Start the client server
+        self._listener = await asyncio.start_server(self.handle_request, **self.address, start_serving=True)
 
-            logger.info(f"Serving on {self.address}")
-            await self.server_connection()
+        logger.info(f"Serving on {self.address}")
+        await self.server_connection()
 
 
 class AppPayload:

src/meta/ipc/server.py

@@ -4,7 +4,7 @@ import logging
 import string
 import random
 
-from ..logger import log_context, log_app, logging_context, setup_main_logger
+from ..logger import log_context, log_app, setup_main_logger, set_logging_context, log_wrap
 from ..config import conf
 
 logger = logging.getLogger(__name__)
@@ -75,45 +75,45 @@ class AppServer:
         """
         Register and hold a new client connection.
         """
-        with logging_context(action=f"CONN {appid}"):
-            reader, writer = connection
-            # Add the new client
-            self.clients[appid] = (address, connection)
+        set_logging_context(action=f"CONN {appid}")
+        reader, writer = connection
+        # Add the new client
+        self.clients[appid] = (address, connection)
 
-            # Send the new client a client list
-            peers = self.peer_list()
-            writer.write(pickle.dumps(peers))
-            writer.write(b'\n')
-            await writer.drain()
+        # Send the new client a client list
+        peers = self.peer_list()
+        writer.write(pickle.dumps(peers))
+        writer.write(b'\n')
+        await writer.drain()
 
-            # Announce the new client to everyone
-            await self.broadcast('new_peer', (), {'appid': appid, 'address': address})
+        # Announce the new client to everyone
+        await self.broadcast('new_peer', (), {'appid': appid, 'address': address})
 
-            # Keep the connection open until socket closed or EOF (indicating client death)
-            try:
-                await reader.read()
-            finally:
-                # Connection ended or it broke
-                logger.info(f"Lost client '{appid}'")
-                await self.deregister_client(appid)
+        # Keep the connection open until socket closed or EOF (indicating client death)
+        try:
+            await reader.read()
+        finally:
+            # Connection ended or it broke
+            logger.info(f"Lost client '{appid}'")
+            await self.deregister_client(appid)
 
     async def handle_connection(self, reader, writer):
         data = await reader.readline()
         route, args, kwargs = pickle.loads(data)
 
         rqid = short_uuid()
+        
+        set_logging_context(context=f"RQID: {rqid}", action=f"ROUTE {route}")
+        logger.info(f"AppServer handling request on route '{route}' with args {args} and kwargs {kwargs}")
 
-        with logging_context(context=f"RQID: {rqid}", action=f"ROUTE {route}"):
-            logger.info(f"AppServer handling request on route '{route}' with args {args} and kwargs {kwargs}")
-
-            if route in self.routes:
-                # Execute route
-                try:
-                    await self.routes[route]((reader, writer), *args, **kwargs)
-                except Exception:
-                    logger.exception(f"AppServer recieved exception during route '{route}'")
-            else:
-                logger.warning(f"AppServer recieved unknown route '{route}'. Ignoring.")
+        if route in self.routes:
+            # Execute route
+            try:
+                await self.routes[route]((reader, writer), *args, **kwargs)
+            except Exception:
+                logger.exception(f"AppServer recieved exception during route '{route}'")
+        else:
+            logger.warning(f"AppServer recieved unknown route '{route}'. Ignoring.")
 
     def peer_list(self):
         return {appid: address for appid, (address, _) in self.clients.items()}
@@ -122,27 +122,27 @@ class AppServer:
         self.clients.pop(appid, None)
         await self.broadcast('drop_peer', (), {'appid': appid})
 
+    @log_wrap(action="broadcast")
     async def broadcast(self, route, args, kwargs):
-        with logging_context(action="broadcast"):
-            logger.debug(f"Sending broadcast on route '{route}' with args {args} and kwargs {kwargs}.")
-            payload = pickle.dumps((route, args, kwargs))
-            if self.clients:
-                await asyncio.gather(
-                    *(self._send(appid, payload) for appid in self.clients),
-                    return_exceptions=True
-                )
+        logger.debug(f"Sending broadcast on route '{route}' with args {args} and kwargs {kwargs}.")
+        payload = pickle.dumps((route, args, kwargs))
+        if self.clients:
+            await asyncio.gather(
+                *(self._send(appid, payload) for appid in self.clients),
+                return_exceptions=True
+            )
 
     async def message_client(self, appid, route, args, kwargs):
         """
         Send a message to client `appid` along `route` with given arguments.
         """
-        with logging_context(action=f"MSG {appid}"):
-            logger.debug(f"Sending '{route}' to '{appid}' with args {args} and kwargs {kwargs}.")
-            if appid not in self.clients:
-                raise ValueError(f"Client '{appid}' is not connected.")
+        set_logging_context(action=f"MSG {appid}")
+        logger.debug(f"Sending '{route}' to '{appid}' with args {args} and kwargs {kwargs}.")
+        if appid not in self.clients:
+            raise ValueError(f"Client '{appid}' is not connected.")
 
-            payload = pickle.dumps((route, args, kwargs))
-            return await self._send(appid, payload)
+        payload = pickle.dumps((route, args, kwargs))
+        return await self._send(appid, payload)
 
     async def _send(self, appid, payload):
         """
@@ -162,18 +162,19 @@ class AppServer:
 
     async def start(self, address):
         log_app.set("APPSERVER")
-        with logging_context(stack=["SERV"]):
-            server = await asyncio.start_server(self.handle_connection, **address)
-            logger.info(f"Serving on {address}")
-            async with server:
-                await server.serve_forever()
+        set_logging_context(stack=("SERV",))
+        server = await asyncio.start_server(self.handle_connection, **address)
+        logger.info(f"Serving on {address}")
+        async with server:
+            await server.serve_forever()
 
 
 async def start_server():
     setup_main_logger()
     address = {'host': '127.0.0.1', 'port': '5000'}
     server = AppServer()
-    await server.start(address)
+    task = asyncio.create_task(server.start(address))
+    await task
 
 
 if __name__ == '__main__':