Fix async processing messages

Author: SmileyChris

graphql_ws/aiohttp.py

@@ -44,9 +44,7 @@ async def _handle(self, ws, request_context=None):
             except ConnectionClosedException:
                 break
 
-            connection_context.remember_task(
-                self.on_message(connection_context, message), loop=self.loop
-            )
+            self.on_message(connection_context, message)
         await self.on_close(connection_context)
 
     async def handle(self, ws, request_context=None):

graphql_ws/base_async.py

@@ -81,8 +81,8 @@ def closed(self):
     async def close(self, code):
         ...
 
-    def remember_task(self, task, loop=None):
-        self.pending_tasks.add(asyncio.ensure_future(task, loop=loop))
+    def remember_task(self, task):
+        self.pending_tasks.add(task)
         # Clear completed tasks
         self.pending_tasks -= WeakSet(
             task for task in self.pending_tasks if task.done()
@@ -102,9 +102,9 @@ async def handle(self, ws, request_context=None):
 
     def process_message(self, connection_context, parsed_message):
         task = asyncio.ensure_future(
-            super().process_message(connection_context, parsed_message)
+            super().process_message(connection_context, parsed_message), loop=self.loop
         )
-        connection_context.pending.add(task)
+        connection_context.remember_task(task)
         return task
 
     async def send_message(self, *args, **kwargs):

graphql_ws/websockets_lib.py

@@ -40,9 +40,7 @@ async def _handle(self, ws, request_context):
             except ConnectionClosedException:
                 break
 
-            connection_context.remember_task(
-                self.on_message(connection_context, message), loop=self.loop
-            )
+            self.on_message(connection_context, message)
         await self.on_close(connection_context)
 
     async def handle(self, ws, request_context=None):