Initial commit: Core packages

odoo-bringout-oca-ocb-bus/bus/models/bus.py

@@ -0,0 +1,238 @@
+# -*- coding: utf-8 -*-
+import contextlib
+import datetime
+import json
+import logging
+import math
+import os
+import random
+import selectors
+import threading
+import time
+from psycopg2 import InterfaceError, sql
+
+import odoo
+from odoo import api, fields, models
+from odoo.service.server import CommonServer
+from odoo.tools.misc import DEFAULT_SERVER_DATETIME_FORMAT
+from odoo.tools import date_utils
+
+_logger = logging.getLogger(__name__)
+
+# longpolling timeout connection
+TIMEOUT = 50
+
+# custom function to call instead of NOTIFY postgresql command (opt-in)
+ODOO_NOTIFY_FUNCTION = os.environ.get('ODOO_NOTIFY_FUNCTION')
+
+
+def get_notify_payload_max_length(default=8000):
+    try:
+        length = int(os.environ.get('ODOO_NOTIFY_PAYLOAD_MAX_LENGTH', default))
+    except ValueError:
+        _logger.warning("ODOO_NOTIFY_PAYLOAD_MAX_LENGTH has to be an integer, "
+                        "defaulting to %d bytes", default)
+        length = default
+    return length
+
+
+# max length in bytes for the NOTIFY query payload
+NOTIFY_PAYLOAD_MAX_LENGTH = get_notify_payload_max_length()
+
+
+#----------------------------------------------------------
+# Bus
+#----------------------------------------------------------
+def json_dump(v):
+    return json.dumps(v, separators=(',', ':'), default=date_utils.json_default)
+
+def hashable(key):
+    if isinstance(key, list):
+        key = tuple(key)
+    return key
+
+
+def channel_with_db(dbname, channel):
+    if isinstance(channel, models.Model):
+        return (dbname, channel._name, channel.id)
+    if isinstance(channel, str):
+        return (dbname, channel)
+    return channel
+
+
+def get_notify_payloads(channels):
+    """
+    Generates the json payloads for the imbus NOTIFY.
+    Splits recursively payloads that are too large.
+
+    :param list channels:
+    :return: list of payloads of json dumps
+    :rtype: list[str]
+    """
+    if not channels:
+        return []
+    payload = json_dump(channels)
+    if len(channels) == 1 or len(payload.encode()) < NOTIFY_PAYLOAD_MAX_LENGTH:
+        return [payload]
+    else:
+        pivot = math.ceil(len(channels) / 2)
+        return (get_notify_payloads(channels[:pivot]) +
+                get_notify_payloads(channels[pivot:]))
+
+
+class ImBus(models.Model):
+
+    _name = 'bus.bus'
+    _description = 'Communication Bus'
+
+    channel = fields.Char('Channel')
+    message = fields.Char('Message')
+
+    @api.autovacuum
+    def _gc_messages(self):
+        timeout_ago = datetime.datetime.utcnow()-datetime.timedelta(seconds=TIMEOUT*2)
+        domain = [('create_date', '<', timeout_ago.strftime(DEFAULT_SERVER_DATETIME_FORMAT))]
+        records = self.search(domain, limit=models.GC_UNLINK_LIMIT)
+        if len(records) >= models.GC_UNLINK_LIMIT:
+            self.env.ref('base.autovacuum_job')._trigger()
+        return records.unlink()
+
+    @api.model
+    def _sendmany(self, notifications):
+        channels = set()
+        values = []
+        for target, notification_type, message in notifications:
+            channel = channel_with_db(self.env.cr.dbname, target)
+            channels.add(channel)
+            values.append({
+                'channel': json_dump(channel),
+                'message': json_dump({
+                    'type': notification_type,
+                    'payload': message,
+                })
+            })
+        self.sudo().create(values)
+        if channels:
+            # We have to wait until the notifications are commited in database.
+            # When calling `NOTIFY imbus`, notifications will be fetched in the
+            # bus table. If the transaction is not commited yet, there will be
+            # nothing to fetch, and the websocket will return no notification.
+            @self.env.cr.postcommit.add
+            def notify():
+                with odoo.sql_db.db_connect('postgres').cursor() as cr:
+                    if ODOO_NOTIFY_FUNCTION:
+                        query = sql.SQL("SELECT {}('imbus', %s)").format(sql.Identifier(ODOO_NOTIFY_FUNCTION))
+                    else:
+                        query = "NOTIFY imbus, %s"
+                    payloads = get_notify_payloads(list(channels))
+                    if len(payloads) > 1:
+                        _logger.info("The imbus notification payload was too large, "
+                                     "it's been split into %d payloads.", len(payloads))
+                    for payload in payloads:
+                        cr.execute(query, (payload,))
+
+    @api.model
+    def _sendone(self, channel, notification_type, message):
+        self._sendmany([[channel, notification_type, message]])
+
+    @api.model
+    def _poll(self, channels, last=0):
+        # first poll return the notification in the 'buffer'
+        if last == 0:
+            timeout_ago = datetime.datetime.utcnow()-datetime.timedelta(seconds=TIMEOUT)
+            domain = [('create_date', '>', timeout_ago.strftime(DEFAULT_SERVER_DATETIME_FORMAT))]
+        else:  # else returns the unread notifications
+            domain = [('id', '>', last)]
+        channels = [json_dump(channel_with_db(self.env.cr.dbname, c)) for c in channels]
+        domain.append(('channel', 'in', channels))
+        notifications = self.sudo().search_read(domain)
+        # list of notification to return
+        result = []
+        for notif in notifications:
+            result.append({
+                'id': notif['id'],
+                'message': json.loads(notif['message']),
+            })
+        return result
+
+
+#----------------------------------------------------------
+# Dispatcher
+#----------------------------------------------------------
+
+class BusSubscription:
+    def __init__(self, channels, last):
+        self.last_notification_id = last
+        self.channels = channels
+
+
+class ImDispatch(threading.Thread):
+    def __init__(self):
+        super().__init__(daemon=True, name=f'{__name__}.Bus')
+        self._channels_to_ws = {}
+
+    def subscribe(self, channels, last, db, websocket):
+        """
+        Subcribe to bus notifications. Every notification related to the
+        given channels will be sent through the websocket. If a subscription
+        is already present, overwrite it.
+        """
+        channels = {hashable(channel_with_db(db, c)) for c in channels}
+        for channel in channels:
+            self._channels_to_ws.setdefault(channel, set()).add(websocket)
+        outdated_channels = websocket._channels - channels
+        self._clear_outdated_channels(websocket, outdated_channels)
+        websocket.subscribe(channels, last)
+        with contextlib.suppress(RuntimeError):
+            if not self.is_alive():
+                self.start()
+
+    def unsubscribe(self, websocket):
+        self._clear_outdated_channels(websocket, websocket._channels)
+
+    def _clear_outdated_channels(self, websocket, outdated_channels):
+        """ Remove channels from channel to websocket map. """
+        for channel in outdated_channels:
+            self._channels_to_ws[channel].remove(websocket)
+            if not self._channels_to_ws[channel]:
+                self._channels_to_ws.pop(channel)
+
+    def loop(self):
+        """ Dispatch postgres notifications to the relevant websockets """
+        _logger.info("Bus.loop listen imbus on db postgres")
+        with odoo.sql_db.db_connect('postgres').cursor() as cr, \
+             selectors.DefaultSelector() as sel:
+            cr.execute("listen imbus")
+            cr.commit()
+            conn = cr._cnx
+            sel.register(conn, selectors.EVENT_READ)
+            while not stop_event.is_set():
+                if sel.select(TIMEOUT):
+                    conn.poll()
+                    channels = []
+                    while conn.notifies:
+                        channels.extend(json.loads(conn.notifies.pop().payload))
+                    # relay notifications to websockets that have
+                    # subscribed to the corresponding channels.
+                    websockets = set()
+                    for channel in channels:
+                        websockets.update(self._channels_to_ws.get(hashable(channel), []))
+                    for websocket in websockets:
+                        websocket.trigger_notification_dispatching()
+
+    def run(self):
+        while not stop_event.is_set():
+            try:
+                self.loop()
+            except Exception as exc:
+                if isinstance(exc, InterfaceError) and stop_event.is_set():
+                    continue
+                _logger.exception("Bus.loop error, sleep and retry")
+                time.sleep(TIMEOUT)
+
+# Partially undo a2ed3d3d5bdb6025a1ba14ad557a115a86413e65
+# IMDispatch has a lazy start, so we could initialize it anyway
+# And this avoids the Bus unavailable error messages
+dispatch = ImDispatch()
+stop_event = threading.Event()
+CommonServer.on_stop(stop_event.set)