diff --git a/addons/bus/tests/test_websocket_caryall.py b/addons/bus/tests/test_websocket_caryall.py index 16d50346143..1c57f8be534 100644 --- a/addons/bus/tests/test_websocket_caryall.py +++ b/addons/bus/tests/test_websocket_caryall.py @@ -1,9 +1,12 @@ # Part of Odoo. See LICENSE file for full copyright and licensing details. import gc +from collections import defaultdict from datetime import timedelta from freezegun import freeze_time +from unittest.mock import patch +from odoo.api import Environment from odoo.tests import common from .common import WebsocketCase from ..websocket import ( @@ -18,6 +21,26 @@ from ..websocket import ( @common.tagged('post_install', '-at_install') class TestWebsocketCaryall(WebsocketCase): + def test_lifecycle_hooks(self): + events = [] + with patch.object(Websocket, '_event_callbacks', defaultdict(set)): + @Websocket.onopen + def onopen(env, websocket): # pylint: disable=unused-variable + self.assertIsInstance(env, Environment) + self.assertIsInstance(websocket, Websocket) + events.append('open') + + @Websocket.onclose + def onclose(env, websocket): # pylint: disable=unused-variable + self.assertIsInstance(env, Environment) + self.assertIsInstance(websocket, Websocket) + events.append('close') + + ws = self.websocket_connect() + ws.close(CloseCode.CLEAN) + self.wait_remaining_websocket_connections() + self.assertEqual(events, ['open', 'close']) + def test_instances_weak_set(self): gc.collect() first_ws = self.websocket_connect() diff --git a/addons/bus/websocket.py b/addons/bus/websocket.py index 8f4613efdc7..7b5ad4f962d 100644 --- a/addons/bus/websocket.py +++ b/addons/bus/websocket.py @@ -9,15 +9,18 @@ import struct import selectors import threading import time -from collections import deque -from contextlib import suppress +from collections import defaultdict, deque +from contextlib import closing, suppress from enum import IntEnum from itertools import count from weakref import WeakSet from werkzeug.exceptions import BadRequest, HTTPException +from odoo import api from odoo.http import Response +from odoo.modules.registry import Registry +from odoo.service import model as service_model from odoo.service.server import CommonServer from odoo.tools import config @@ -103,6 +106,16 @@ class RateLimitExceededException(Exception): """ +# ------------------------------------------------------ +# WEBSOCKET LIFECYCLE +# ------------------------------------------------------ + + +class LifecycleEvent(IntEnum): + OPEN = 0 + CLOSE = 1 + + # ------------------------------------------------------ # WEBSOCKET # ------------------------------------------------------ @@ -201,6 +214,7 @@ class CloseFrame(Frame): class Websocket: _instances = WeakSet() + _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 @@ -214,7 +228,9 @@ class Websocket: # How many seconds between each request. RL_DELAY = float(config['websocket_rate_limit_delay']) - def __init__(self, socket): + def __init__(self, socket, session): + # Session linked to the current websocket connection. + self._session = session self._socket = socket self._close_sent = False self._close_received = False @@ -227,6 +243,7 @@ class Websocket: self._incoming_frame_timestamps = deque(maxlen=type(self).RL_BURST) self.state = ConnectionState.OPEN type(self)._instances.add(self) + self._trigger_lifecycle_event(LifecycleEvent.OPEN) # ------------------------------------------------------ # PUBLIC METHODS @@ -284,6 +301,16 @@ class Websocket: else: self._terminate() + @classmethod + def onopen(cls, func): + cls._event_callbacks[LifecycleEvent.OPEN].add(func) + return func + + @classmethod + def onclose(cls, func): + cls._event_callbacks[LifecycleEvent.CLOSE].add(func) + return func + # ------------------------------------------------------ # PRIVATE METHODS # ------------------------------------------------------ @@ -485,6 +512,7 @@ class Websocket: self._selector.close() self._socket.close() self.state = ConnectionState.CLOSED + self._trigger_lifecycle_event(LifecycleEvent.CLOSE) def _handle_control_frame(self, frame): if frame.opcode is Opcode.PING: @@ -547,6 +575,25 @@ class Websocket: 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. + """ + registry = Registry(self._session.db) + with closing(registry.cursor()) as cr: + env = api.Environment(cr, self._session.uid, self._session.context) + for callback in type(self)._event_callbacks[event_type]: + try: + service_model.retrying(functools.partial(callback, env, self), env) + except Exception: + _logger.warning( + 'Error during Websocket %s callback', + LifecycleEvent(event_type).name, + exc_info=True + ) + class TimeoutReason(IntEnum): KEEP_ALIVE = 0 @@ -640,7 +687,7 @@ class WebsocketConnectionHandler: response = cls._get_handshake_response(request.httprequest.headers) response.call_on_close(functools.partial( cls._serve_forever, - Websocket(request.httprequest.environ['socket']), + Websocket(request.httprequest.environ['socket'], request.session), )) return response