fix #545: memory leak after close session

Author: wang0618

pywebio/platform/adaptor/ws.py

@@ -13,19 +13,18 @@
 logger = logging.getLogger(__name__)
 
 
-class _state:
-    # only used in reconnect enabled
+# used to store global state when reconnect enabled
+class _reconnect_state:
     # used to clean up session
     detached_sessions = LRUDict()  # session_id -> detached timestamp. In increasing order of the time
 
     # unclosed and unexpired session
-    # only used in reconnect enabled
     # used to clean up session
     # used to retrieve session by id when new connection
     unclosed_sessions: Dict[str, Session] = {}  # session_id -> session
 
     # the messages that can't deliver to browser when session close due to connection lost
-    undelivered_messages: Dict[str, list] = {}  # session_id -> unhandled message list
+    session_will_messages: Dict[str, list] = {}  # session_id -> unhandled message list
 
     # used to get the active conn in session's callbacks
     active_connections: Dict[str, 'WebSocketConnection'] = {}  # session_id -> WSHandler
@@ -34,24 +33,24 @@ class _state:
 
 
 def set_expire_second(sec):
-    _state.expire_second = max(_state.expire_second, sec)
+    _reconnect_state.expire_second = max(_reconnect_state.expire_second, sec)
 
 
 def clean_expired_sessions():
-    while _state.detached_sessions:
-        session_id, detached_ts = _state.detached_sessions.popitem(last=False)  # 弹出最早过期的session
+    while _reconnect_state.detached_sessions:
+        session_id, detached_ts = _reconnect_state.detached_sessions.popitem(last=False)  # 弹出最早过期的session
 
-        if time.time() < detached_ts + _state.expire_second:
+        if time.time() < detached_ts + _reconnect_state.expire_second:
             # this session is not expired
-            _state.detached_sessions[session_id] = detached_ts  # restore
-            _state.detached_sessions.move_to_end(session_id, last=False)  # move to head
+            _reconnect_state.detached_sessions[session_id] = detached_ts  # restore
+            _reconnect_state.detached_sessions.move_to_end(session_id, last=False)  # move to head
             break
 
         # clean this session
         logger.debug("session %s expired" % session_id)
-        _state.active_connections.pop(session_id, None)
-        _state.undelivered_messages.pop(session_id, None)
-        session = _state.unclosed_sessions.pop(session_id, None)
+        _reconnect_state.active_connections.pop(session_id, None)
+        _reconnect_state.session_will_messages.pop(session_id, None)
+        session = _reconnect_state.unclosed_sessions.pop(session_id, None)
         if session:
             session.close(nonblock=True)
 
@@ -61,7 +60,7 @@ def clean_expired_sessions():
 
 async def session_clean_task():
     global _session_clean_task_started
-    if _session_clean_task_started or not _state.expire_second:
+    if _session_clean_task_started or not _reconnect_state.expire_second:
         return
 
     _session_clean_task_started = True
@@ -72,7 +71,7 @@ async def session_clean_task():
         except Exception:
             logger.exception("Error when clean expired sessions")
 
-        await asyncio.sleep(_state.expire_second // 2)
+        await asyncio.sleep(_reconnect_state.expire_second // 2)
 
 
 class WebSocketConnection(abc.ABC):
@@ -118,30 +117,30 @@ def __init__(self, connection: WebSocketConnection, application, reconnectable:
         if self.session_id in ('NEW', None):  # 初始请求，创建新 Session
             self._init_session(application)
             if reconnectable:
+                _reconnect_state.active_connections[self.session_id] = self.connection
+                _reconnect_state.unclosed_sessions[self.session_id] = self.session
                 # set session id to client, so the client can send it back to server to recover a session when it
                 # resumes form a connection lost
                 connection.write_message(dict(command='set_session_id', spec=self.session_id))
-        elif self.session_id not in _state.unclosed_sessions:  # session is expired
+        elif self.session_id not in _reconnect_state.unclosed_sessions:  # session is expired
             bye_msg = dict(command='close_session')
-            for m in _state.undelivered_messages.get(self.session_id, [bye_msg]):
+            for m in _reconnect_state.session_will_messages.get(self.session_id, [bye_msg]):
                 try:
                     connection.write_message(m)
                 except Exception:
                     logger.exception("Error in sending message via websocket")
-        else:
-            self.session = _state.unclosed_sessions[self.session_id]
-            _state.detached_sessions.pop(self.session_id, None)
-            _state.active_connections[self.session_id] = connection
+        else:  # resumes form a connection lost
+            self.session = _reconnect_state.unclosed_sessions[self.session_id]
+            _reconnect_state.detached_sessions.pop(self.session_id, None)
+            _reconnect_state.active_connections[self.session_id] = connection
             # send the latest messages to client
-            self._send_msg_to_client(self.session)
+            self._send_msg_to_client()
 
         logger.debug('session id: %s' % self.session_id)
 
     def _init_session(self, application):
         session_info = self.connection.make_session_info()
         self.session_id = random_str(24)
-        # todo: only set item when reconnection enabled
-        _state.active_connections[self.session_id] = self.connection
 
         if iscoroutinefunction(application) or isgeneratorfunction(application):
             self.session = CoroutineBasedSession(
@@ -154,12 +153,20 @@ def _init_session(self, application):
                 on_task_command=self._send_msg_to_client,
                 on_session_close=self._close_from_session,
                 loop=self.ioloop)
-        _state.unclosed_sessions[self.session_id] = self.session
 
-    def _send_msg_to_client(self, session):
-        # self.connection may not be active,
-        # here we need the active connection for this session
-        conn = _state.active_connections.get(self.session_id)
+    def _get_active_connection(self) -> Optional[WebSocketConnection]:
+        # when reconnect enabled, the active connection for this session is in _reconnect_state.active_connections,
+        # otherwise, it's self.connection.
+        if self.reconnectable:
+            conn = _reconnect_state.active_connections.get(self.session_id)
+        else:
+            conn = self.connection
+
+        return conn
+
+    def _send_msg_to_client(self, session: Session = None):
+        conn = self._get_active_connection()
+        session = session or self.session
 
         if not conn or conn.closed():
             return
@@ -175,17 +182,13 @@ def _send_msg_to_client(self, session):
                 logger.exception("Error in sending message via websocket")
 
     def _close_from_session(self):
-        session = _state.unclosed_sessions[self.session_id]
-        if self.session_id in _state.active_connections:
-            # send the undelivered messages to client
-            self._send_msg_to_client(session=session)
-        else:
-            _state.undelivered_messages[self.session_id] = session.get_task_commands()
-
-        conn = _state.active_connections.pop(self.session_id, None)
-        _state.unclosed_sessions.pop(self.session_id, None)
+        conn = self._get_active_connection()
         if conn and not conn.closed():
+            self._send_msg_to_client()
             conn.close()
+        elif self.reconnectable:  # no active connection, and reconnect is enabled
+            _reconnect_state.session_will_messages[self.session_id] = self.session.get_task_commands()
+        self.session = None
 
     def send_client_data(self, data):
         if isinstance(data, bytes):
@@ -197,12 +200,14 @@ def send_client_data(self, data):
         self.session.send_client_event(event)
 
     def notify_connection_lost(self):
-        _state.active_connections.pop(self.session_id, None)
-        if not self.reconnectable:
+        logger.debug("WebSocket closed")
+        if not self.reconnectable and self.session:
             # when the connection lost is caused by `on_session_close()`, it's OK to close the session here though.
             # because the `session.close()` is reentrant
             self.session.close(nonblock=True)
-        else:
-            if self.session_id in _state.unclosed_sessions:
-                _state.detached_sessions[self.session_id] = time.time()
-        logger.debug("WebSocket closed")
+            self.session = None  # reset the reference
+            return
+
+        _reconnect_state.active_connections.pop(self.session_id, None)
+        if self.session_id in _reconnect_state.unclosed_sessions:
+            _reconnect_state.detached_sessions[self.session_id] = time.time()

pywebio/session/coroutinebased.py

@@ -153,9 +153,14 @@ def get_task_commands(self):
     def _cleanup(self):
         for t in list(self.coros.values()):  # t.close() may cause self.coros changed size
             t.step(SessionClosedException, throw_exp=True)
+            # in case that the task catch the SessionClosedException, we need to close it manually
             t.close()
         self.coros = {}  # delete session tasks
 
+        # reset the reference, to avoid circular reference
+        self._on_session_close = None
+        self._on_task_command = None
+
     def close(self, nonblock=False):
         """关闭当前Session。由Backend调用"""
         if self.closed():
@@ -295,7 +300,6 @@ def __init__(self, coro, session: CoroutineBasedSession, on_coro_stop=None):
         """
         self.session = session
         self.coro = coro
-        self.coro_id = None
         self.result = None
         self.task_closed = False  # 任务完毕/取消
         self.on_coro_stop = on_coro_stop or (lambda _: None)
@@ -322,21 +326,23 @@ def step(self, result=None, throw_exp=False):
             except StopIteration as e:
                 if len(e.args) == 1:
                     self.result = e.args[0]
-                self.task_closed = True
+                self.close()
                 logger.debug('Task[%s] finished', self.coro_id)
-                self.on_coro_stop(self)
             except Exception as e:
                 if not isinstance(e, SessionException):
                     self.session.on_task_exception()
-                self.task_closed = True
-                self.on_coro_stop(self)
+                self.close()
+
+        if coro_yield is None:
+            return
 
         future = None
         if isinstance(coro_yield, WebIOFuture):
             if coro_yield.coro:
                 future = asyncio.run_coroutine_threadsafe(coro_yield.coro, asyncio.get_event_loop())
-        elif coro_yield is not None:
+        else:
             future = coro_yield
+
         if not self.session.closed() and hasattr(future, 'add_done_callback'):
             future.add_done_callback(self._wakeup)
             self.pending_futures[id(future)] = future
@@ -350,14 +356,18 @@ def close(self):
         if self.task_closed:
             return
 
-        logger.debug('Task[%s] closed', self.coro_id)
+        self.task_closed = True
+
         self.coro.close()
         while self.pending_futures:
             _, f = self.pending_futures.popitem()
             f.cancel()
 
-        self.task_closed = True
         self.on_coro_stop(self)
+        self.on_coro_stop = None  # avoid circular reference
+        self.session = None
+
+        logger.debug('Task[%s] closed', self.coro_id)
 
     def __del__(self):
         if not self.task_closed:

pywebio/session/threadbased.py

@@ -173,6 +173,10 @@ def _trigger_close_event(self):
             self._on_session_close()
 
     def _cleanup(self, nonblock=False):
+        # reset the reference, to avoid circular reference
+        self._on_session_close = None
+        self._on_task_command = None
+
         cls = type(self)
         if not nonblock:
             self.unhandled_task_msgs.wait_empty(8)