diff --git a/addons/bus/tests/test_websocket_caryall.py b/addons/bus/tests/test_websocket_caryall.py index 5c8e70dc7d3..ab913ba8c2e 100644 --- a/addons/bus/tests/test_websocket_caryall.py +++ b/addons/bus/tests/test_websocket_caryall.py @@ -17,6 +17,7 @@ except ImportError: from odoo.api import Environment from odoo.tests import common, new_test_user from .common import WebsocketCase +from .. import websocket as websocket_module from ..models.bus import dispatch from ..websocket import ( CloseCode, @@ -32,7 +33,7 @@ from ..websocket import ( class TestWebsocketCaryall(WebsocketCase): def test_lifecycle_hooks(self): events = [] - with patch.object(Websocket, '_event_callbacks', defaultdict(set)): + with patch.object(Websocket, '_Websocket__event_callbacks', defaultdict(set)): @Websocket.onopen def onopen(env, websocket): # pylint: disable=unused-variable self.assertIsInstance(env, Environment) @@ -51,10 +52,10 @@ class TestWebsocketCaryall(WebsocketCase): self.assertEqual(events, ['open', 'close']) def test_instances_weak_set(self): - with patch.object(Websocket, "_instances", WeakSet()): + with patch.object(websocket_module, "_websocket_instances", WeakSet()): first_ws = self.websocket_connect() second_ws = self.websocket_connect() - self.assertEqual(len(Websocket._instances), 2) + self.assertEqual(len(websocket_module._websocket_instances), 2) first_ws.close(CloseCode.CLEAN) second_ws.close(CloseCode.CLEAN) self.wait_remaining_websocket_connections() @@ -62,7 +63,7 @@ class TestWebsocketCaryall(WebsocketCase): # collected. Stop it now. self._serve_forever_patch.stop() gc.collect() - self.assertEqual(len(Websocket._instances), 0) + self.assertEqual(len(websocket_module._websocket_instances), 0) def test_timeout_manager_no_response_timeout(self): with freeze_time('2022-08-19') as frozen_time: @@ -202,7 +203,7 @@ class TestWebsocketCaryall(WebsocketCase): self.assertEqual(mock.call_args[0][2], client_last_notification_id) def test_no_cursor_when_no_callback_for_lifecycle_event(self): - with patch.object(Websocket, '_event_callbacks', defaultdict(set)): + with patch.object(Websocket, '_Websocket__event_callbacks', defaultdict(set)): with patch('odoo.addons.bus.websocket.acquire_cursor') as mock: self.websocket_connect() self.assertFalse(mock.called) diff --git a/addons/bus/websocket.py b/addons/bus/websocket.py index d7654f7c36e..58164ec157d 100644 --- a/addons/bus/websocket.py +++ b/addons/bus/websocket.py @@ -210,9 +210,11 @@ class CloseFrame(Frame): super().__init__(Opcode.CLOSE, payload) +_websocket_instances = WeakSet() + + class Websocket: - _instances = WeakSet() - _event_callbacks = defaultdict(set) + __event_callbacks = defaultdict(set) # Maximum size for a message in bytes, whether it is sent as one # frame or many fragmented ones. MESSAGE_MAX_SIZE = 2 ** 20 @@ -229,7 +231,8 @@ class Websocket: def __init__(self, sock, session): # Session linked to the current websocket connection. self._session = session - self._socket = sock + self._db = session.db + self.__socket = sock self._close_sent = False self._close_received = False self._timeout_manager = TimeoutManager() @@ -237,19 +240,19 @@ class Websocket: self._incoming_frame_timestamps = deque(maxlen=type(self).RL_BURST) # Used to notify the websocket that bus notifications are # available. - self._notif_sock_w, self._notif_sock_r = socket.socketpair() + self.__notif_sock_w, self.__notif_sock_r = socket.socketpair() self._channels = set() self._last_notif_sent_id = 0 # Websocket start up - self._selector = ( + self.__selector = ( selectors.PollSelector() if odoo.evented and hasattr(selectors, 'PollSelector') else selectors.DefaultSelector() ) - self._selector.register(self._socket, selectors.EVENT_READ) - self._selector.register(self._notif_sock_r, selectors.EVENT_READ) + self.__selector.register(self.__socket, selectors.EVENT_READ) + self.__selector.register(self.__notif_sock_r, selectors.EVENT_READ) self.state = ConnectionState.OPEN - type(self)._instances.add(self) + _websocket_instances.add(self) self._trigger_lifecycle_event(LifecycleEvent.OPEN) # ------------------------------------------------------ @@ -261,7 +264,7 @@ class Websocket: try: readables = { selector_key[0].fileobj for selector_key in - self._selector.select(type(self).INACTIVITY_TIMEOUT) + self.__selector.select(type(self).INACTIVITY_TIMEOUT) } if self._timeout_manager.has_timed_out() and self.state is ConnectionState.OPEN: self.disconnect( @@ -273,9 +276,9 @@ class Websocket: if not readables: self._send_ping_frame() continue - if self._notif_sock_r in readables: + if self.__notif_sock_r in readables: self._dispatch_bus_notifications() - if self._socket in readables: + if self.__socket in readables: message = self._process_next_message() if message is not None: yield message @@ -298,12 +301,12 @@ class Websocket: @classmethod def onopen(cls, func): - cls._event_callbacks[LifecycleEvent.OPEN].add(func) + cls.__event_callbacks[LifecycleEvent.OPEN].add(func) return func @classmethod def onclose(cls, func): - cls._event_callbacks[LifecycleEvent.CLOSE].add(func) + cls.__event_callbacks[LifecycleEvent.CLOSE].add(func) return func def subscribe(self, channels, last): @@ -324,11 +327,11 @@ class Websocket: return readables = { selector_key[0].fileobj for selector_key in - self._selector.select(0) + self.__selector.select(0) } - if self._notif_sock_r not in readables: + if self.__notif_sock_r not in readables: # Send a random bit to mark the socket as readable. - self._notif_sock_w.send(b'x') + self.__notif_sock_w.send(b'x') # ------------------------------------------------------ # PRIVATE METHODS @@ -356,7 +359,7 @@ class Websocket: """ Pull n bytes from the socket """ data = bytearray() while len(data) < n: - received_data = self._socket.recv(n - len(data)) + received_data = self.__socket.recv(n - len(data)) if not received_data: raise ConnectionClosed() data.extend(received_data) @@ -498,7 +501,7 @@ class Websocket: struct.pack('!BBQ', first_byte, 127, payload_length) ) output.extend(frame.payload) - self._socket.sendall(output) + self.__socket.sendall(output) self._timeout_manager.acknowledge_frame_sent(frame) if not isinstance(frame, CloseFrame): return @@ -508,7 +511,7 @@ class Websocket: return self._terminate() # After sending a control frame indicating the connection # should be closed, a peer does not send any further data. - self._selector.unregister(self._notif_sock_r) + self.__selector.unregister(self.__notif_sock_r) def _send_close_frame(self, code, reason=None): """ Send a close frame. """ @@ -525,17 +528,17 @@ class Websocket: def _terminate(self): """ Close the underlying TCP socket. """ with suppress(OSError, TimeoutError): - self._socket.shutdown(socket.SHUT_WR) + self.__socket.shutdown(socket.SHUT_WR) # Call recv until obtaining a return value of 0 indicating # the other end has performed an orderly shutdown. A timeout # is set to ensure the connection will be closed even if # the other end does not close the socket properly. - self._socket.settimeout(1) - while self._socket.recv(4096): + self.__socket.settimeout(1) + while self.__socket.recv(4096): pass - self._selector.unregister(self._socket) - self._selector.close() - self._socket.close() + self.__selector.unregister(self.__socket) + self.__selector.close() + self.__socket.close() self.state = ConnectionState.CLOSED dispatch.unsubscribe(self) self._trigger_lifecycle_event(LifecycleEvent.CLOSE) @@ -602,24 +605,17 @@ class Websocket: raise RateLimitExceededException() self._incoming_frame_timestamps.append(now) - @classmethod - def _kick_all(cls): - """ Disconnect all the websocket instances. """ - for websocket in cls._instances: - if websocket.state is ConnectionState.OPEN: - websocket.disconnect(CloseCode.GOING_AWAY) - def _trigger_lifecycle_event(self, event_type): """ Trigger a lifecycle event that is, call every function registered for this event type. Every callback is given both the environment and the related websocket. """ - if not type(self)._event_callbacks[event_type]: + if not type(self).__event_callbacks[event_type]: return - with closing(acquire_cursor(self._session.db)) as cr: + with closing(acquire_cursor(self._db)) as cr: env = api.Environment(cr, self._session.uid, self._session.context) - for callback in type(self)._event_callbacks[event_type]: + for callback in type(self).__event_callbacks[event_type]: try: service_model.retrying(functools.partial(callback, env, self), env) except Exception: @@ -644,7 +640,7 @@ class Websocket: if session.uid is not None and not check_session(session, env): raise SessionExpiredException() # Mark the notification request as processed. - self._notif_sock_r.recv(1) + self.__notif_sock_r.recv(1) notifications = env['bus.bus']._poll(self._channels, self._last_notif_sent_id) if not notifications: return @@ -838,11 +834,11 @@ class WebsocketConnectionHandler: try: response = cls._get_handshake_response(request.httprequest.headers) socket = request.httprequest.environ['socket'] - response.call_on_close(functools.partial( - cls._serve_forever, - Websocket(socket, request.session), - request.db, - request.httprequest + session, db, httprequest = request.session, request.db, request.httprequest + response.call_on_close(lambda: cls._serve_forever( + Websocket(socket, session), + db, + httprequest, )) # Force save the session. Session must be persisted to handle # WebSocket authentication. @@ -943,4 +939,11 @@ class WebsocketConnectionHandler: _logger.exception("Exception occurred during websocket request handling") -CommonServer.on_stop(Websocket._kick_all) +def _kick_all(): + """ Disconnect all the websocket instances. """ + for websocket in _websocket_instances: + if websocket.state is ConnectionState.OPEN: + websocket.disconnect(CloseCode.GOING_AWAY) + + +CommonServer.on_stop(_kick_all)