Merge pull request #72 from jhjaggars/handle-cancel

adding except blocks for task cancelation

Author: matburt

receptor/connection/base.py

@@ -1,8 +1,7 @@
-import logging
-
 import asyncio
-from collections.abc import AsyncIterator
+import logging
 from abc import abstractmethod, abstractproperty
+from collections.abc import AsyncIterator
 
 from ..messages.envelope import FramedBuffer
 
@@ -24,23 +23,27 @@ async def send(self, bytes_):
 
 
 async def watch_queue(conn, buf):
-    while not conn.closed:
-        try:
-            msg = await asyncio.wait_for(buf.get(), 5.0)
-            if not msg:
+    try:
+        while not conn.closed:
+            try:
+                msg = await asyncio.wait_for(buf.get(), 5.0)
+                if not msg:
+                    return await conn.close()
+            except asyncio.TimeoutError:
+                continue
+            except Exception:
+                logger.exception("watch_queue: error getting data from buffer")
+                continue
+
+            try:
+                await conn.send(msg)
+            except Exception:
+                logger.exception("watch_queue: error received trying to write")
+                await buf.put(msg)
                 return await conn.close()
-        except asyncio.TimeoutError:
-            continue
-        except Exception:
-            logger.exception("watch_queue: error getting data from buffer")
-            continue
-
-        try:
-            await conn.send(msg)
-        except Exception:
-            logger.exception("watch_queue: error received trying to write")
-            await buf.put(msg)
-            return await conn.close()
+    except asyncio.CancelledError:
+        logger.debug("watch_queue: cancel request received")
+        await conn.close()
 
 
 class Worker:
@@ -63,6 +66,8 @@ async def receive(self):
                 if self.conn.closed:
                     break
                 await self.buf.put(msg)
+        except asyncio.CancelledError:
+            logger.debug("receive: cancel request received")
         except Exception:
             logger.exception("receive")
 

receptor/receptor.py

@@ -92,6 +92,9 @@ async def message_handler(self, buf):
         while True:
             try:
                 data = await buf.get()
+            except asyncio.CancelledError:
+                logger.debug("message_handler: cancel request received")
+                break
             except Exception:
                 logger.exception("message_handler")
                 break