rucio 35.7.0__py3-none-any.whl

rucio/daemons/cache/consumer.py

@@ -0,0 +1,197 @@
+# 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.
+
+"""
+Cache consumer is a daemon to retrieve rucio cache operation information to synchronize rucio catalog.
+"""
+
+import json
+import logging
+import threading
+import time
+from traceback import format_exc
+from typing import TYPE_CHECKING, Optional
+
+import rucio.db.sqla.util
+from rucio.common import exception
+from rucio.common.config import config_get, config_get_bool, config_get_int, config_get_list
+from rucio.common.logging import formatted_logger, setup_logging
+from rucio.common.stomp_utils import StompConnectionManager
+from rucio.common.types import InternalScope, LoggerFunction
+from rucio.core.monitor import MetricManager
+from rucio.core.rse import get_rse_id
+from rucio.core.volatile_replica import add_volatile_replicas, delete_volatile_replicas
+
+if TYPE_CHECKING:
+    from types import FrameType
+
+    from stomp import Connection
+    from stomp.utils import Frame
+
+logging.getLogger("stomp").setLevel(logging.CRITICAL)
+
+METRICS = MetricManager(module=__name__)
+GRACEFUL_STOP = threading.Event()
+DAEMON_NAME = 'cache-consumer'
+
+
+class AMQConsumer:
+    """
+    class Consumer
+    """
+
+    def __init__(
+            self,
+            broker: str,
+            conn: "Connection",
+            logger: "LoggerFunction"
+    ):
+        """
+        __init__
+        """
+        self.__broker = broker
+        self.__conn = conn
+        self.__logger = logger
+
+    @METRICS.count_it
+    def on_heartbeat_timeout(self) -> None:
+        self.__conn.disconnect()
+
+    @METRICS.count_it
+    def on_error(self, frame: "Frame") -> None:
+        """
+        on_error
+        """
+        self.__logger(logging.ERROR, 'Message receive error: [%s] %s' % (self.__broker, frame.body))
+
+    @METRICS.count_it
+    def on_message(self, frame: "Frame") -> None:
+        """
+        on_message
+        """
+        try:
+            msg = json.loads(frame.body)  # type: ignore
+            self.__logger(logging.DEBUG, 'Message received: %s ' % msg)
+            if isinstance(msg, dict) and 'operation' in msg.keys():
+                for f in msg['files']:
+                    f['scope'] = InternalScope(f['scope'])
+                if 'rse_id' in msg:
+                    rse_id = msg['rse_id']
+                else:
+                    rse_id = get_rse_id(rse=msg['rse'], vo=msg.get('vo', 'def'))
+
+                rse_vo_str = msg['rse']
+                if 'vo' in msg and msg['vo'] != 'def':
+                    rse_vo_str = '{} on {}'.format(rse_vo_str, msg['vo'])
+                if msg['operation'] == 'add_replicas':
+                    self.__logger(logging.INFO, 'add_replicas to RSE %s: %s ' % (rse_vo_str, str(msg['files'])))
+                    add_volatile_replicas(rse_id=rse_id, replicas=msg['files'])
+                elif msg['operation'] == 'delete_replicas':
+                    self.__logger(logging.INFO, 'delete_replicas to RSE %s: %s ' % (rse_vo_str, str(msg['files'])))
+                    delete_volatile_replicas(rse_id=rse_id, replicas=msg['files'])
+            else:
+                self.__logger(logging.DEBUG, 'Check failed: %s %s '
+                              % (isinstance(msg, dict), 'operation' in msg.keys()))
+        except:
+            self.__logger(logging.ERROR, str(format_exc()))
+
+
+def consumer(id_: int, num_thread: int = 1) -> None:
+    """
+    Main loop to consume messages from the Rucio Cache producer.
+    """
+
+    logger = formatted_logger(logging.log, DAEMON_NAME + ' %s')
+
+    logger(logging.INFO, 'Rucio Cache consumer starting')
+
+    brokers = config_get_list('messaging-cache', 'brokers')
+
+    use_ssl = config_get_bool('messaging-cache', 'use_ssl', default=True, raise_exception=False)
+    if not use_ssl:
+        username = config_get('messaging-cache', 'username')
+        password = config_get('messaging-cache', 'password')
+    destination = config_get('messaging-cache', 'destination')
+    subscription_id = 'rucio-cache-messaging'
+
+    vhost = config_get('messaging-cache', 'broker_virtual_host', raise_exception=False)
+    port = config_get_int('messaging-cache', 'port')
+    reconnect_attempts = config_get_int('messaging-cache', 'reconnect_attempts', default=100)
+    ssl_key_file = config_get('messaging-cache', 'ssl_key_file', raise_exception=False)
+    ssl_cert_file = config_get('messaging-cache', 'ssl_cert_file', raise_exception=False)
+
+    stomp_conn_mngr = StompConnectionManager()
+    conns, _ = stomp_conn_mngr.re_configure(
+        brokers=brokers,
+        port=port,
+        use_ssl=use_ssl,
+        vhost=vhost,
+        reconnect_attempts=reconnect_attempts,
+        ssl_key_file=ssl_key_file,
+        ssl_cert_file=ssl_cert_file,
+        timeout=None,
+        logger=logger
+    )
+
+    logger(logging.INFO, 'consumer started')
+
+    while not GRACEFUL_STOP.is_set():
+        for conn in conns:
+            if not conn.is_connected():
+                host_port = conn.transport._Transport__host_and_ports[0]
+
+                logger(logging.INFO, 'connecting to %s' % host_port[0])
+                METRICS.counter('reconnect.{host}').labels(host=host_port[0]).inc()
+                conn.set_listener('rucio-cache-consumer', AMQConsumer(broker=host_port, conn=conn, logger=logger))
+                if not use_ssl:
+                    conn.connect(username, password)
+                else:
+                    conn.connect()
+
+                conn.subscribe(destination=destination, ack='auto', id=subscription_id)
+        time.sleep(1)
+
+    logger(logging.INFO, 'graceful stop requested')
+    stomp_conn_mngr.disconnect()
+    logger(logging.INFO, 'graceful stop done')
+
+
+def stop(signum: Optional[int] = None, frame: Optional["FrameType"] = None) -> None:
+    """
+    Graceful exit.
+    """
+
+    GRACEFUL_STOP.set()
+
+
+def run(num_thread: int = 1) -> None:
+    """
+    Starts up the rucio cache consumer thread
+    """
+    setup_logging(process_name=DAEMON_NAME)
+
+    if rucio.db.sqla.util.is_old_db():
+        raise exception.DatabaseException('Database was not updated, daemon won\'t start')
+
+    logging.info('starting consumer thread')
+    threads = [threading.Thread(target=consumer, kwargs={'id_': i, 'num_thread': num_thread})
+               for i in range(0, num_thread)]
+
+    [t.start() for t in threads]
+
+    logging.info('waiting for interrupts')
+
+    # Interruptible joins require a timeout.
+    while threads[0].is_alive():
+        [t.join(timeout=3.14) for t in threads]

rucio/daemons/common.py

@@ -0,0 +1,415 @@
+# 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 datetime
+import functools
+import logging
+import os
+import queue
+import socket
+import threading
+import time
+from collections.abc import Callable, Generator, Iterator, Sequence
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, Union
+
+from rucio.common.logging import formatted_logger
+from rucio.common.utils import PriorityQueue
+from rucio.core import heartbeat as heartbeat_core
+from rucio.core.monitor import MetricManager
+
+if TYPE_CHECKING:
+    from rucio.common.types import LoggerFunction
+
+T = TypeVar('T')
+METRICS = MetricManager(module=__name__)
+
+
+class HeartbeatHandler:
+    """
+    Simple contextmanager which sets a heartbeat and associated logger on entry and cleans up the heartbeat on exit.
+    """
+
+    def __init__(self, executable: str, renewal_interval: int):
+        """
+        :param executable: the executable name which will be set in heartbeats
+        :param renewal_interval: the interval at which the heartbeat will be renewed in the database.
+        Calls to live() in-between intervals will reuse the locally cached heartbeat.
+        """
+        self.executable = executable
+        self._hash_executable = None
+        self.renewal_interval = renewal_interval
+        self.older_than = renewal_interval * 10 if renewal_interval and renewal_interval > 0 else None  # 10 was chosen without any particular reason
+
+        self.hostname = socket.getfqdn()
+        self.pid = os.getpid()
+        self.hb_thread = threading.current_thread()
+
+        self.logger = logging.log
+        self.last_heart_beat = None
+        self.last_time = None
+        self.last_payload = None
+
+    def __enter__(self) -> 'HeartbeatHandler':
+        heartbeat_core.sanity_check(executable=self.executable, hostname=self.hostname)
+        self.live()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self.last_heart_beat:
+            heartbeat_core.die(self.executable, self.hostname, self.pid, self.hb_thread)
+            if self.logger:
+                self.logger(logging.INFO, 'Heartbeat cleaned up')
+
+    @property
+    def hash_executable(self) -> str:
+        if not self._hash_executable:
+            self._hash_executable = heartbeat_core.calc_hash(self.executable)
+        return self._hash_executable
+
+    @property
+    def short_executable(self) -> str:
+        return min(self.executable, self.hash_executable, key=len)
+
+    def live(
+            self,
+            force_renew: bool = False,
+            payload: Optional[str] = None
+    ) -> tuple[int, int, Callable]:
+        """
+        :return: a tuple: <the number of the current worker>, <total number of workers>, <decorated logger>
+        """
+        if force_renew \
+                or not self.last_time \
+                or not self.last_heart_beat \
+                or self.last_time < datetime.datetime.now() - datetime.timedelta(seconds=self.renewal_interval) \
+                or self.last_payload != payload:
+            if self.older_than:
+                self.last_heart_beat = heartbeat_core.live(self.executable, self.hostname, self.pid, self.hb_thread, payload=payload, older_than=self.older_than)
+            else:
+                self.last_heart_beat = heartbeat_core.live(self.executable, self.hostname, self.pid, self.hb_thread, payload=payload)
+
+            prefix = '[%i/%i]: ' % (self.last_heart_beat['assign_thread'], self.last_heart_beat['nr_threads'])
+            self.logger = formatted_logger(logging.log, prefix + '%s')
+
+            if not self.last_time:
+                self.logger(logging.DEBUG, 'First heartbeat set')
+            else:
+                self.logger(logging.DEBUG, 'Heartbeat renewed')
+            self.last_time = datetime.datetime.now()
+            self.last_payload = payload
+
+        return self.last_heart_beat['assign_thread'], self.last_heart_beat['nr_threads'], self.logger
+
+
+def _activity_looper(
+        once: bool,
+        sleep_time: int,
+        activities: Optional[Sequence[str]],
+        heartbeat_handler: HeartbeatHandler,
+) -> Generator[tuple[str, float], tuple[float, bool], None]:
+    """
+    Generator which loops (either once, or indefinitely) over all activities while ensuring that `sleep_time`
+    passes between handling twice the same activity.
+
+    Returns an activity and how much time the calling context must sleep before handling that activity
+    and expects to get in return the time when the activity started to be executed and whether next
+    execution must be immediate.
+    """
+
+    # For each activity, the priority queue will keep the next absolute time when that
+    # activity must be handled.
+    activity_next_exe_time = PriorityQueue()
+
+    # On startup, we schedule to immediately handle all activities.
+    now = time.time()
+    for activity in activities or [None]:
+        activity_next_exe_time[activity] = now
+
+    while activity_next_exe_time:
+        activity = activity_next_exe_time.top()
+        desired_exe_time = activity_next_exe_time[activity]
+
+        if once:
+            time_to_sleep = 0
+            activity_next_exe_time.pop()
+        else:
+            time_to_sleep = desired_exe_time - time.time()
+
+        logger = heartbeat_handler.logger
+        if time_to_sleep > 0:
+            if activity:
+                logger(logging.DEBUG, 'Switching to activity %s and sleeping %s seconds', activity, time_to_sleep)
+            else:
+                logger(logging.DEBUG, 'Sleeping %s seconds', time_to_sleep)
+        else:
+            if activity:
+                logger(logging.DEBUG, 'Switching to activity %s', activity)
+            else:
+                logger(logging.DEBUG, 'Starting next iteration')
+
+        # The calling context notifies us when the activity actually got handled. And if sleeping is desired.
+        actual_exe_time, must_sleep = yield activity, time_to_sleep
+
+        if not once:
+            if must_sleep:
+                time_diff = time.time() - actual_exe_time
+                time_to_sleep = max(1, sleep_time - time_diff)
+                activity_next_exe_time[activity] = time.time() + time_to_sleep
+            else:
+                activity_next_exe_time[activity] = time.time() + 1
+
+
+def db_workqueue(
+        once: bool,
+        graceful_stop: threading.Event,
+        executable: str,
+        partition_wait_time: int,
+        sleep_time: int,
+        activities: Optional[Sequence[str]] = None,
+) -> Callable[
+    [
+        Callable[
+            ...,
+            Union[bool, tuple[bool, T], None]
+        ]
+    ],
+    Callable[
+        [],
+        Iterator[Union[T, None]]
+    ]
+]:
+    """
+    Used to wrap a function for interacting with the database as a work queue: i.e. to select
+    a set of rows and perform some work on those rows while ensuring that two instances running in parallel don't
+    work on the same set of rows. The last condition is ensured by using heartbeats to keep track of currently
+    active workers.
+
+    :param once: Whether to stop after one iteration
+    :param graceful_stop: the threading.Event() object used for graceful stop of the daemon
+    :param executable: the name of the executable used for heartbeats
+    :param partition_wait_time: time to wait for database partition rebalancing before starting the actual daemon loop
+    :param sleep_time: time to sleep between the iterations of the daemon
+    :param activities: optional list of activities on which to work. The run_once_fnc will be called on activities one by one.
+    """
+
+    def _decorate(run_once_fnc: Callable[..., Optional[Union[bool, tuple[bool, T]]]]) -> Callable[[], Iterator[Optional[T]]]:
+
+        @functools.wraps(run_once_fnc)
+        def _generator() -> Iterator[T]:
+
+            with HeartbeatHandler(executable=executable, renewal_interval=sleep_time - 1) as heartbeat_handler:
+                logger = heartbeat_handler.logger
+                logger(logging.INFO, 'started')
+
+                if partition_wait_time:
+                    graceful_stop.wait(partition_wait_time)
+                    _, _, logger = heartbeat_handler.live(force_renew=True)
+
+                activity_loop = _activity_looper(once=once, sleep_time=sleep_time, activities=activities, heartbeat_handler=heartbeat_handler)
+                activity, time_to_sleep = next(activity_loop, (None, None))
+                while time_to_sleep is not None:
+                    if graceful_stop.is_set():
+                        break
+
+                    if time_to_sleep > 0:
+                        graceful_stop.wait(time_to_sleep)
+
+                    _, _, logger = heartbeat_handler.live()
+
+                    must_sleep = True
+                    start_time = time.time()
+                    try:
+                        result = run_once_fnc(heartbeat_handler=heartbeat_handler, activity=activity)
+
+                        # Handle return values already existing in the code
+                        # TODO: update all existing daemons to always explicitly return (must_sleep, ret_value)
+                        if result is None:
+                            must_sleep = True
+                            ret_value = None
+                        elif isinstance(result, bool):
+                            must_sleep = result
+                            ret_value = None
+                        else:
+                            must_sleep, ret_value = result
+
+                        if ret_value is not None:
+                            yield ret_value
+                    except Exception as e:
+                        METRICS.counter('exceptions.{exception}').labels(exception=e.__class__.__name__).inc()
+                        logger(logging.CRITICAL, "Exception", exc_info=True)
+                        if once:
+                            raise
+
+                    try:
+                        activity, time_to_sleep = activity_loop.send((start_time, must_sleep))
+                    except StopIteration:
+                        break
+
+                if not once:
+                    logger(logging.INFO, 'Graceful stop requested')
+
+        return _generator
+
+    return _decorate
+
+
+def run_daemon(
+        once: bool,
+        graceful_stop: threading.Event,
+        executable: str,
+        partition_wait_time: int,
+        sleep_time: int,
+        run_once_fnc: Callable[..., Optional[Union[bool, tuple[bool, Any]]]],
+        activities: Optional[list[str]] = None
+) -> None:
+    """
+    Run the daemon loop and call the function run_once_fnc at each iteration
+    """
+
+    daemon = db_workqueue(
+        once=once,
+        graceful_stop=graceful_stop,
+        executable=executable,
+        partition_wait_time=partition_wait_time,
+        sleep_time=sleep_time,
+        activities=activities,
+    )(run_once_fnc)
+
+    for _ in daemon():
+        pass
+
+
+class ProducerConsumerDaemon(Generic[T]):
+    """
+    Daemon which connects N producers with M consumers via a queue.
+    """
+
+    def __init__(
+            self,
+            producers: Sequence[Callable[[], Iterator[T]]],
+            consumers: Sequence[Callable[..., None]],
+            graceful_stop: threading.Event,
+            logger: "LoggerFunction" = logging.log
+    ):
+        self.producers = producers
+        self.consumers = consumers
+
+        self.queue = queue.Queue()
+        self.lock = threading.Lock()
+        self.graceful_stop = graceful_stop
+        self.active_producers = 0
+        self.producers_done_event = threading.Event()
+        self.logger = logger
+
+    def _produce(
+            self,
+            it: Callable[[], Iterator[T]],
+            wait_for_consumers: bool = False
+    ) -> None:
+        """
+        Iterate over the generator function and put the extracted elements into the queue.
+
+        Perform a graceful shutdown when graceful_stop is set.
+        """
+
+        i = it()
+        with self.lock:
+            self.active_producers += 1
+        try:
+            while not self.graceful_stop.is_set():
+                if self.queue.qsize() > len(self.consumers):
+                    self.graceful_stop.wait(1)
+                    continue
+
+                try:
+                    product = next(i)
+                    self.queue.put(product)
+                except StopIteration:
+                    break
+                except Exception as e:
+                    METRICS.counter('exceptions.{exception}').labels(exception=e.__class__.__name__).inc()
+                    self.logger(logging.CRITICAL, "Exception", exc_info=True)
+        finally:
+            with self.lock:
+                self.active_producers -= 1
+                if not self.active_producers > 0:
+                    self.producers_done_event.set()
+
+            if wait_for_consumers:
+                self.queue.join()
+
+    def _consume(
+            self,
+            fnc: Callable[[T], Any]
+    ) -> None:
+        """
+        Wait for elements to arrive via the queue and call the given function on each element.
+
+        If producers_done_event is set, handle all remaining elements from the queue and exit gracefully.
+        """
+        while not self.producers_done_event.is_set() or self.queue.unfinished_tasks:
+            try:
+                product = self.queue.get_nowait()
+            except queue.Empty:
+                self.producers_done_event.wait(1)
+                continue
+
+            try:
+                fnc(product)
+            except Exception as e:
+                METRICS.counter('exceptions.{exception}').labels(exception=e.__class__.__name__).inc()
+                self.logger(logging.CRITICAL, "Exception", exc_info=True)
+            finally:
+                self.queue.task_done()
+
+    def run(self) -> None:
+
+        producer_threads: list[threading.Thread] = []
+        for i, producer in enumerate(self.producers):
+            thread = threading.Thread(
+                target=self._produce,
+                name=f'producer-{i}-{producer.__name__}',
+                kwargs={
+                    'it': producer,
+                    'wait_for_consumers': True
+                }
+            )
+            thread.start()
+            producer_threads.append(thread)
+
+        consumer_threads: list[threading.Thread] = []
+        for i, consumer in enumerate(self.consumers):
+            thread = threading.Thread(
+                target=self._consume,
+                name=f'consumer-{i}-{consumer.__name__}',
+                kwargs={
+                    'fnc': consumer,
+                }
+            )
+            thread.start()
+            consumer_threads.append(thread)
+
+        logging.info('waiting for interrupts')
+
+        while producer_threads:
+            for thread in producer_threads:
+                thread.join(timeout=3.14)
+            producer_threads = [thread for thread in producer_threads if thread.is_alive()]
+
+        self.producers_done_event.set()
+
+        while consumer_threads:
+            for thread in consumer_threads:
+                thread.join(timeout=3.14)
+            consumer_threads = [thread for thread in consumer_threads if thread.is_alive()]

rucio/daemons/conveyor/__init__.py

@@ -0,0 +1,13 @@
+# 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.