From 256ea2bd28290e8f059fd40d251a5841afee0226 Mon Sep 17 00:00:00 2001 From: tsm-odoo Date: Wed, 17 Aug 2022 12:25:59 +0000 Subject: [PATCH] [IMP] bus, core, website: add lifecycle events to websocket server This commit is part of the websocket integration in Odoo. It focuses on implementing websocket lifecycle events. Two lifecycle hooks are currently available: - onopen: called after the websocket opening - onclose: called after the websocket closure. Those callbacks will be passed an environment and the websocket related to the lifecycle event. In order to subscribe to websocket lifecycle events, one must decorate a free functions with either `@Websocket.onopen` or `@Websocket.onclose`. Part-of: odoo/odoo#75510 --- addons/bus/tests/test_websocket_caryall.py | 23 +++++++++ addons/bus/websocket.py | 55 ++++++++++++++++++++-- 2 files changed, 74 insertions(+), 4 deletions(-) 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