Fix #561 pkasyncio.create_task keep ref to avoid GC bug (#562)

Author: robnagler

pykern/api/server.py

@@ -184,6 +184,8 @@ def handle_on_close(self):
         self.destroy()
 
     async def handle_on_message(self, msg):
+        if self._destroyed:
+            return
         m = None
         try:
             m = _ServerMsg(self)
@@ -232,8 +234,8 @@ async def get(self, *args, **kwargs):
 
     async def on_message(self, msg):
         try:
-            # WebSocketHandler only allows one on_message at a time.
-            asyncio.create_task(self.pykern_connection.handle_on_message(msg))
+            # WebSocketHandler only allows one on_message at a time
+            pykern.pkasyncio.create_task(self.pykern_connection.handle_on_message(msg))
         except Exception as e:
             pkdlog("exception={} stack={}", e, pkdexc())
 

pykern/pkasyncio.py

@@ -18,6 +18,8 @@
 
 _cfg = None
 
+_background_tasks = set()
+
 
 class Loop:
     def __init__(self):
@@ -134,7 +136,16 @@ def cfg_port(value):
 
 
 def create_task(coro):
-    return asyncio.create_task(coro)
+    """Create a task
+
+    Keeps a global reference to the task so to avoid the garbage
+    collector running before the task is run.
+    https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
+    """
+    t = asyncio.create_task(coro)
+    _background_tasks.add(t)
+    t.add_done_callback(_background_tasks.discard)
+    return t
 
 
 async def sleep(secs):