rucio-clients 35.7.0__py3-none-any.whl → 37.0.0rc2__py3-none-any.whl

rucio/common/stomp_utils.py

@@ -15,145 +15,409 @@
 """
 Common utility functions for stomp connections
 """
-
+import json
 import logging
+import random
 import socket
+from collections import namedtuple
+from copy import deepcopy
+from functools import partial
 from time import monotonic
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+from stomp import Connection12
+from stomp.exception import ConnectFailedException, NotConnectedException
+from stomp.listener import HeartbeatListener
 
-from stomp import Connection
+from rucio.common.config import config_get, config_get_bool, config_get_float, config_get_int, config_get_list
+from rucio.common.logging import formatted_logger
+from rucio.core.monitor import MetricManager
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Sequence
-    from typing import Any
+    from collections.abc import Iterable, Iterator
+
+    from stomp.connect import Frame
+
+    from rucio.common.types import LoggerFunction
+
+
+METRICS = MetricManager(module=__name__)
+
+
+class Connection(Connection12):
+    """
+    Connection class.
+
+    Wraps Stomp Connection but knows the brokers without accessing
+    hidden variables from the Transport.
+    """
+    def __init__(self, host_and_ports: list[tuple[str, int]], **kwargs):
+        """
+        Initialise.
+
+        Args:
+            host_and_ports: brokers list
+
+        Kwargs:
+            Arguments to pass to the Constructor12 base class.
+        """
+        super().__init__(host_and_ports=host_and_ports, **kwargs)
+        self._brokers = host_and_ports
+
+    @property
+    def brokers(self) -> list[tuple[str, int]]:
+        """
+        List brokers.
+
+        Returns:
+            All assigned brokers in (host, port) format.
+        """
+        return self._brokers
+
 
-    LoggerFunction = Callable[..., Any]
+class ListenerBase(HeartbeatListener):
+    """Listener Base."""
 
+    _logger = formatted_logger(logging.log, 'ListenerBase %s')
 
-def resolve_ips(fqdns: "Sequence[str]", logger: "LoggerFunction" = logging.log):
-    logger(logging.DEBUG, 'resolving dns aliases: %s' % fqdns)
-    resolved = []
-    for fqdn in fqdns:
-        addrinfos = socket.getaddrinfo(fqdn, 0, socket.AF_INET, 0, socket.IPPROTO_TCP)
-        resolved.extend(ai[4][0] for ai in addrinfos)
-    logger(logging.DEBUG, 'dns aliases resolved to %s', resolved)
-    return resolved
+    def __init__(self,
+                 conn: Connection,
+                 logger: "None | LoggerFunction" = None,
+                 **kwargs):
+        """
+        Initialise.
+
+        Args:
+            conn: The connection object that is using this listener
+            logger: Logger to use. Defaults to logging.getLogger(__name__).getChild(__qualname__).
+
+        Kwargs:
+            Arguments to pass to the stomp.ConnectionListener base class.
+        """
+        super().__init__(transport=conn.transport, **kwargs)
+        self._conn = conn
+        if logger is not None:
+            self._logger = logger
+
+    @METRICS.count_it
+    def on_heartbeat_timeout(self):
+        self._conn.disconnect()
+
+    @METRICS.count_it
+    def on_error(self, frame: "Frame"):
+        """
+        on_error
+        """
+        self._logger(logging.ERROR, 'Message receive error: [%s] %s', self._conn.brokers[0][0], frame.body)
+
+
+StompConfig = namedtuple("StompConfig", ('brokers', 'use_ssl', 'port', 'vhost',
+                                         'destination', 'key_file', 'cert_file',
+                                         'username', 'password', 'nonssl_port',
+                                         'reconnect_attempts_max', 'timeout', 'heartbeats'))
 
 
 class StompConnectionManager:
+    """Stomp Connection Manager."""
+
+    _logger = formatted_logger(logging.log, 'StompConnectionManager %s')
+
+    def __init__(self,
+                 config_section: str,
+                 logger: "None | LoggerFunction" = None):
+        """
+        Initialise.
+
+        Args:
+            config_section: The name of the config section for this manager to parse for configuration.
+            logger: logger to use. Defaults to logging.getLogger(__name__).getChild(__qualname__).
+        """
+        if logger is not None:
+            self._logger = logger
+        self._config = self._parse_config(config_section)
+        self._listener_factory = None
+        self._conns = []
+        for broker in self._config.brokers:
+            conn = Connection(host_and_ports=[broker],
+                              vhost=self._config.vhost,
+                              reconnect_attempts_max=self._config.reconnect_attempts_max,
+                              timeout=self._config.timeout,
+                              heartbeats=self._config.heartbeats)
+            if self._config.use_ssl:
+                conn.set_ssl(cert_file=self._config.cert_file, key_file=self._config.key_file)
+            self._conns.append(conn)
+
+    @property
+    def config(self) -> StompConfig:
+        """
+        Get the config.
+
+        Returns:
+            config object.
+        """
+        return deepcopy(self._config)
+
+    def set_listener_factory(self, name: str, listener_cls: type, **kwargs) -> None:
+        """
+        Setup listener factory
+
+        This method will setup a factory to create a name and listener for the arguments to
+        connection.set_listener based on pre-defined argument values.
 
-    def __init__(self):
-        self._brokers = None
-        self._port = None
-        self._use_ssl = None
-        self._vhost = None
-        self._reconnect_attempts = None
-        self._ssl_key_file = None
-        self._timeout = None
-        self._heartbeats = None
+        Args:
+            name: Listener name
+            listener_cls: Listener class.
+        """
+        def create_listener(name, listener_factory, conn):
+            return name, listener_factory(conn=conn)
+        self._listener_factory = partial(create_listener,
+                                         name=name,
+                                         listener_factory=partial(listener_cls, logger=self._logger, **kwargs))
 
-        self._connections = {}
+    def _parse_config(self, config_section: str) -> StompConfig:
+        """
+        Parse config section.
 
-    def is_stalled(self, connection: Connection, *, logger: "LoggerFunction" = logging.log):
-        if not connection.is_connected():
-            return True
+        Args:
+            config_section: The name of the config section for this manager to parse for configuration.
 
-        if self._heartbeats and getattr(connection, 'received_heartbeat') and connection.received_heartbeat:
-            heartbeat_period_seconds = max(0, self._heartbeats[0], self._heartbeats[1]) / 1000
+        Raises:
+            RuntimeError: If cannot parse config sections 'brokers' or 'use_ssl' or if misconfigured.
 
-            if not heartbeat_period_seconds:
-                return False
+        Returns:
+            Stomp manager configuration object.
+        """
+        try:
+            brokers = config_get(config_section, 'brokers')
+        except Exception as exc:
+            self._logger(logging.ERROR, "Could not load brokers from configuration")
+            raise RuntimeError('Could not load brokers from configuration') from exc
 
-            now = monotonic()
-            if connection.received_heartbeat + 10 * heartbeat_period_seconds < now:
-                logger(logging.WARNING, "Stomp connection missed heartbeats for a long time")
-                return True
+        try:
+            use_ssl = config_get_bool(config_section, 'use_ssl')
+        except Exception as exc:
+            self._logger(logging.ERROR, "could not find use_ssl in configuration -- please update your rucio.cfg")
+            raise RuntimeError('could not find use_ssl in configuration -- please update your rucio.cfg') from exc
 
-        return False
+        port = config_get_int(config_section, 'port')
+        vhost = config_get(config_section, 'broker_virtual_host', raise_exception=False)
+        destination = config_get(config_section, "destination")
+        key_file = config_get(config_section, 'ssl_key_file', default=None, raise_exception=False)
+        cert_file = config_get(config_section, 'ssl_cert_file', default=None, raise_exception=False)
+        username = config_get(config_section, 'username', default=None, raise_exception=False)
+        password = config_get(config_section, 'password', default=None, raise_exception=False)
+        nonssl_port = config_get_int(config_section, 'nonssl_port', default=0, raise_exception=False)
+        timeout = config_get_float(config_section, 'timeout', default=None, raise_exception=False)
+        heartbeats = tuple(config_get_list(config_section, 'heartbeats', default=[0., 0.], raise_exception=False))
+        reconnect_attempts = config_get_int(config_section, 'reconnect_attempts', default=100)
+        if use_ssl and (key_file is None or cert_file is None):
+            self._logger(logging.ERROR, "If use_ssl is True in config you must provide both 'ssl_cert_file' "
+                                        "and 'ssl_key_file'")
+            raise RuntimeError("If use_ssl is True in config you must provide both 'ssl_cert_file' and 'ssl_key_file'")
+        if not use_ssl and (username is None or password is None or nonssl_port == 0):
+            self._logger(logging.ERROR, "If use_ssl is False in config you must provide "
+                                        "'username', 'password' and 'nonssl_port'")
+            raise RuntimeError("If use_ssl is False in config you must provide "
+                               "'username', 'password' and 'nonssl_port'")
+        return StompConfig(brokers=self._resolve_host_and_port(brokers, port if use_ssl else nonssl_port),
+                           use_ssl=use_ssl,
+                           port=port, vhost=vhost,
+                           destination=destination, key_file=key_file, cert_file=cert_file,
+                           username=username, password=password, nonssl_port=nonssl_port,
+                           reconnect_attempts_max=reconnect_attempts, timeout=timeout, heartbeats=heartbeats)
 
-    def disconnect(self):
-        for conn in self._connections.values():
+    def _resolve_host_and_port(self, fqdns: "str | Iterable[str]", port: int) -> list[tuple[str, int]]:
+        """
+        Resolve host and port.
+
+        Args:
+            fqdns: fully qualified domain name(s)
+            port: port
+
+        Returns:
+            list of (host, port) tuples.
+        """
+        if isinstance(fqdns, str):
+            fqdns = fqdns.split(',')
+
+        hosts_and_ports = []
+        for fqdn in fqdns:
+            try:
+                addrinfos = socket.getaddrinfo(fqdn.strip(), port, socket.AF_INET, 0, socket.IPPROTO_TCP)
+            except socket.gaierror as exc:
+                self._logger(logging.ERROR, "[broker] Cannot resolve domain name %s (%s)", fqdn.strip(), str(exc))
+                continue
+
+            hosts_and_ports.extend(addrinfo[4] for addrinfo in addrinfos)
+        if not hosts_and_ports:
+            self._logger(logging.WARNING, "[broker] No resolved brokers")
+        return hosts_and_ports
+
+    def _is_stalled(self, conn: Connection) -> bool:
+        """
+        Determine if a connection is stalled.
+
+        Args:
+            conn: The Connection object
+
+        Returns:
+            Whether the connection has stalled.
+        """
+        received_heartbeat = getattr(conn, 'received_heartbeat', None)
+        if received_heartbeat is None or not any(self._config.heartbeats):
+            return False
+
+        heartbeat_period_seconds = max(0, self._config.heartbeats[0], self._config.heartbeats[1]) / 1000
+        if heartbeat_period_seconds == 0.:
+            return False
+
+        now = monotonic()
+        if received_heartbeat + 10 * heartbeat_period_seconds >= now:
+            return False
+
+        return True
+
+    def connect(self) -> "Iterator[Connection]":
+        """
+        Connect.
+
+        Yields:
+            Each connection object after ensuring it's connected.
+        """
+        config = self._config
+        params = {'wait': True, "heartbeats": self._config.heartbeats}
+        self._logger(logging.WARNING, 'heartbeats: %s', self._config.heartbeats)
+        if config.use_ssl:
+            params.update(username=config.username, password=config.password)
+
+        for conn in self._conns:
+            if self._is_stalled(conn):
+                try:
+                    conn.disconnect()
+                except Exception:
+                    self._logger(logging.ERROR, "[broker] Stalled connection could not be disconnected")
             if not conn.is_connected():
-                conn.disconnect()
+                self._logger(logging.INFO, 'connecting to %s:%s', *conn.brokers[0])
+                METRICS.counter('reconnect.{host}').labels(host=conn.brokers[0][0]).inc()
+                if self._listener_factory is not None:
+                    conn.set_listener(*self._listener_factory(conn=conn))
+
+                try:
+                    conn.connect(**params)
+                except ConnectFailedException as error:
+                    self._logger(logging.WARNING, "[broker] Could not deliver message due to "
+                                                  "ConnectFailedException: %s", str(error))
+                    continue
+                except Exception as error:
+                    self._logger(logging.ERROR, "[broker] Could not connect: %s", str(error))
+                    continue
+            try:
+                yield conn
+            except Exception:
+                self._logger(logging.ERROR, "[broker] Error in yielded code, skipping to next connection.")
+
+    def deliver_messages(self, messages: "Iterable[dict[str, Any]]") -> list[int]:
+        """
+        Deliver messages.
+
+        Args:
+            messages: Messages to deliver.
+
+        Returns:
+            delivered message ids, ready for deletion.
+        """
+        config = self._config
+        conn = random.sample(list(self.connect()), 1)[0]
+        to_delete = []
+        for message in messages:
+            try:
+                body = json.dumps({"event_type": str(message["event_type"]).lower(),
+                                   "payload": message["payload"],
+                                   "created_at": str(message["created_at"])})
+            except ValueError:
+                self._logger(logging.ERROR, "[broker] Cannot serialize payload to JSON: %s", str(message["payload"]))
+                to_delete.append(message["id"])
+                continue
+
+            try:
+                conn.send(
+                    body=body,
+                    destination=config.destination,
+                    headers={"persistent": "true",
+                             "event_type": str(message["event_type"]).lower()}
+                )
+                to_delete.append(message["id"])
+            except NotConnectedException as error:
+                self._logger(logging.WARNING, "[broker] Could not deliver message due to NotConnectedException: %s",
+                             str(error))
+                continue
+            except Exception as error:
+                self._logger(logging.ERROR, "[broker] Could not deliver message: %s", str(error))
+                continue
+
+            msg_event_type = str(message["event_type"]).lower()
+            msg_payload = message.get("payload", {})
+            if msg_event_type.startswith("transfer") or msg_event_type.startswith("stagein"):
+                self._logger(logging.DEBUG,
+                             "[broker] - event_type: %s, scope: %s, name: %s, rse: %s, request-id: %s, "
+                             "transfer-id: %s, created_at: %s",
+                             msg_event_type,
+                             msg_payload.get("scope", None),
+                             msg_payload.get("name", None),
+                             msg_payload.get("dst-rse", None),
+                             msg_payload.get("request-id", None),
+                             msg_payload.get("transfer-id", None),
+                             str(message["created_at"]))
+
+            elif msg_event_type.startswith("dataset"):
+                self._logger(logging.DEBUG,
+                             "[broker] - event_type: %s, scope: %s, name: %s, rse: %s, rule-id: %s, created_at: %s)",
+                             msg_event_type,
+                             msg_payload.get("scope", None),
+                             msg_payload.get("name", None),
+                             msg_payload.get("rse", None),
+                             msg_payload.get("rule_id", None),
+                             str(message["created_at"]))
+
+            elif msg_event_type.startswith("deletion"):
+                if "url" not in msg_payload:
+                    msg_payload["url"] = "unknown"
+                self._logger(logging.DEBUG,
+                             "[broker] - event_type: %s, scope: %s, name: %s, rse: %s, url: %s, created_at: %s)",
+                             msg_event_type,
+                             msg_payload.get("scope", None),
+                             msg_payload.get("name", None),
+                             msg_payload.get("rse", None),
+                             msg_payload.get("url", None),
+                             str(message["created_at"]))
+            else:
+                self._logger(logging.DEBUG, "[broker] Other message: %s", message)
+
+        return to_delete
+
+    def subscribe(self, id_: str, ack: str, destination: "None | str" = None, **kwargs) -> None:
+        """
+        Subscribe
+
+        Args:
+            id_: The identifier to uniquely identify the subscription
+            ack: Either auto, client or client-individual
+            destination: The topic or queue to subscribe to. If None then
+                         destination is taken from the rucio config Defaults to None.
+
+        Kwargs:
+            Arguments to pass to the Construction objects subscribe method.
+        """
+        if destination is None:
+            destination = self._config.destination
+        for conn in self.connect():
+            conn.subscribe(destination=destination,
+                           id=id_, ack=ack, **kwargs)
 
-    def re_configure(
-            self,
-            brokers: "Sequence[str]",
-            port: int,
-            use_ssl: bool,
-            vhost,
-            reconnect_attempts: int,
-            ssl_key_file,
-            ssl_cert_file,
-            timeout,
-            heartbeats=(0, 1000),
-            *,
-            logger: "LoggerFunction" = logging.log
-    ) -> tuple[list, list]:
-
-        configuration_changed = any([
-            self._brokers != brokers,
-            self._port != port,
-            self._use_ssl != use_ssl,
-            self._vhost != vhost,
-            self._reconnect_attempts != reconnect_attempts,
-            self._ssl_key_file != ssl_key_file,
-            self._timeout != timeout,
-            self._heartbeats != heartbeats,
-        ])
-        if configuration_changed:
-            self._brokers = brokers
-            self._port = port
-            self._use_ssl = use_ssl
-            self._vhost = vhost
-            self._reconnect_attempts = reconnect_attempts
-            self._ssl_key_file = ssl_key_file
-            self._timeout = timeout
-            self._heartbeats = heartbeats
-
-        current_remotes = set(self._connections)
-        desired_remotes = set((ip, port) for ip in resolve_ips(brokers, logger=logger))
-
-        if configuration_changed:
-            # Re-create all connections
-            to_delete = current_remotes
-            to_create = desired_remotes
-        else:
-            to_delete = current_remotes.difference(desired_remotes)
-            to_create = desired_remotes.difference(current_remotes)
-
-            for remote in current_remotes.intersection(desired_remotes):
-                conn = self._connections[remote]
-
-                if self.is_stalled(conn, logger=logger):
-                    # Re-create stalled connections
-                    to_delete.add(remote)
-                    to_create.add(remote)
-
-        deleted_conns = []
-        for remote in to_delete:
-            conn = self._connections.pop(remote)
-            if conn.is_connected():
+    def disconnect(self) -> None:
+        """Disconnect."""
+        for conn in self._conns:
+            try:
                 conn.disconnect()
-            deleted_conns.append(to_delete)
-
-        created_conns = []
-        for remote in to_create:
-            conn = Connection(
-                host_and_ports=[remote],
-                vhost=vhost,
-                timeout=timeout,
-                heartbeats=heartbeats,
-                reconnect_attempts_max=reconnect_attempts
-            )
-            if use_ssl:
-                conn.set_ssl(key_file=ssl_key_file, cert_file=ssl_cert_file)
-            self._connections[remote] = conn
-            created_conns.append(conn)
-
-        if not to_delete and not to_create:
-            logger(logging.INFO, "Stomp connections didn't change")
-        else:
-            logger(logging.INFO, f"Stomp connections refreshed. Deleted: {list(to_delete)}. Added: {list(to_create)}")
-
-        return created_conns, deleted_conns
+            except Exception:
+                self._logger(logging.ERROR, "[broker] Could not disconnect")

rucio/common/test_rucio_server.py

@@ -36,27 +36,33 @@ def get_scope_and_rses():
 
     :return: A tuple (scope, rses) for the rucio client where scope is mock/test and rses is a list or (None, [None]) if no suitable rse exists.
     """
-    cmd = "rucio list-rses --rses 'test_container_xrd=True'"
+    cmd = "rucio rse list 'test_container_xrd=True'"
     print(cmd)
     exitcode, out, err = execute(cmd)
     print(out, err)
     rses = out.split()
     if len(rses) == 0:
         return None, [None]
-    return 'test', rses
+
+    scope = 'test'
+    account = 'root'
+    cmd = f"rucio scope add {scope} --account {account}"
+    _, out, err = execute(cmd)
+    print(out, err)
+    return scope, rses
 
 
 def delete_rules(did):
     # get the rules for the file
     print('Deleting rules')
-    cmd = "rucio list-rules {0} | grep {0} | cut -f1 -d\\ ".format(did)
+    cmd = "rucio rule list {0} | grep {0} | cut -f1 -d\\ ".format(did)
     print(cmd)
     exitcode, out, err = execute(cmd)
     print(out, err)
     rules = out.split()
     # delete the rules for the file
     for rule in rules:
-        cmd = "rucio delete-rule {0}".format(rule)
+        cmd = "rucio rule remove {0}".format(rule)
         print(cmd)
         exitcode, out, err = execute(cmd)
 
@@ -112,7 +118,7 @@ class TestRucioServer(unittest.TestCase):
         self.assertEqual(exitcode, 0)
 
         # List the files
-        cmd = 'rucio list-files {0}:{1}'.format(self.scope, tmp_dsn)
+        cmd = 'rucio did content list --did {0}:{1}'.format(self.scope, tmp_dsn)
         print(self.marker + cmd)
         exitcode, out, err = execute(cmd)
         print(out)
@@ -120,7 +126,7 @@ class TestRucioServer(unittest.TestCase):
         self.assertEqual(exitcode, 0)
 
         # List the replicas
-        cmd = 'rucio list-file-replicas {0}:{1}'.format(self.scope, tmp_dsn)
+        cmd = 'rucio replica list file {0}:{1}'.format(self.scope, tmp_dsn)
         print(self.marker + cmd)
         exitcode, out, err = execute(cmd)
         print(out)