Replace uwsgi messaging with oxenmq messages

The mule starts the main OMQ instance, but each worker also has one on
hand that is connected to the main one and can be used to send arbitrary
data (via bencoding) to the mule.

Author: jagerman

sogs/base_config.py

@@ -22,6 +22,10 @@
 # match, without the leading 05 byte).
 OMQ_LISTEN = 'tcp://*:22028'
 
+# Address (usually a unix socket) where the main oxenmq process listens for internal connections
+# from other uwsgi workers (to deliver event notifications).  This must not be publicly accessible!
+OMQ_INTERNAL = 'ipc://./omq.sock'
+
 # The log level.
 LOG_LEVEL = logging.WARNING
 

sogs/cleanup.py

@@ -5,7 +5,6 @@
 from .web import app
 from . import db
 from . import config
-from .signal import Signal
 
 # Cleanup interval, in seconds.
 INTERVAL = 10

sogs/filtration.py

@@ -4,9 +4,9 @@
 
 
 def should_drop_message_with_body(body):
-    body = message_body(body)
     """return true if we should drop a message given its body"""
     if os.path.exists(config.BAD_WORDS_FILE):
+        body = message_body(body)
         with open(config.BAD_WORDS_FILE, 'r') as f:
             for line in f:
                 word = line.rstrip()

sogs/legacy_routes.py

@@ -8,7 +8,7 @@
 from . import config
 from . import http
 from . import filtration
-from .signal import send_signal, Signal
+from .omq import send_mule
 
 import os
 import time
@@ -403,7 +403,7 @@ def handle_legacy_delete_messages(ids=None):
             [room.id, *ids],
         )
 
-    send_signal(Signal.MESSAGE_DELETED)
+    send_mule("messages_deleted", ids)
 
     return jsonify({'status_code': 200})
 

sogs/model.py

@@ -3,7 +3,7 @@
 from . import config
 from . import db
 from . import utils
-from .signal import send_signal, Signal
+from .omq import send_mule
 
 import time
 
@@ -532,7 +532,7 @@ def add_post_to_room(user_id, room_id, data, sig, rate_limit_size=5, rate_limit_
         row = result.fetchone()
         msg = {'timestamp': utils.convert_time(row['posted']), 'server_id': row['id']}
 
-    send_signal(Signal.MESSAGE_POSTED)
+    send_mule("message_posted", msg["server_id"])
 
     return msg
 

sogs/mule.py

@@ -1,35 +1,30 @@
 import uwsgi
 import traceback
 import oxenmq
+from oxenc import bt_deserialize
+import time
 from datetime import timedelta
+import functools
 
 from .web import app
-from .signal import Signal
 from . import cleanup
 from . import config
 from . import crypto
+from .omq import omq
 
 # This is the uwsgi "mule" that handles things not related to serving HTTP requests:
 # - it holds the oxenmq instance (with its own interface into sogs)
 # - it handles cleanup jobs (e.g. periodic deletions)
 
-omq = None
-
 
 def run():
-    if uwsgi.mule_id() != 1:
-        app.logger.critical(
-            "Error: sogs.mule must be the first uwsgi mule (mule1), not mule{}".format(
-                uwsgi.mule_id()
-            )
-        )
-        raise RuntimeError("Invalid uwsgi configuration")
-
     try:
-        setup_omq()
+        app.logger.info("OxenMQ mule started.")
 
-        run_loop()
-    except:
+        while True:
+            time.sleep(1)
+
+    except Exception:
         app.logger.error("mule died via exception:\n{}".format(traceback.format_exc()))
 
 
@@ -38,13 +33,18 @@ def allow_conn(addr, pk, sn):
     return oxenmq.AuthLevel.basic
 
 
+def admin_conn(addr, pk, sn):
+    return oxenmq.AuthLevel.admin
+
+
+def inproc_fail(connid, reason):
+    raise RuntimeError(f"Couldn't connect mule to itself: {reason}")
+
+
 def setup_omq():
-    global omq
-    omq = oxenmq.OxenMQ(
-        privkey=crypto._privkey.encode(),
-        pubkey=crypto.server_pubkey.encode(),
-        log_level=oxenmq.LogLevel.fatal,
-    )
+    global omq, mule_conn
+
+    app.logger.debug("Mule setting up omq")
     if isinstance(config.OMQ_LISTEN, list):
         listen = config.OMQ_LISTEN
     elif config.OMQ_LISTEN is None:
@@ -54,41 +54,52 @@ def setup_omq():
     for addr in listen:
         omq.listen(addr, curve=True, allow_connection=allow_conn)
         app.logger.info(f"OxenMQ listening on {addr}")
+
+    # Internal socket for workers to talk to us:
+    omq.listen(config.OMQ_INTERNAL, curve=False, allow_connection=admin_conn)
+
+    # Periodic database cleanup timer:
     omq.add_timer(cleanup.cleanup, timedelta(seconds=cleanup.INTERVAL))
-    omq.start()
 
+    # Commands other workers can send to us, e.g. for notifications of activity for us to know about
+    worker = omq.add_category("worker", access_level=oxenmq.AuthLevel.admin)
+    worker.add_command("message_posted", message_posted)
+    worker.add_command("messages_deleted", messages_deleted)
+    worker.add_command("message_edited", message_edited)
 
-def run_loop():
-    app.logger.info("mule started!")
+    app.logger.debug("Mule starting omq")
+    omq.start()
 
-    callbacks = {Signal.MESSAGE_POSTED: message_posted, Signal.MESSAGE_DELETED: message_deleted}
+    # Connect mule to itself so that if something the mule does wants to send something to the mule
+    # it will work.  (And so be careful not to recurse!)
+    app.logger.debug("Mule connecting to self")
+    mule_conn = omq.connect_inproc(on_success=None, on_failure=inproc_fail)
 
-    app.logger.info("mule started 2!")
-    while True:
-        app.logger.info("mule started looping")
-        msg = uwsgi.mule_get_msg()
-        app.logger.info("mule ax {}".format(msg))
+
+def log_exceptions(f):
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
         try:
-            sig = Signal(int(msg.decode()))
-        except ValueError:
-            app.logger.error(f"mule received unregistered uwsgi mule message {msg}")
-            continue
+            return f(*args, **kwargs)
+        except Exception as e:
+            app.logger.error(f"{f.__name__} raised exception: {e}")
+            raise
+
+    return wrapper
+
 
-        if sig in callbacks:
-            try:
-                callbacks[sig]()
-            except Exception as e:
-                app.logger.error(
-                    f"An exception occured while mule was processing uwsgi signal {sig}:\n{e}"
-                )
-        else:
-            app.logger.error(f"mule received uwsgi signal {sig} but that signal has no handler!")
-    app.logger.info("mule done")
+@log_exceptions
+def message_posted(m: oxenmq.Message):
+    id = bt_deserialize(m.data()[0])
+    app.logger.warning(f"FIXME: mule -- message posted stub, id={id}")
 
 
-def message_posted():
-    app.logger.warning("FIXME: mule -- message posted stub")
+@log_exceptions
+def messages_deleted(m: oxenmq.Message):
+    ids = bt_deserialize(m.data()[0])
+    app.logger.warning(f"FIXME: mule -- message delete stub, deleted messages: {ids}")
 
 
-def message_deleted():
-    app.logger.warning("FIXME: mule -- message delete stub")
+@log_exceptions
+def message_edited(m: oxenmq.Message):
+    app.logger.warning("FIXME: mule -- message edited stub")

sogs/omq.py

@@ -0,0 +1,63 @@
+# Common oxenmq object; this is used by workers and the oxenmq mule.  We create, but do not start,
+# this pre-forking.
+
+import oxenmq
+from oxenc import bt_serialize
+
+from . import crypto, config
+from .postfork import postfork
+
+omq = None
+mule_conn = None
+
+
+def make_omq():
+    omq = oxenmq.OxenMQ(privkey=crypto._privkey.encode(), pubkey=crypto.server_pubkey.encode())
+
+    # We have multiple workers talking to the mule, so we *must* use ephemeral ids to not replace each
+    # others' connections.
+    omq.ephemeral_routing_id = True
+
+    return omq
+
+
+# Postfork for workers: we start oxenmq and connect to the mule process
+@postfork
+def start_oxenmq():
+    try:
+        import uwsgi
+    except ModuleNotFoundError:
+        return
+
+    global omq, mule_conn
+
+    omq = make_omq()
+
+    if uwsgi.mule_id() != 0:
+        from . import mule
+
+        mule.setup_omq()
+        return
+
+    from .web import app  # Imported here to avoid circular import
+
+    app.logger.debug(f"Starting oxenmq connection to mule in worker {uwsgi.worker_id()}")
+
+    omq.start()
+    app.logger.debug(f"Started, connecting to mule")
+    mule_conn = omq.connect_remote(oxenmq.Address(config.OMQ_INTERNAL))
+
+    app.logger.debug(f"worker {uwsgi.worker_id()} connected to mule OMQ")
+
+
+def send_mule(command, *args, prefix="worker."):
+    """
+    Sends a command to the mule from a worker (or possibly from the mule itself).  The command will
+    be prefixed with "worker." (unless overridden).
+
+    Any args will be bt-serialized and send as message parts.
+    """
+    if prefix:
+        command = prefix + command
+
+    omq.send(mule_conn, command, *(bt_serialize(data) for data in args))