rucio 37.0.0rc3__py3-none-any.whl → 37.1.0__py3-none-any.whl

rucio/common/utils.py

@@ -852,27 +852,24 @@ def get_bytes_value_from_string(input_string: str) -> Union[bool, int]:
     :param input_string: String containing a value and an unit
     :return: Integer value representing the value in bytes
     """
-    result = re.findall('^([0-9]+)([A-Za-z]+)$', input_string)
+    unit_multipliers = {
+        'b': 1,
+        'kb': 10**3,
+        'mb': 10**6,
+        'gb': 10**9,
+        'tb': 10**12,
+        'pb': 10**15,
+    }
+
+    result = re.findall(r'^([0-9]+)([A-Za-z]+)$', input_string)
     if result:
         value = int(result[0][0])
         unit = result[0][1].lower()
-        if unit == 'b':
-            value = value
-        elif unit == 'kb':
-            value = value * 1000
-        elif unit == 'mb':
-            value = value * 1000000
-        elif unit == 'gb':
-            value = value * 1000000000
-        elif unit == 'tb':
-            value = value * 1000000000000
-        elif unit == 'pb':
-            value = value * 1000000000000000
-        else:
+        multiplier = unit_multipliers.get(unit)
+        if multiplier is None:
             return False
-        return value
-    else:
-        return False
+        return value * multiplier
+    return False
 
 
 def parse_did_filter_from_string(input_string: str) -> tuple[dict[str, Any], str]:

rucio/core/account_limit.py

@@ -18,13 +18,16 @@ from sqlalchemy.exc import NoResultFound
 from sqlalchemy.sql import func, literal, select
 from sqlalchemy.sql.expression import and_, or_
 
-from rucio.core.account import get_all_rse_usages_per_account
+from rucio.common.exception import AccountNotFound
+from rucio.core.account import account_exists, get_all_rse_usages_per_account
 from rucio.core.rse import get_rse_name
 from rucio.core.rse_expression_parser import parse_expression
 from rucio.db.sqla import models
 from rucio.db.sqla.session import read_session, transactional_session
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable
+
     from sqlalchemy.orm import Session
 
     from rucio.common.types import InternalAccount, RSEAccountUsageDict, RSEGlobalAccountUsageDict, RSELocalAccountUsageDict, RSEResolvedGlobalAccountLimitDict
@@ -91,26 +94,39 @@ def get_rse_account_usage(rse_id: str, *, session: "Session") -> list["RSEAccoun
 
 
 @read_session
-def get_global_account_limits(account: Optional["InternalAccount"] = None, *, session: "Session") -> dict[str, "RSEResolvedGlobalAccountLimitDict"]:
+def get_global_account_limit(account: Optional["InternalAccount"] = None, rse_expression: Optional[str] = None, *,
+                             session: "Session") -> Optional[Union[int, float, dict[str, "RSEResolvedGlobalAccountLimitDict"]]]:
     """
-    Returns the global account limits for the account.
+    Returns the global account limit for the given account and RSE expression.
+    If no RSE expression is provided, returns all limits for the given account.
 
-    :param account:  Account to check the limit for.
-    :param session:  Database session in use.
-    :return:         Dict {'MOCK': {'resolved_rses': ['MOCK'], 'limit': 10, 'resolved_rse_ids': [123]}}.
+    Note: The account parameter must always be specified when fetching global account limit for an account.
+
+    :param account:         Account to check the limit for (required unless fetching all limits for all accounts is supported).
+    :param rse_expression:  Specific RSE expression to check the limit for (optional).
+    :param session:         Database session in use.
+    :return:                Limit in Bytes for a single RSE expression, or a dictionary of all limits {'MOCK': {'resolved_rses': ['MOCK'], 'limit': 10, 'resolved_rse_ids': [123]}}.
     """
     if account:
-        stmt = select(
-            models.AccountGlobalLimit
-        ).where(
-            models.AccountGlobalLimit.account == account
-        )
-        global_account_limits = session.execute(stmt).scalars().all()
-    else:
-        stmt = select(
-            models.AccountGlobalLimit
-        )
-        global_account_limits = session.execute(stmt).scalars().all()
+        if not account_exists(account, session=session):
+            raise AccountNotFound(f"Account {account} does not exist")
+        if rse_expression:
+            # Fetch limit for a single RSE expression
+            try:
+                stmt = select(models.AccountGlobalLimit).where(
+                    and_(models.AccountGlobalLimit.account == account,
+                         models.AccountGlobalLimit.rse_expression == rse_expression)
+                )
+                global_account_limit = session.execute(stmt).scalar_one()
+                return float("inf") if global_account_limit.bytes == -1 else global_account_limit.bytes
+            except NoResultFound:
+                return None
+
+    # Fetch all global limits for the account (or all accounts if no account specified)
+    stmt = select(models.AccountGlobalLimit)
+    if account:
+        stmt = stmt.where(models.AccountGlobalLimit.account == account)
+    global_account_limits = session.execute(stmt).scalars().all()
 
     resolved_global_account_limits = {}
     for limit in global_account_limits:
@@ -118,9 +134,7 @@ def get_global_account_limits(account: Optional["InternalAccount"] = None, *, se
             resolved_rses = parse_expression(limit['rse_expression'], filter_={'vo': account.vo}, session=session)
         else:
             resolved_rses = parse_expression(limit['rse_expression'], session=session)
-        limit_in_bytes = limit['bytes']
-        if limit_in_bytes == -1:
-            limit_in_bytes = float('inf')
+        limit_in_bytes = float('inf') if limit['bytes'] == -1 else limit['bytes']
         resolved_global_account_limits[limit['rse_expression']] = {
             'resolved_rses': [resolved_rse['rse'] for resolved_rse in resolved_rses],
             'resolved_rse_ids': [resolved_rse['id'] for resolved_rse in resolved_rses],
@@ -130,69 +144,32 @@ def get_global_account_limits(account: Optional["InternalAccount"] = None, *, se
 
 
 @read_session
-def get_global_account_limit(account: "InternalAccount", rse_expression: str, *, session: "Session") -> Union[int, float, None]:
+def get_local_account_limit(account: "InternalAccount", rse_ids: Optional[Union[str, "Iterable[str]"]] = None, *, session: "Session") -> Optional[Union[int, float, dict[str, int]]]:
     """
-    Returns the global account limit for the account on the rse expression.
-
-    :param account:         Account to check the limit for.
-    :param rse_expression:  RSE expression to check the limit for.
-    :param session:         Database session in use.
-    :return:                Limit in Bytes.
-    """
-    try:
-        stmt = select(
-            models.AccountGlobalLimit
-        ).where(
-            and_(models.AccountGlobalLimit.account == account,
-                 models.AccountGlobalLimit.rse_expression == rse_expression)
-        )
-        global_account_limit = session.execute(stmt).scalar_one()
-        if global_account_limit.bytes == -1:
-            return float("inf")
-        else:
-            return global_account_limit.bytes
-    except NoResultFound:
-        return None
-
-
-@read_session
-def get_local_account_limit(account: "InternalAccount", rse_id: str, *, session: "Session") -> Union[int, float, None]:
-    """
-    Returns the account limit for the account on the rse.
-
-    :param account:  Account to check the limit for.
-    :param rse_id:   RSE id to check the limit for.
-    :param session:  Database session in use.
-    :return:         Limit in Bytes.
-    """
-    try:
-        stmt = select(
-            models.AccountLimit
-        ).where(
-            and_(models.AccountLimit.account == account,
-                 models.AccountLimit.rse_id == rse_id)
-        )
-        account_limit = session.execute(stmt).scalar_one()
-        if account_limit.bytes == -1:
-            return float("inf")
-        else:
-            return account_limit.bytes
-    except NoResultFound:
-        return None
-
-
-@read_session
-def get_local_account_limits(account: "InternalAccount", rse_ids: Optional[list[str]] = None, *, session: "Session") -> dict[str, int]:
-    """
-    Returns the account limits for the account on the list of rses.
+    Returns the local account limit for a given RSE or list of RSEs.
 
     :param account:  Account to check the limit for.
-    :param rse_ids:  List of RSE ids to check the limit for.
+    :param rse_ids:  Single RSE id or an iterable of RSE ids to check the limit for.
     :param session:  Database session in use.
-    :return:         Dictionary {'rse_id': bytes, ...}.
+    :return:         Limit in Bytes (int/float) for a single RSE or
+                     Dictionary {'rse_id': bytes, ...} for multiple RSEs.
     """
+    if not account_exists(account, session=session):
+        raise AccountNotFound(f"Account {account} does not exist")
+    if isinstance(rse_ids, str):  # Single RSE case
+        try:
+            stmt = select(models.AccountLimit).where(
+                and_(models.AccountLimit.account == account, models.AccountLimit.rse_id == rse_ids)
+            )
+            account_limit = session.execute(stmt).scalar_one()
+            return float("inf") if account_limit.bytes == -1 else account_limit.bytes
+        except NoResultFound:
+            return None
 
+    # Multiple RSE case or no RSE specified
     account_limits = {}
+
+    # If rse_ids is a list of RSEs
     if rse_ids:
         rse_id_clauses = []
         for rse_id in rse_ids:
@@ -339,14 +316,14 @@ def get_local_account_usage(account: "InternalAccount", rse_id: Optional[str] =
     )
     if not rse_id:
         # All RSESs
-        limits = get_local_account_limits(account=account, session=session)
+        limits = get_local_account_limit(account=account, session=session)
         counters = {c.rse_id: c for c in session.execute(stmt).scalars().all()}
     else:
         # One RSE
         stmt.where(
             models.AccountUsage.rse_id == rse_id
         )
-        limits = get_local_account_limits(account=account, rse_ids=[rse_id], session=session)
+        limits = get_local_account_limit(account=account, rse_ids=[rse_id], session=session)
         counters = {c.rse_id: c for c in session.execute(stmt).scalars().all()}
     result_list = []
 
@@ -385,7 +362,7 @@ def get_global_account_usage(account: "InternalAccount", rse_expression: Optiona
     result_list = []
     if not rse_expression:
         # All RSE Expressions
-        limits = get_global_account_limits(account=account, session=session)
+        limits = get_global_account_limit(account=account, session=session)
         all_rse_usages = {usage['rse_id']: (usage['bytes'], usage['files']) for usage in get_all_rse_usages_per_account(account=account, session=session)}
         for rse_expression, limit in limits.items():
             usage = 0

rucio/core/did_meta_plugins/filter_engine.py

@@ -102,9 +102,7 @@ class FilterEngine:
             filters, _ = parse_did_filter_from_string_fe(filters, omit_name=True)
         elif isinstance(filters, dict):
             filters = [filters]
-        elif isinstance(filters, list):
-            filters = filters
-        else:
+        elif not isinstance(filters, list):
             raise exception.DIDFilterSyntaxError("Input filters are of an unrecognised type.")
 
         filters = self._make_input_backwards_compatible(filters=filters)

rucio/core/rse_selector.py

@@ -17,7 +17,7 @@ from typing import TYPE_CHECKING, Optional
 
 from rucio.common.exception import InsufficientAccountLimit, InsufficientTargetRSEs, InvalidRuleWeight, RSEOverQuota
 from rucio.core.account import get_all_rse_usages_per_account, get_usage, has_account_attribute
-from rucio.core.account_limit import get_global_account_limits, get_local_account_limit
+from rucio.core.account_limit import get_global_account_limit, get_local_account_limit
 from rucio.core.rse import get_rse_limits, has_rse_attribute, list_rse_attributes
 from rucio.core.rse_counter import get_counter as get_rse_counter
 from rucio.core.rse_expression_parser import parse_expression
@@ -86,7 +86,7 @@ class RSESelector:
                 rse['space_left'] = float('inf')
                 rses_with_enough_quota.append(rse)
         else:
-            global_quota_limit = get_global_account_limits(account=account, session=session)
+            global_quota_limit = get_global_account_limit(account=account, session=session)
             all_rse_usages = {usage['rse_id']: usage['bytes'] for usage in get_all_rse_usages_per_account(account=account, session=session)}
             for rse in self.rses:
                 if rse['mock_rse']:
@@ -96,7 +96,7 @@ class RSESelector:
                 else:
                     # check local quota
                     local_quota_left = None
-                    quota_limit = get_local_account_limit(account=account, rse_id=rse['rse_id'], session=session)
+                    quota_limit = get_local_account_limit(account=account, rse_ids=rse['rse_id'], session=session)
                     if quota_limit is None:
                         local_quota_left = 0
                     else:

rucio/core/rule_grouping.py

@@ -1617,4 +1617,3 @@ def apply_rule(did, rule, rses, source_rses, rseselector, *, session: "Session",
         account_counter.increase(rse_id=rse_id, account=rule.account, files=account_counters_files[rse_id], bytes_=account_counters_bytes[rse_id], session=session)
     session.flush()
 
-    return

rucio/daemons/cache/consumer.py

@@ -21,12 +21,13 @@ import logging
 import threading
 import time
 from traceback import format_exc
-from typing import TYPE_CHECKING
+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 ListenerBase, StompConnectionManager
+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
@@ -35,6 +36,7 @@ from rucio.core.volatile_replica import add_volatile_replicas, delete_volatile_r
 if TYPE_CHECKING:
     from types import FrameType
 
+    from stomp import Connection
     from stomp.utils import Frame
 
 logging.getLogger("stomp").setLevel(logging.CRITICAL)
@@ -44,11 +46,35 @@ GRACEFUL_STOP = threading.Event()
 DAEMON_NAME = 'cache-consumer'
 
 
-class AMQConsumer(ListenerBase):
+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:
         """
@@ -56,7 +82,7 @@ class AMQConsumer(ListenerBase):
         """
         try:
             msg = json.loads(frame.body)  # type: ignore
-            self._logger(logging.DEBUG, 'Message received: %s', msg)
+            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'])
@@ -67,42 +93,81 @@ class AMQConsumer(ListenerBase):
 
                 rse_vo_str = msg['rse']
                 if 'vo' in msg and msg['vo'] != 'def':
-                    rse_vo_str = f"{rse_vo_str} on {msg['vo']}"
+                    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']))
+                    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']))
+                    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())
+                self.__logger(logging.DEBUG, 'Check failed: %s %s '
+                              % (isinstance(msg, dict), 'operation' in msg.keys()))
         except:
-            self._logger(logging.ERROR, str(format_exc()))
+            self.__logger(logging.ERROR, str(format_exc()))
 
 
-def consumer(id_: int, num_thread: int = 1, logger: LoggerFunction = logging.log) -> None:
+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')
 
-    conn_mgr = StompConnectionManager(config_section='messaging-cache', logger=logger)
+    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')
 
-    conn_mgr.set_listener_factory('rucio-cache-consumer', AMQConsumer, heartbeats=conn_mgr.config.heartbeats)
-
     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_mgr.subscribe(id_='rucio-cache-messaging', ack='auto')
+                conn.subscribe(destination=destination, ack='auto', id=subscription_id)
         time.sleep(1)
 
     logger(logging.INFO, 'graceful stop requested')
-    conn_mgr.disconnect()
+    stomp_conn_mngr.disconnect()
     logger(logging.INFO, 'graceful stop done')
 
 
-def stop(signum: "int | None" = None, frame: "FrameType | None" = None) -> None:
+def stop(signum: Optional[int] = None, frame: Optional["FrameType"] = None) -> None:
     """
     Graceful exit.
     """
@@ -115,19 +180,18 @@ def run(num_thread: int = 1) -> None:
     Starts up the rucio cache consumer thread
     """
     setup_logging(process_name=DAEMON_NAME)
-    logger = formatted_logger(logging.log, DAEMON_NAME + ' %s')
 
     if rucio.db.sqla.util.is_old_db():
         raise exception.DatabaseException('Database was not updated, daemon won\'t start')
 
-    logger(logging.INFO, 'starting consumer thread')
-    threads = []
-    for i in range(num_thread):
-        con_thread = threading.Thread(target=consumer, kwargs={'id_': i, 'num_thread': num_thread, 'logger': logger})
-        con_thread.start()
-        threads.append(con_thread)
+    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]
 
-    logger(logging.INFO, 'waiting for interrupts')
+    logging.info('waiting for interrupts')
 
-    while [thread.join(timeout=3.14) for thread in threads if thread.is_alive()]:
-        pass
+    # Interruptible joins require a timeout.
+    while threads[0].is_alive():
+        [t.join(timeout=3.14) for t in threads]