rucio-clients 37.0.0rc1__py3-none-any.whl

rucio/common/stomp_utils.py

@@ -0,0 +1,423 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+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, Any
+
+from stomp import Connection12
+from stomp.exception import ConnectFailedException, NotConnectedException
+from stomp.listener import HeartbeatListener
+
+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 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
+
+
+class ListenerBase(HeartbeatListener):
+    """Listener Base."""
+
+    _logger = formatted_logger(logging.log, 'ListenerBase %s')
+
+    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.
+
+        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))
+
+    def _parse_config(self, config_section: str) -> StompConfig:
+        """
+        Parse config section.
+
+        Args:
+            config_section: The name of the config section for this manager to parse for configuration.
+
+        Raises:
+            RuntimeError: If cannot parse config sections 'brokers' or 'use_ssl' or if misconfigured.
+
+        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
+
+        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
+
+        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 _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():
+                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 disconnect(self) -> None:
+        """Disconnect."""
+        for conn in self._conns:
+            try:
+                conn.disconnect()
+            except Exception:
+                self._logger(logging.ERROR, "[broker] Could not disconnect")

rucio/common/stopwatch.py

@@ -0,0 +1,55 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import time
+from typing import Optional
+
+
+class Stopwatch:
+    """Stopwatch to measure time durations.
+
+    Note: The stopwatch is started on initialization.
+    """
+
+    _t_start: float
+    _t_end: Optional[float]
+
+    def __init__(self) -> None:
+        self.restart()
+
+    def _now(self) -> float:
+        # TODO: change to time.monotonic_ns() if python 3.6 support is dropped.
+        return time.monotonic()
+
+    def restart(self) -> None:
+        """Resets and starts the stopwatch."""
+        self._t_start = self._now()
+        self._t_end = None
+
+    def stop(self) -> None:
+        """Stops the stopwatch."""
+        self._t_end = self._now()
+
+    @property
+    def elapsed(self) -> float:
+        """Returns the total number of elapsed seconds."""
+        if self._t_end is None:
+            return self._now() - self._t_start
+        else:
+            return self._t_end - self._t_start
+
+    def __float__(self) -> float:
+        """Returns the total number of elapsed seconds."""
+        return self.elapsed

rucio/common/test_rucio_server.py

@@ -0,0 +1,154 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from os import remove
+from os.path import basename
+
+from rucio.common.utils import execute
+from rucio.common.utils import generate_uuid as uuid
+
+
+def file_generator(size=2048, namelen=10):
+    """ Create a bogus file and returns it's name.
+    :param size: size in bytes
+    :returns: The name of the generated file.
+    """
+    fn = '/tmp/rucio_testfile_' + uuid()
+    execute('dd if=/dev/urandom of={0} count={1} bs=1'.format(fn, size))
+    return fn
+
+
+def get_scope_and_rses():
+    """
+    Check if xrd containers rses for xrootd are available in the testing environment.
+
+    :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 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]
+
+    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 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 rule remove {0}".format(rule)
+        print(cmd)
+        exitcode, out, err = execute(cmd)
+
+
+class TestRucioServer(unittest.TestCase):
+
+    def setUp(self):
+        self.marker = '$ > '
+        self.scope, self.rses = get_scope_and_rses()
+        self.rse = self.rses[0]
+        self.generated_dids = []
+
+    def tearDown(self):
+        for did in self.generated_dids:
+            delete_rules(did)
+        self.generated_dids = []
+
+    def test_ping(self):
+        """CLIENT (USER): rucio ping"""
+        cmd = 'rucio ping'
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        self.assertEqual(exitcode, 0)
+
+    def test_whoami(self):
+        """CLIENT (USER): rucio whoami"""
+        cmd = 'rucio whoami'
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out, err)
+        self.assertEqual(exitcode, 0)
+
+    def test_upload_download(self):
+        """CLIENT(USER): rucio upload files to dataset/download dataset"""
+        if self.rse is None:
+            return
+
+        tmp_file1 = file_generator()
+        tmp_file2 = file_generator()
+        tmp_file3 = file_generator()
+        tmp_dsn = 'tests.rucio_client_test_server_' + uuid()
+
+        # Adding files to a new dataset
+        cmd = 'rucio upload --rse {0} --scope {1} {2} {3} {4} {1}:{5}'.format(self.rse, self.scope, tmp_file1, tmp_file2, tmp_file3, tmp_dsn)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out)
+        print(err)
+        remove(tmp_file1)
+        remove(tmp_file2)
+        remove(tmp_file3)
+        self.assertEqual(exitcode, 0)
+
+        # List the files
+        cmd = 'rucio did content list --did {0}:{1}'.format(self.scope, tmp_dsn)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out)
+        print(err)
+        self.assertEqual(exitcode, 0)
+
+        # List the replicas
+        cmd = 'rucio replica list file {0}:{1}'.format(self.scope, tmp_dsn)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out)
+        print(err)
+        self.assertEqual(exitcode, 0)
+
+        # Downloading dataset
+        cmd = 'rucio download --dir /tmp/ {0}:{1}'.format(self.scope, tmp_dsn)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(out)
+        print(err)
+        # The files should be there
+        cmd = 'ls /tmp/{0}/rucio_testfile_*'.format(tmp_dsn)
+        print(self.marker + cmd)
+        exitcode, out, err = execute(cmd)
+        print(err, out)
+        self.assertEqual(exitcode, 0)
+
+        # cleaning
+        remove('/tmp/{0}/'.format(tmp_dsn) + basename(tmp_file1))
+        remove('/tmp/{0}/'.format(tmp_dsn) + basename(tmp_file2))
+        remove('/tmp/{0}/'.format(tmp_dsn) + basename(tmp_file3))
+        added_dids = ['{0}:{1}'.format(self.scope, did) for did in (basename(tmp_file1), basename(tmp_file2), basename(tmp_file3), tmp_dsn)]
+        self.generated_dids += added_dids