Improved cleanup handling after WSClose is received from server

Author: Codemation

easyrpc/proxy.py

@@ -4,17 +4,19 @@
 import asyncio
 from traceback import format_exc 
 from concurrent.futures._base import CancelledError
+from aiohttp._websocket.models import WSMessage, WSMsgType
 
 from aiohttp import ClientSession
 import aiohttp
+
 from easyrpc.register import (
     create_proxy_from_config, 
     Coroutine, 
     Generator, 
     AsyncGenerator, 
     async_generator_asend
 )
-from easyrpc.auth import encode
+from easyrpc.auth import encode, decode
 from easyrpc.origin import Origin
 from easyrpc.generator import RpcGenerator
 from easyrpc.exceptions import (
@@ -250,6 +252,7 @@ async def get_origin_registered_functions(self):
             self.server.origin(func, namespace=self.namespace)
 
     async def cleanup_proxy_session(self):
+        self.log.warning(f"cleanup_proxy_session called")
         if not self.session_id in self.client_connections:
             return
         try:
@@ -328,14 +331,24 @@ async def ws_receiver():
             try:
                 while True:
                     message = await ws.receive()
+                    self.log.debug(f"ws_receiver got message: {message}")
+
+                    if message.type == WSMsgType.CLOSE:
+                        self.log.info(f"Server sent WSCLOSE")
+                        break
+
                     if message.data == None:
                         break
 
                     if self.serialization == 'json':
                         if 'error' in message.data and not 'ws_action' in message.data:
                             break
 
-                    message = self.deserialize(message.data)
+                    self.log.info(message.data)
+                    try:
+                        message: WSMessage = self.deserialize(message.data)
+                    except Exception as e:
+                        self.log.warning(f"error deserializing message: {repr(e)} - message: {message.data}")
 
                     if not 'ws_action' in message:
                         continue
@@ -433,9 +446,9 @@ async def ws_receiver():
                             })
 
             except Exception as e:
-                if not isinstance(e, CancelledError):
-                    self.log.exception(f"error with ws_receiver")
-            await self.cleanup_proxy_session()
+                self.log.info(f"ws_receiver exiting: reason - {repr(e)}")
+            finally:
+                await self.cleanup_proxy_session()
         return ws_receiver
     async def get_proxy_ws_session(self):
         """