---
# local user to be used
-botuser: jabberbot
\ No newline at end of file
+botuser: jabberbot
+# message broker backend: "amqp" (RabbitMQ, default) or "nats" (NATS + JetStream)
+broker_backend: amqp
+# whether to install and run a local NATS server on this host (see tasks/nats_server.yml)
+nats_local: false
+nats_version: "2.14.3"
\ No newline at end of file
scope: user
name: urlbot-worker.service
state: restarted
+
+- name: restart nats-server
+ become: true
+ systemd:
+ name: nats-server.service
+ state: restarted
---
+- name: provision local nats server
+ include_tasks: nats_server.yml
+ when: broker_backend == 'nats' and nats_local | default(false)
+
- name: setup virtualenv for the chatbot
shell:
cmd: "python3 -m venv {{ venv_chatbot }}"
- key: "password"
value: "{{password}}"
- key: "rooms"
- value: "{{rooms | join(', ')}}"
+ value: "{{rooms | join(', ')}},"
- key: "src-url"
value: "{{botrepo_view}}"
- key: "bot_nickname"
value: "{{bot_owner_email}}"
- key: "amqp_uri"
value: "{{ amqp_uri }}"
+ - key: "broker_backend"
+ value: "{{ broker_backend }}"
# TODO: detectlanguage_api_key, sudoers and giphy_key
- key: "giphy_key"
value: "{{giphy_key}}"
- restart worker
- restart chatbot
+- name: set nats configuration
+ lineinfile: dest=~/urlbot-v3/local_config.ini create=yes line="nats_uri = {{ nats_uri }}" regexp="^nats_uri.="
+ when: broker_backend == 'nats'
+ tags: [render_config]
+ notify:
+ - restart worker
+ - restart chatbot
+
- name: create directory for systemd unit files
shell: mkdir -p ~/.config/systemd/user/ creates=~/.config/systemd/user/
- name: unitfile support
template:
- src: "wait-for-rabbitmq.sh"
- dest: "~/wait-for-rabbitmq.sh"
+ src: "{{ 'wait-for-nats.sh' if broker_backend == 'nats' else 'wait-for-rabbitmq.sh' }}"
+ dest: "~/wait-for-broker.sh"
mode: u+x
- name: unitfile
--- /dev/null
+---
+# Installs and runs a local NATS+JetStream server on this host. Only
+# included when broker_backend == 'nats' and nats_local is true - opt-in per
+# host, unlike RabbitMQ which this role has always expected to be
+# pre-provisioned externally.
+
+- name: nats - create system user
+ become: true
+ user:
+ name: nats
+ system: true
+ shell: /usr/sbin/nologin
+ create_home: false
+
+- name: nats - create JetStream storage directory
+ become: true
+ file:
+ path: /var/lib/nats/jetstream
+ state: directory
+ owner: nats
+ group: nats
+ mode: "0750"
+
+- name: nats - determine release architecture
+ set_fact:
+ nats_release_arch: "{{ {
+ 'x86_64': 'amd64',
+ 'aarch64': 'arm64',
+ 'armv7l': 'arm7',
+ 'armv6l': 'arm6',
+ }[ansible_architecture] }}"
+
+- name: nats - download server release
+ become: true
+ get_url:
+ url: "https://github.com/nats-io/nats-server/releases/download/v{{ nats_version }}/nats-server-v{{ nats_version }}-linux-{{ nats_release_arch }}.tar.gz"
+ dest: "/tmp/nats-server-v{{ nats_version }}-linux-{{ nats_release_arch }}.tar.gz"
+ register: nats_download
+
+- name: nats - unpack and install binary
+ become: true
+ unarchive:
+ src: "/tmp/nats-server-v{{ nats_version }}-linux-{{ nats_release_arch }}.tar.gz"
+ dest: /tmp
+ remote_src: true
+ when: nats_download is changed
+
+- name: nats - copy binary into place
+ become: true
+ copy:
+ src: "/tmp/nats-server-v{{ nats_version }}-linux-{{ nats_release_arch }}/nats-server"
+ dest: /usr/local/bin/nats-server
+ mode: "0755"
+ remote_src: true
+ when: nats_download is changed
+ notify:
+ - restart nats-server
+
+- name: nats - install systemd unit
+ become: true
+ template:
+ src: nats-server.service
+ dest: /etc/systemd/system/nats-server.service
+ register: nats_unit
+
+- name: nats - enable and start service
+ become: true
+ systemd:
+ name: nats-server
+ daemon_reload: true
+ enabled: true
+ state: started
--- /dev/null
+[Unit]
+Description=NATS server (JetStream enabled)
+After=network.target
+
+[Service]
+ExecStart=/usr/local/bin/nats-server -js -sd /var/lib/nats/jetstream
+User=nats
+Group=nats
+Restart=always
+
+[Install]
+WantedBy=multi-user.target
Description=jabber bot entertaining and supporting activity on jabber MUCs
[Service]
-ExecStartPre=/home/{{ botuser }}/wait-for-rabbitmq.sh
+ExecStartPre=/home/{{ botuser }}/wait-for-broker.sh
ExecStart={{ venv_chatbot }}/bin/urlbotd-chat
WorkingDirectory=/home/{{ botuser }}/urlbot-v3/
StandardOutput=journal+console
Description=jabber bot entertaining and supporting activity on jabber MUCs
[Service]
-ExecStartPre=/home/{{ botuser }}/wait-for-rabbitmq.sh
+ExecStartPre=/home/{{ botuser }}/wait-for-broker.sh
ExecStart={{ venv_worker }}/bin/urlbotd-worker
WorkingDirectory=/home/{{ botuser }}/urlbot-v3/
StandardOutput=journal+console
--- /dev/null
+#!/bin/bash
+
+while ! echo > /dev/tcp/{{ nats_uri | urlsplit('hostname') }}/{{ nats_uri | urlsplit('port') or 4222 }}; do
+ echo "waiting for remote service..."
+ sleep 5s
+done
dependencies = [
"slixmpp>=1.10.0",
"pika",
+ "nats-py",
"configobj",
"requests",
]
[project.optional-dependencies]
chatbot = []
worker = ["lxml", "plyvel"]
+ai = ["torch", "transformers"]
[project.scripts]
urlbotd-chat = "distbot.bot.bot:run"
# -*- coding: utf-8 -*-
import asyncio
-import functools
import json
import logging
import time
from asyncio import TimerHandle
from typing import List
-import pika
-from pika.channel import Channel
-
-from distbot.common.action import Action, send_action
+from distbot.common.action import Action
+from distbot.common.broker import get_async_broker
from distbot.common.config import conf_set, conf_get
from distbot.common.message import process_message
def __init__(self, bot, queue_name):
self._bot = bot
- self._connection = None
+ self._broker = get_async_broker()
self._channel = None
self._queue_name = queue_name
- self._consumer_tag = None
self.loop = asyncio.get_event_loop()
self.idle_task = self.loop.create_task(self.tick())
- def open_connection(self, connection):
- logger.debug("opened connection for actionworker")
- self._connection = connection
- connection.channel(on_open_callback=self.open_channel)
-
- def open_channel(self, channel: Channel):
- logger.debug("opened channel")
- # self.add_on_channel_close_callback()
- # self.setup_exchange(self.EXCHANGE)
- self._channel = channel
- channel.add_on_close_callback(lambda *_: self._bot.disconnect())
- channel.queue_declare(queue=self._queue_name, durable=True, callback=lambda _: self.start_consuming())
- def cancel(_):
- logger.debug("cancel channel")
- channel.close()
- channel.add_on_cancel_callback(lambda _: cancel)
-
- def start_consuming(self):
- """This method sets up the consumer by first calling
- add_on_cancel_callback so that the object is notified if RabbitMQ
- cancels the consumer. It then issues the Basic.Consume RPC command
- which returns the consumer tag that is used to uniquely identify the
- consumer with RabbitMQ. We keep the value to use it when we want to
- cancel consuming. The on_message method is passed in as a callback pika
- will invoke when a message is fully received.
-
- """
- logger.info('Issuing consumer related RPC commands')
- # self.add_on_cancel_callback()
- self._consumer_tag = self._channel.basic_consume(
- queue=self._queue_name,
- on_message_callback=self.on_message
+ async def start(self):
+ logger.debug("connecting actionworker to durable queue %s", self._queue_name)
+ self._channel = await self._broker.connect_durable_queue(
+ self._queue_name, self.on_message, on_channel_closed=self._bot.disconnect,
)
- logger.info("Setup consumer for actions on ch %s on con %s with tag %s", self._channel, self._connection,
- self._consumer_tag)
+ logger.info("actionworker consuming from %s", self._queue_name)
- def on_message(self, ch, method, properties: pika.spec.BasicProperties, body):
+ def on_message(self, ch, method, properties, body):
logger.info('Received message # %s from %s: %s',
method.delivery_tag, properties.app_id, body)
body = json.loads(body.decode("utf-8"))
break
return action_item
+ def _publish_action(self, action: Action):
+ self._channel.basic_publish(
+ exchange='',
+ routing_key=self._queue_name,
+ body=action.serialize()
+ )
+
def schedule_action(self, action: Action):
logger.info("scheduling event: %s", action.serialize())
if action.mutex and self.find_scheduled_action_by_mutex(action.mutex):
logger.info("not scheduling that event (prevented by mutex)")
raise RuntimeError("not scheduling that event (prevented by mutex)")
- handle: TimerHandle = self.loop.call_later(action.time - time.time(), functools.partial(send_action, self._queue_name), action)
+ handle: TimerHandle = self.loop.call_later(action.time - time.time(), self._publish_action, action)
self.event_list.append(handle)
@staticmethod
conf_set('request_counter', request_counter + 1)
def die(self):
- if self._connection:
- self._connection.close()
+ if self._channel:
+ asyncio.ensure_future(self._broker.close())
import re
import shlex
-import pika
-from pika.adapters.asyncio_connection import AsyncioConnection
-
import slixmpp
from slixmpp.exceptions import IqError, IqTimeout
from slixmpp.stanza import Message
self.action_worker.die()
def initialize_actionworker(self):
- connection = AsyncioConnection(
- pika.URLParameters(conf_get("amqp_uri")),
- on_open_callback=self.action_worker.open_connection,
- on_open_error_callback=lambda con, err: logger.error("Could not connect: %s", err),
- on_close_callback=lambda con, e, x=None: logger.info("closing connection of actionworker (%s)", e)
- )
- logger.debug("connection state: %s", connection.connection_state)
+ asyncio.ensure_future(self.action_worker.start())
def disconnect(self, reconnect=False, wait=None, send_close=True):
logger.info("Stopping all workers...")
from functools import partial
from typing import Optional
-import pika
-import pika.channel
-import pika.exceptions
-
from distbot.common.action import Action, send_action
+from distbot.common.broker import SyncBroker, get_sync_broker
from distbot.common.config import conf_get
logger = logging.getLogger(__name__)
else:
self.usage = "(reaction only)"
self.used_routing_key = None
- self.connection: Optional[pika.BlockingConnection] = None
- self.channel = None
+ self.broker: Optional[SyncBroker] = None
def init_channel(self):
- self.connection = pika.BlockingConnection(pika.URLParameters(conf_get("amqp_uri")))
- self.channel = self.connection.channel()
-
- self.channel.exchange_declare(exchange='classifier', exchange_type='topic')
+ self.broker = get_sync_broker()
+ self.broker.connect()
def init_queue(self):
- result = self.channel.queue_declare(exclusive=True, queue="")
- self.queue = result.method.queue
-
- for binding_key in self.binding_keys:
- logger.info("Registering plugin %s for %s to queue %s", self.get_subclass_name(), binding_key, self.queue)
- self.channel.queue_bind(
- exchange='classifier',
- queue=self.queue,
- routing_key=binding_key
- )
+ self.broker.register_classifier_consumer(self.binding_keys, self.callback)
def callback(
self,
- ch: pika.channel.Channel,
- method: pika.spec.Basic.Deliver,
- properties: pika.spec.BasicProperties,
+ ch,
+ method,
+ properties,
body: bytes,
):
logger.debug("Reacting on %s in %s", str(method.routing_key), self.get_subclass_name())
return self.__class__.__name__
def die(self):
- self.connection.add_callback_threadsafe(callback=self.channel.stop_consuming)
+ self.broker.stop()
def run(self):
self.init_channel()
self.register_plugin()
self.init_queue()
- self.channel.basic_consume(queue=self.queue, on_message_callback=self.callback)
- self.channel.start_consuming()
+ self.broker.run_forever()
def parse_body(self, msg: dict) -> Action | None:
raise NotImplementedError()
def register_plugin(self):
- self.channel.queue_declare(queue='plugin_registry')
- self.channel.basic_publish(
- exchange='',
- routing_key='plugin_registry',
- body=json.dumps(self.get_declaration()).encode("utf-8")
+ self.broker.publish_queue(
+ 'plugin_registry',
+ json.dumps(self.get_declaration()).encode("utf-8"),
+ durable=False,
+ max_age_seconds=120,
)
def get_declaration(self):
import json
from copy import deepcopy
-import pika
-from distbot.common.config import conf_get
+from distbot.common.broker import get_sync_broker
class Action:
def send_action(actionqueue, action: Action):
- connection = pika.BlockingConnection(
- pika.URLParameters(conf_get("amqp_uri"))
- )
- channel = connection.channel()
- channel.queue_declare(queue=actionqueue, durable=True)
-
- channel.basic_publish(
- exchange='',
- routing_key=actionqueue,
- body=action.serialize()
- )
+ broker = get_sync_broker()
+ broker.connect()
+ try:
+ broker.publish_queue(actionqueue, action.serialize())
+ finally:
+ broker.close()
--- /dev/null
+# -*- coding: utf-8 -*-
+"""Broker abstraction: lets the AMQP (RabbitMQ) and NATS backends be selected
+at runtime via the `broker_backend` config key, without call sites caring
+which one is in use.
+
+Two interfaces are needed because the codebase already has two distinct
+connection models:
+
+- `SyncBroker`: blocking, used by one instance per owning thread (a plugin
+ `Worker` thread, or the process-wide singleton used by common/message.py
+ and common/action.py).
+- `AsyncBroker`: native asyncio, used only by action_worker.py, which already
+ runs on bot.py's slixmpp event loop.
+
+Callback signatures deliberately mirror pika's own callback shape
+(`on_message(ch, method, properties, body)`, with `method.routing_key`,
+`method.delivery_tag`, and `ch.basic_ack(delivery_tag=...)`) so that
+`Worker.callback()` and `ActionWorker.on_message()` do not need to change at
+all when the backend changes - only connection/consumer setup differs.
+"""
+from abc import ABC, abstractmethod
+from typing import Callable, List, Protocol
+
+from distbot.common.config import conf_get
+
+OnMessage = Callable[..., None]
+
+
+class SyncBroker(ABC):
+ @abstractmethod
+ def connect(self) -> None: ...
+
+ @abstractmethod
+ def close(self) -> None: ...
+
+ @abstractmethod
+ def publish(self, subject: str, body: bytes) -> None:
+ """Fire-and-forget onto the shared classifier fanout."""
+
+ @abstractmethod
+ def publish_queue(self, queue: str, body: bytes, durable: bool = True, max_age_seconds: int = None) -> None:
+ """Point-to-point publish directly onto a named queue.
+
+ max_age_seconds only matters for backends with a durability layer
+ that needs an explicit retention window (JetStream); ignored
+ otherwise (a plain AMQP queue keeps messages until consumed either
+ way)."""
+
+ @abstractmethod
+ def register_classifier_consumer(self, binding_keys: List[str], on_message: OnMessage) -> None:
+ """Register (does not block) a handler for classifier messages whose
+ subject matches any of binding_keys. Multiple calls accumulate."""
+
+ @abstractmethod
+ def register_queue_consumer(self, queue: str, on_message: OnMessage, max_age_seconds: int = None) -> None:
+ """Register a plain point-to-point queue consumer."""
+
+ @abstractmethod
+ def run_forever(self) -> None:
+ """Blocks the calling thread, dispatching to all registered consumers
+ until stop() is called (from another thread)."""
+
+ @abstractmethod
+ def stop(self) -> None:
+ """Thread-safe; unblocks run_forever()."""
+
+
+class DurableChannel(Protocol):
+ def basic_publish(self, exchange: str, routing_key: str, body: bytes) -> None: ...
+
+
+class AsyncBroker(ABC):
+ @abstractmethod
+ async def connect_durable_queue(
+ self, queue: str, on_message: OnMessage, on_channel_closed: Callable[[], None] = None,
+ ) -> DurableChannel:
+ """Connects, ensures the durable queue exists, starts consuming, and
+ returns a duck-typed channel exposing basic_publish(exchange,
+ routing_key, body) for fire-and-forget publishes onto that queue."""
+
+ @abstractmethod
+ async def close(self) -> None: ...
+
+
+def get_sync_broker() -> SyncBroker:
+ if conf_get("broker_backend") == "nats":
+ from distbot.common.broker_nats import NatsSyncBroker
+ return NatsSyncBroker()
+ from distbot.common.broker_amqp import AmqpSyncBroker
+ return AmqpSyncBroker()
+
+
+def get_async_broker() -> AsyncBroker:
+ if conf_get("broker_backend") == "nats":
+ from distbot.common.broker_nats import NatsAsyncBroker
+ return NatsAsyncBroker()
+ from distbot.common.broker_amqp import AmqpAsyncBroker
+ return AmqpAsyncBroker()
--- /dev/null
+# -*- coding: utf-8 -*-
+import asyncio
+import logging
+from typing import Callable, List
+
+import pika
+from pika.adapters.asyncio_connection import AsyncioConnection
+
+from distbot.common.broker import AsyncBroker, DurableChannel, OnMessage, SyncBroker
+from distbot.common.config import conf_get
+
+logger = logging.getLogger(__name__)
+
+
+class AmqpSyncBroker(SyncBroker):
+ """Thin wrapper around the exact pika calls the codebase used directly
+ before the broker abstraction existed - a behavior-preserving refactor."""
+
+ def __init__(self):
+ self.connection = None
+ self.channel = None
+
+ def connect(self) -> None:
+ self.connection = pika.BlockingConnection(pika.URLParameters(conf_get("amqp_uri")))
+ self.channel = self.connection.channel()
+ self.channel.exchange_declare(exchange='classifier', exchange_type='topic')
+
+ def close(self) -> None:
+ if self.connection:
+ self.connection.close()
+
+ def publish(self, subject: str, body: bytes) -> None:
+ self.channel.basic_publish(exchange='classifier', routing_key=subject, body=body)
+
+ def publish_queue(self, queue: str, body: bytes, durable: bool = True, max_age_seconds: int = None) -> None:
+ # max_age_seconds is a JetStream-only retention knob; AMQP has no
+ # equivalent and keeps messages until consumed regardless.
+ self.channel.queue_declare(queue=queue, durable=durable)
+ self.channel.basic_publish(exchange='', routing_key=queue, body=body)
+
+ def register_classifier_consumer(self, binding_keys: List[str], on_message: OnMessage) -> None:
+ result = self.channel.queue_declare(exclusive=True, queue="")
+ queue = result.method.queue
+ for binding_key in binding_keys:
+ logger.info("Registering consumer for %s on queue %s", binding_key, queue)
+ self.channel.queue_bind(exchange='classifier', queue=queue, routing_key=binding_key)
+ self.channel.basic_consume(queue=queue, on_message_callback=on_message)
+
+ def register_queue_consumer(self, queue: str, on_message: OnMessage, max_age_seconds: int = None) -> None:
+ self.channel.queue_declare(queue=queue)
+ self.channel.basic_consume(queue=queue, on_message_callback=on_message)
+
+ def run_forever(self) -> None:
+ self.channel.start_consuming()
+
+ def stop(self) -> None:
+ self.connection.add_callback_threadsafe(callback=self.channel.stop_consuming)
+
+
+class AmqpAsyncBroker(AsyncBroker):
+ def __init__(self):
+ self._connection = None
+ self._channel = None
+
+ async def connect_durable_queue(
+ self, queue: str, on_message: OnMessage, on_channel_closed: Callable[[], None] = None,
+ ) -> DurableChannel:
+ loop = asyncio.get_event_loop()
+ connected: asyncio.Future = loop.create_future()
+
+ def _on_open(connection):
+ logger.debug("opened connection for actionworker")
+ connection.channel(on_open_callback=_on_channel_open)
+
+ def _on_channel_open(channel):
+ logger.debug("opened channel")
+ self._channel = channel
+ if on_channel_closed:
+ channel.add_on_close_callback(lambda *_: on_channel_closed())
+
+ def _cancel(_):
+ logger.debug("cancel channel")
+ channel.close()
+ channel.add_on_cancel_callback(lambda _: _cancel)
+
+ def _queue_declared(_):
+ channel.basic_consume(queue=queue, on_message_callback=on_message)
+ if not connected.done():
+ connected.set_result(channel)
+ channel.queue_declare(queue=queue, durable=True, callback=_queue_declared)
+
+ def _on_open_error(_conn, err):
+ logger.error("Could not connect: %s", err)
+ if not connected.done():
+ connected.set_exception(RuntimeError(str(err)))
+
+ def _on_close(_conn, exc, _x=None):
+ logger.info("closing connection of actionworker (%s)", exc)
+
+ self._connection = AsyncioConnection(
+ pika.URLParameters(conf_get("amqp_uri")),
+ on_open_callback=_on_open,
+ on_open_error_callback=_on_open_error,
+ on_close_callback=_on_close,
+ )
+ return await connected
+
+ async def close(self) -> None:
+ if self._connection:
+ self._connection.close()
--- /dev/null
+# -*- coding: utf-8 -*-
+"""NATS/JetStream implementation of the broker abstraction (common/broker.py).
+
+Design notes (see the migration plan for the full rationale):
+
+- Binding-key patterns are AMQP topic-exchange syntax baked into ~20 plugins
+ (some, e.g. fun.py's "me.#.<nick>.#", have no native NATS-subject-algebra
+ equivalent - '#' can appear mid-pattern and matches zero-or-more words,
+ while NATS '>' can only be a trailing token and matches one-or-more). So
+ the classifier fanout does NOT attempt subject-algebra translation: every
+ subscriber subscribes to the single broad subject "classifier.>" and
+ filters in-process with distbot.common.routing.matches(), the same regex
+ matcher already used to simulate AMQP routing in the test harness.
+
+- nats-py has no blocking/sync client, but worker.py's thread-per-plugin
+ model doesn't need to change: each NatsSyncBroker instance runs its own
+ asyncio event loop on a dedicated background thread and marshals every
+ call onto it, so the plugin thread itself stays a plain blocking caller.
+
+- Named point-to-point queues (action_processing, plugin_registry) always
+ go through JetStream, not core pub/sub, because core NATS has no
+ buffering for a subscriber that hasn't attached yet - unlike AMQP's real
+ (if possibly non-durable) queues. plugin_registry gets a short max_age so
+ stale registrations from restarted/removed plugins don't linger.
+"""
+import asyncio
+import logging
+import threading
+from typing import Callable, List
+
+import nats
+import nats.js.errors
+
+from distbot.common.broker import AsyncBroker, DurableChannel, OnMessage, SyncBroker
+from distbot.common.config import conf_get
+from distbot.common.routing import matches
+
+logger = logging.getLogger(__name__)
+
+CLASSIFIER_PREFIX = "classifier."
+
+
+def _to_str(value) -> str:
+ return value.decode("utf-8") if isinstance(value, bytes) else value
+
+
+def _to_bytes(value) -> bytes:
+ return value.encode("utf-8") if isinstance(value, str) else value
+
+
+def _stream_name_for(queue: str) -> str:
+ return queue.upper()
+
+
+def _durable_name_for(queue: str) -> str:
+ return f"{queue}-consumer"
+
+
+class _Method:
+ def __init__(self, routing_key, delivery_tag=None):
+ self.routing_key = routing_key
+ self.delivery_tag = delivery_tag
+
+
+class _Properties:
+ app_id = None
+
+
+class _NoAckChannel:
+ """Used for classifier-fanout dispatch: core NATS pub/sub has no ack
+ concept, matching AMQP's fanout-to-exclusive-queue semantics today."""
+
+ def basic_publish(self, exchange: str, routing_key: str, body: bytes) -> None:
+ raise NotImplementedError("classifier dispatch channel is receive-only")
+
+ def basic_ack(self, delivery_tag=None) -> None:
+ pass
+
+
+class _JetStreamAckChannel:
+ def __init__(self, msg):
+ self._msg = msg
+
+ def basic_ack(self, delivery_tag=None) -> None:
+ asyncio.ensure_future(self._msg.ack())
+
+
+async def _ensure_stream(js, queue: str, max_age_seconds: int = None):
+ stream_name = _stream_name_for(queue)
+ config = {"name": stream_name, "subjects": [queue]}
+ if max_age_seconds:
+ config["max_age"] = max_age_seconds
+ try:
+ await js.stream_info(stream_name)
+ except nats.js.errors.NotFoundError:
+ await js.add_stream(**config)
+
+
+class NatsSyncBroker(SyncBroker):
+ """One instance per owning thread. Used both for one-shot connect/publish/
+ close call sites (common/message.py, common/action.py) and for the
+ long-running per-plugin-thread consumer model (worker.py).
+
+ nats-py is asyncio-only, but callers of this class are plain sync code -
+ including, critically, callbacks invoked *from inside* another
+ NatsSyncBroker's own dispatch (Worker.callback() calling send_action(),
+ which builds a brand new NatsSyncBroker of its own). A naive
+ "new_event_loop() + run_until_complete()" per call breaks the moment the
+ calling thread already has a loop running (asyncio forbids nesting), so
+ instead every instance runs its own persistent event loop on a dedicated
+ background thread, and every operation is marshalled onto it via
+ run_coroutine_threadsafe(...).result() - safe to call from any thread,
+ async or not, since it never touches the calling thread's own loop."""
+
+ def __init__(self):
+ self._loop = None
+ self._thread = None
+ self._nc = None
+ self._js = None
+ self._classifier_consumers = []
+ self._queue_consumers = []
+ self._subs = []
+ self._stop_wait_event = None
+
+ def connect(self) -> None:
+ self._loop = asyncio.new_event_loop()
+ self._thread = threading.Thread(target=self._loop.run_forever, daemon=True)
+ self._thread.start()
+ self._call(self._connect())
+
+ def _call(self, coro):
+ return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
+
+ async def _connect(self):
+ self._nc = await nats.connect(conf_get("nats_uri"))
+ self._js = self._nc.jetstream()
+
+ def close(self) -> None:
+ if self._nc is not None:
+ self._call(self._nc.close())
+ if self._loop is not None:
+ self._loop.call_soon_threadsafe(self._loop.stop)
+ self._thread.join(timeout=5)
+ self._loop.close()
+
+ def publish(self, subject: str, body: bytes) -> None:
+ full_subject = CLASSIFIER_PREFIX + _to_str(subject)
+ self._call(self._nc.publish(full_subject, _to_bytes(body)))
+
+ def publish_queue(self, queue: str, body: bytes, durable: bool = True, max_age_seconds: int = None) -> None:
+ self._call(self._publish_queue(queue, body, max_age_seconds))
+
+ async def _publish_queue(self, queue: str, body: bytes, max_age_seconds: int):
+ await _ensure_stream(self._js, queue, max_age_seconds)
+ await self._js.publish(queue, _to_bytes(body))
+
+ def register_classifier_consumer(self, binding_keys: List[str], on_message: OnMessage) -> None:
+ self._classifier_consumers.append((binding_keys, on_message))
+
+ def register_queue_consumer(self, queue: str, on_message: OnMessage, max_age_seconds: int = None) -> None:
+ self._queue_consumers.append((queue, on_message, max_age_seconds))
+
+ def run_forever(self) -> None:
+ self._call(self._start_consumers())
+ self._stop_wait_event = threading.Event()
+ self._stop_wait_event.wait()
+
+ async def _start_consumers(self):
+ if self._classifier_consumers:
+ sub = await self._nc.subscribe(CLASSIFIER_PREFIX + ">", cb=self._on_classifier_message)
+ self._subs.append(sub)
+
+ for queue, handler, max_age_seconds in self._queue_consumers:
+ await _ensure_stream(self._js, queue, max_age_seconds)
+ sub = await self._js.subscribe(
+ queue, durable=_durable_name_for(queue), cb=self._make_queue_callback(handler), manual_ack=True,
+ )
+ self._subs.append(sub)
+
+ async def _on_classifier_message(self, msg):
+ subject = msg.subject[len(CLASSIFIER_PREFIX):]
+ method = _Method(routing_key=subject)
+ for binding_keys, handler in self._classifier_consumers:
+ if any(matches(binding_key, subject) for binding_key in binding_keys):
+ handler(_NoAckChannel(), method, _Properties(), msg.data)
+ break
+
+ def stop(self) -> None:
+ if self._stop_wait_event is not None:
+ self._stop_wait_event.set()
+
+ def _make_queue_callback(self, handler: OnMessage):
+ async def _cb(msg):
+ method = _Method(routing_key=msg.subject)
+ handler(_JetStreamAckChannel(msg), method, _Properties(), msg.data)
+ return _cb
+
+
+class _NatsDurableChannel:
+ """Duck-typed as common/broker.py's DurableChannel: exposes
+ basic_publish(exchange, routing_key, body) so action_worker.py's
+ _publish_action() doesn't need to know which backend it's talking to."""
+
+ def __init__(self, js, queue: str):
+ self._js = js
+ self._queue = queue
+
+ def basic_publish(self, exchange: str, routing_key: str, body: bytes) -> None:
+ asyncio.ensure_future(self._js.publish(self._queue, _to_bytes(body)))
+
+
+class NatsAsyncBroker(AsyncBroker):
+ def __init__(self):
+ self._nc = None
+ self._js = None
+ self._sub = None
+
+ async def connect_durable_queue(
+ self, queue: str, on_message: OnMessage, on_channel_closed: Callable[[], None] = None,
+ ) -> DurableChannel:
+ async def _on_closed():
+ # fires once nats-py gives up reconnecting entirely - the closest
+ # analog to pika's channel-close callback, which today fires on
+ # any lost connection (this app doesn't implement reconnection).
+ on_channel_closed()
+
+ self._nc = await nats.connect(
+ conf_get("nats_uri"),
+ closed_cb=_on_closed if on_channel_closed else None,
+ )
+ self._js = self._nc.jetstream()
+ await _ensure_stream(self._js, queue)
+
+ async def _cb(msg):
+ method = _Method(routing_key=msg.subject)
+ on_message(_JetStreamAckChannel(msg), method, _Properties(), msg.data)
+
+ self._sub = await self._js.subscribe(queue, durable=_durable_name_for(queue), cb=_cb, manual_ack=True)
+ return _NatsDurableChannel(self._js, queue)
+
+ async def close(self) -> None:
+ if self._nc is not None:
+ await self._nc.close()
hist_max_time = integer(default=10*60)
amqp_uri = string(default="amqp://guest:guest@localhost:5672/%2F")
+broker_backend = option("amqp", "nats", default="amqp")
+nats_uri = string(default="nats://localhost:4222")
import shlex
from typing import List
-import pika
-from distbot.common.config import conf_get
+from distbot.common.broker import get_sync_broker
from slixmpp import Message, Presence
from slixmpp.jid import JID
def process_message(routing_key, body: str):
- connection = pika.BlockingConnection(pika.URLParameters(conf_get("amqp_uri")))
- channel = connection.channel()
- channel.exchange_declare(exchange='classifier', exchange_type='topic')
-
- connection = pika.BlockingConnection(
- pika.URLParameters(conf_get("amqp_uri"))
- )
- channel = connection.channel()
-
logger.debug("Processing message body with routing key {}".format(routing_key))
- channel.basic_publish(
- exchange='classifier',
- routing_key=routing_key,
- body=body
- )
- connection.close()
+ broker = get_sync_broker()
+ broker.connect()
+ try:
+ broker.publish(routing_key, body)
+ finally:
+ broker.close()
--- /dev/null
+# -*- coding: utf-8 -*-
+import re
+
+
+def matches(binding_key: str, subject: str) -> bool:
+ """Match a subject against an AMQP-topic-exchange-style binding key.
+
+ '.' separates words, '*' matches exactly one word, '#' matches zero or
+ more words (including at the end, where a naive translation would
+ require at least one word).
+ """
+ regex = binding_key.replace('.', r'\.').replace('*', '[^.]+').replace('#', '.*')
+ # fix .# leading to \..*
+ regex = re.sub(r"\\\.\.\*$", ".*", regex)
+ return re.fullmatch(regex, subject) is not None
import signal
from time import sleep
-import pika
-
from distbot.bot.worker import Worker
from distbot.bot import worker as worker_mod
-from distbot.common.config import conf_get
from distbot.plugins import (
basic, fun, lookup, url, feeds, muc, translation, searx, queue_management, plugin_help,
morse, meta,
+ vote,
extended, bugtracker, bots, bofh, didyouknow,
votepoll,
youtube,
url: url.ALL,
didyouknow: didyouknow.ALL,
youtube: youtube.ALL,
+ vote: vote.ALL,
# debug: debug.ALL
votepoll: votepoll.ALL,
}
def initialize_workers():
- connection = pika.BlockingConnection(
- pika.URLParameters(conf_get("amqp_uri"))
- )
- channel = connection.channel()
- channel.exchange_declare(exchange='topic_command', exchange_type='topic')
- channel.exchange_declare(exchange='topic_parse', exchange_type='topic')
- channel.queue_declare(queue=WORKER_QUEUE, durable=True)
- channel.queue_declare(queue=ACTION_QUEUE, durable=True)
-
for classes in PLUGIN_MODULES.values():
for cls in classes:
try:
def init_queue(self):
super().init_queue()
- self.channel.basic_consume(queue='plugin_registry', on_message_callback=self.callback_plugins)
+ self.broker.register_queue_consumer('plugin_registry', self.callback_plugins, max_age_seconds=120)
def callback_plugins(self, ch, method, properties, body):
body = json.loads(body.decode("utf-8"))
# -*- coding: utf-8 -*-
-import re
from dataclasses import dataclass
from unittest import mock
from distbot.bot.bot import Bot
from distbot.bot.worker import Worker
+from distbot.common.routing import matches
@dataclass
params = get_amqp_callback_params(msg, sender, bot_nick)
for recipient in self.recipients:
for binding_key in recipient.binding_keys:
- if self.matches(binding_key, params.method.routing_key):
+ if matches(binding_key, params.method.routing_key):
recipient.callback(**params.__dict__)
break
- @staticmethod
- def matches(binding_key, message_routing_key):
- regex = binding_key.replace('.', '\\.').replace('*', '[^.]+').replace('#', '.*')
- # fix .# leading to \..*
- regex = re.sub(r"\\.\.\*$", ".*", regex)
- return re.fullmatch(regex, message_routing_key) is not None
-
@pytest.mark.parametrize("bind,msg_routing_key", [
("nick.dice.#", "nick.dice"),
("nick.dice.#", "nick.dice.2"),
])
def test_local_message_broker(bind, msg_routing_key):
- assert LocalMessageBroker.matches(bind, msg_routing_key)
\ No newline at end of file
+ assert matches(bind, msg_routing_key)
\ No newline at end of file
--- /dev/null
+# -*- coding: utf-8 -*-
+import asyncio
+import json
+import time
+from unittest.mock import Mock
+
+import pytest
+
+import distbot.bot.action_worker as aw_module
+from distbot.bot.action_worker import ActionWorker
+from distbot.common.action import Action
+
+
+def _in(seconds=60, **kwargs):
+ return Action(time=time.time() + seconds, **kwargs)
+
+
+@pytest.fixture(autouse=True)
+def mock_conf(monkeypatch):
+ monkeypatch.setattr(aw_module, 'conf_get', Mock(return_value='0'))
+ monkeypatch.setattr(aw_module, 'conf_set', Mock())
+
+
+@pytest.fixture(autouse=True)
+def no_sleep(monkeypatch):
+ monkeypatch.setattr(aw_module.time, 'sleep', Mock())
+
+
+@pytest.fixture()
+def mock_process_message(monkeypatch):
+ m = Mock()
+ monkeypatch.setattr(aw_module, 'process_message', m)
+ return m
+
+
+@pytest.fixture()
+def aw(monkeypatch):
+ loop = asyncio.new_event_loop()
+ monkeypatch.setattr(asyncio, 'get_event_loop', lambda: loop)
+
+ def _create_task(coro):
+ coro.close() # close coroutine without running it to avoid warning
+ return Mock()
+
+ original_create_task = loop.create_task
+ loop.create_task = _create_task
+ worker = ActionWorker(bot=Mock(), queue_name='test_actions')
+ loop.create_task = original_create_task # run_until_complete() needs the real one
+ worker.event_list = [] # give each test its own list (class attr is shared)
+ worker._channel = Mock() # intercept _publish_action calls
+ yield worker
+ loop.close()
+
+
+# --- schedule_action: timer fires on channel ---
+
+def test_schedule_fires_publish_on_channel(aw):
+ aw._channel = Mock()
+ action = _in(seconds=0, mutex='fire')
+ aw.schedule_action(action)
+ aw.loop.run_until_complete(asyncio.sleep(0))
+ aw._channel.basic_publish.assert_called_once_with(
+ exchange='',
+ routing_key='test_actions',
+ body=action.serialize()
+ )
+
+
+# --- schedule_action ---
+
+def test_schedule_adds_handle(aw):
+ aw.schedule_action(_in(mutex='t1'))
+ assert len(aw.event_list) == 1
+
+
+def test_schedule_blocks_duplicate_mutex(aw):
+ aw.schedule_action(_in(mutex='t1'))
+ with pytest.raises(RuntimeError):
+ aw.schedule_action(_in(mutex='t1'))
+
+
+def test_schedule_allows_different_mutexes(aw):
+ aw.schedule_action(_in(mutex='t1'))
+ aw.schedule_action(_in(mutex='t2'))
+ assert len(aw.event_list) == 2
+
+
+def test_schedule_without_mutex_always_allowed(aw):
+ aw.schedule_action(_in())
+ aw.schedule_action(_in())
+ assert len(aw.event_list) == 2
+
+
+# --- unschedule_action ---
+
+def test_unschedule_cancels_handle(aw):
+ aw.schedule_action(_in(mutex='stopper'))
+ handle = aw.event_list[0]
+
+ aw.unschedule_action(Action(mutex='stopper'))
+
+ assert handle.cancelled()
+
+
+def test_unschedule_unknown_mutex_is_noop(aw):
+ aw.schedule_action(_in(mutex='exists'))
+ aw.unschedule_action(Action(mutex='ghost'))
+ assert not aw.event_list[0].cancelled()
+
+
+def test_unschedule_allows_reschedule_after_cancel(aw):
+ aw.schedule_action(_in(mutex='reuse'))
+ aw.unschedule_action(Action(mutex='reuse'))
+ aw.schedule_action(_in(mutex='reuse')) # should not raise
+
+
+# --- run_action: event scheduling ---
+
+def test_run_action_schedules_nested_event(aw):
+ action = Action(msg='set timer', recipient='room@conf', sender='user', event=_in(mutex='ev1'))
+ aw.run_action(action)
+ assert len(aw.event_list) == 1
+
+
+def test_run_action_event_inherits_recipient_from_parent(aw):
+ event = _in(mutex='ev1')
+ action = Action(msg='hi', recipient='room@conf', sender='user', event=event)
+ aw.run_action(action)
+ scheduled = aw.event_list[0]._args[0]
+ assert scheduled.recipient == 'room@conf'
+
+
+def test_run_action_event_inherits_sender_from_parent(aw):
+ event = _in(mutex='ev1')
+ action = Action(msg='hi', recipient='room@conf', sender='user@room', event=event)
+ aw.run_action(action)
+ scheduled = aw.event_list[0]._args[0]
+ assert scheduled.sender == 'user@room'
+
+
+def test_run_action_event_keeps_explicit_recipient(aw):
+ event = _in(mutex='ev1', recipient='explicit@room')
+ action = Action(recipient='parent@room', sender='user', event=event)
+ aw.run_action(action)
+ scheduled = aw.event_list[0]._args[0]
+ assert scheduled.recipient == 'explicit@room'
+
+
+# --- run_action: stop_event ---
+
+def test_run_action_stop_event_cancels_timer(aw):
+ aw.schedule_action(_in(mutex='cancel_me'))
+ handle = aw.event_list[0]
+
+ action = Action(event=Action(stop_event=True, mutex='cancel_me'))
+ aw.run_action(action)
+
+ assert handle.cancelled()
+
+
+# --- run_action: mutex conflict warning ---
+
+def test_run_action_mutex_conflict_prepends_warning(aw):
+ aw.schedule_action(_in(mutex='conflict'))
+ action = Action(msg='original', event=_in(mutex='conflict'))
+ aw.run_action(action)
+ assert 'Warning' in action.msg
+ assert 'original' in action.msg
+
+
+def test_run_action_no_warning_without_msg(aw):
+ aw.schedule_action(_in(mutex='conflict'))
+ action = Action(event=_in(mutex='conflict'))
+ aw.run_action(action) # should not raise, msg stays None
+ assert action.msg is None
+
+
+# --- run_action: command replay ---
+
+def test_run_action_command_replays_to_broker(aw, mock_process_message):
+ action = Action(command='nick.dice.5', recipient='room@conf')
+ aw.run_action(action)
+ mock_process_message.assert_called_once()
+ call_kwargs = mock_process_message.call_args.kwargs
+ assert call_kwargs['routing_key'] == b'nick.dice.5'
+
+
+def test_run_action_command_body_strips_nick_prefix(aw, mock_process_message):
+ action = Action(command='nick.dice.5', recipient='room@conf')
+ aw.run_action(action)
+ body = json.loads(mock_process_message.call_args.kwargs['body'])
+ assert body['body'] == 'dice 5'
+
+
+# --- run_action: message delivery ---
+
+def test_run_action_sends_msg_via_bot(aw):
+ action = Action(msg='hello', recipient='room@conf')
+ aw.run_action(action)
+ aw._bot.echo.assert_called_once_with('hello', 'room@conf')
+
+
+def test_run_action_sends_priv_msg_to_sender(aw):
+ action = Action(priv_msg='whisper', sender='user@room')
+ aw.run_action(action)
+ aw._bot.echo.assert_called_once_with('whisper', 'user@room')
--- /dev/null
+# -*- coding: utf-8 -*-
+import json
+from unittest.mock import Mock
+
+import pytest
+
+import distbot.bot.worker as worker_module
+from distbot.bot.worker import Worker
+from distbot.common.action import Action
+
+
+class _TestWorker(Worker):
+ binding_keys = ['nick.test.#']
+ usage = "testbot: test"
+ returned_action = None
+
+ def parse_body(self, msg):
+ return self.returned_action
+
+
+@pytest.fixture()
+def w():
+ worker = _TestWorker(actionqueue='test_actions')
+ worker.channel = Mock()
+ return worker
+
+
+@pytest.fixture()
+def send_action(monkeypatch):
+ mock_action = Mock()
+ monkeypatch.setattr(worker_module, 'send_action', mock_action)
+ return mock_action
+
+
+def _make_pika_args(routing_key='nick.test.foo', body_from='user@room/User', body_to='room@conf', body_text='test foo'):
+ ch = Mock()
+ method = Mock()
+ method.routing_key = routing_key
+ method.delivery_tag = 42
+ properties = Mock()
+ body = json.dumps({'from': body_from, 'to': body_to, 'body': body_text}).encode()
+ return ch, method, properties, body
+
+
+def _published_action(send_action):
+ args = send_action.call_args.args
+ assert args[0] == 'test_actions'
+ return args[1]
+
+
+def test_callback_publishes_action(w, send_action):
+ w.returned_action = Action(msg='hello')
+ ch, method, properties, body = _make_pika_args()
+
+ w.callback(ch, method, properties, body)
+
+ send_action.assert_called_once()
+ assert _published_action(send_action).msg == 'hello'
+
+
+def test_callback_sets_sender_and_recipient(w, send_action):
+ w.returned_action = Action(msg='hi')
+ ch, method, properties, body = _make_pika_args(body_from='alice@room/Alice', body_to='room@conf')
+
+ w.callback(ch, method, properties, body)
+
+ action = _published_action(send_action)
+ assert action.sender == 'alice@room/Alice'
+ assert action.recipient == 'room@conf'
+
+
+def test_callback_preserves_explicit_recipient(w, send_action):
+ w.returned_action = Action(msg='hi', recipient='custom@room')
+ ch, method, properties, body = _make_pika_args(body_to='other@room')
+
+ w.callback(ch, method, properties, body)
+
+ assert _published_action(send_action).recipient == 'custom@room'
+
+
+def test_callback_no_publish_when_no_action(w, send_action):
+ w.returned_action = None
+ ch, method, properties, body = _make_pika_args()
+
+ w.callback(ch, method, properties, body)
+
+ send_action.assert_not_called()
+
+
+def test_callback_acks_on_action(w, send_action):
+ w.returned_action = Action(msg='hi')
+ ch, method, properties, body = _make_pika_args()
+
+ w.callback(ch, method, properties, body)
+
+ ch.basic_ack.assert_called_once_with(delivery_tag=42)
+
+
+def test_callback_acks_when_no_action(w):
+ w.returned_action = None
+ ch, method, properties, body = _make_pika_args()
+
+ w.callback(ch, method, properties, body)
+
+ ch.basic_ack.assert_called_once_with(delivery_tag=42)
+
+
+def test_callback_acks_even_if_parse_body_raises(w):
+ w.parse_body = Mock(side_effect=RuntimeError("plugin exploded"))
+ ch, method, properties, body = _make_pika_args()
+
+ w.callback(ch, method, properties, body)
+
+ w.channel.basic_publish.assert_not_called()
+ ch.basic_ack.assert_called_once_with(delivery_tag=42)
projects / urlbot-v3.git / commitdiff