rucio 37.0.0rc1__py3-none-any.whl

rucio/daemons/bb8/common.py

@@ -0,0 +1,759 @@
+# 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 logging
+from datetime import date, datetime, timedelta
+from string import Template
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+
+from requests import get
+from sqlalchemy import BigInteger, and_, case, cast, false, func, or_, select
+from sqlalchemy.orm import Session, aliased
+
+from rucio.common.config import config_get, config_get_bool, config_get_int
+from rucio.common.exception import (
+    DuplicateRule,
+    InsufficientAccountLimit,
+    InsufficientTargetRSEs,
+    RuleNotFound,
+)
+from rucio.common.types import InternalAccount, InternalScope
+from rucio.core.lock import get_dataset_locks
+from rucio.core.rse import get_rse_name, get_rse_vo, list_rse_attributes
+from rucio.core.rse_expression_parser import parse_expression
+from rucio.core.rse_selector import RSESelector
+from rucio.core.rule import add_rule, get_rule, update_rule
+from rucio.db.sqla import models
+from rucio.db.sqla.constants import DIDType, LockState, RuleGrouping, RuleState
+from rucio.db.sqla.session import read_session, transactional_session
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+    from sqlalchemy.engine import Row
+
+    from rucio.common.types import LoggerFunction
+
+
+@transactional_session
+def rebalance_rule(
+    parent_rule: "Mapping[str, Any]",
+    activity: str,
+    rse_expression: str,
+    priority: int,
+    source_replica_expression: str = "*\\bb8-enabled=false",
+    comment: Optional[str] = None,
+    *,
+    session: Session,
+) -> Optional[list[str]]:
+    """
+    Rebalance a replication rule to a new RSE
+    :param parent_rule:                Replication rule to be rebalanced.
+    :param activity:                   Activity to be used for the rebalancing.
+    :param rse_expression:             RSE expression of the new rule.
+    :param priority:                   Priority of the newly created rule.
+    :param source_replica_expression:  Source replica expression of the new rule.
+    :param comment:                    Comment to set on the new rules.
+    :returns:                          The new child rule id.
+    """
+
+    if parent_rule["expires_at"] is None:
+        lifetime = None
+    else:
+        lifetime = (parent_rule["expires_at"] - datetime.utcnow()).days * 24 * 3600 + (
+            parent_rule["expires_at"] - datetime.utcnow()
+        ).seconds
+
+    if parent_rule["grouping"] == RuleGrouping.ALL:
+        grouping = "ALL"
+    elif parent_rule["grouping"] == RuleGrouping.NONE:
+        grouping = "NONE"
+    else:
+        grouping = "DATASET"
+
+    # check if concurrent replica at target rse does not exist
+    concurrent_replica = False
+    for lock in get_dataset_locks(parent_rule["scope"], parent_rule["name"]):
+        lock_rse_expr = lock["rse"]
+        if lock_rse_expr == rse_expression:
+            concurrent_replica = True
+
+    if concurrent_replica:
+        return None
+
+    child_rule = add_rule(
+        dids=[{"scope": parent_rule["scope"], "name": parent_rule["name"]}],
+        account=parent_rule["account"],
+        copies=parent_rule["copies"],
+        rse_expression=rse_expression,
+        grouping=grouping,
+        weight=parent_rule["weight"],
+        lifetime=lifetime,
+        locked=parent_rule["locked"],
+        subscription_id=parent_rule["subscription_id"],
+        source_replica_expression=source_replica_expression,
+        activity=activity,
+        notify=parent_rule["notification"],
+        purge_replicas=parent_rule["purge_replicas"],
+        ignore_availability=False,
+        comment=parent_rule["comments"] if not comment else comment,
+        ask_approval=False,
+        asynchronous=False,
+        ignore_account_limit=True,
+        priority=priority,
+        session=session,
+    )[0]
+
+    update_rule(
+        rule_id=parent_rule["id"],
+        options={"child_rule_id": child_rule, "lifetime": 0},
+        session=session,
+    )
+    return child_rule
+
+
+def __dump_url(
+        rse_id: str,
+        logger: "LoggerFunction" = logging.log
+) -> Union[list[str], Literal[False]]:
+    """
+    getting potential urls of the dump over last week
+    :param rse_id:                     RSE where the dump is released.
+    :param logger:                     Logger.
+    """
+
+    rse = get_rse_name(rse_id=rse_id)
+    vo = get_rse_vo(rse_id=rse_id)
+
+    # get the date of the most recent dump
+    today = date.today()
+    dump_dates = []
+    dump_production_day = config_get(
+        "bb8", "dump_production_day", raise_exception=False, default=None
+    )
+    if dump_production_day is None:
+        for idx in range(0, 7):
+            dump_date = today - timedelta(idx)
+            dump_dates.append(dump_date.strftime("%d-%m-%Y"))
+    else:
+        weekdays = {
+            "Sunday": 6,
+            "Monday": 0,
+            "Tuesday": 1,
+            "Wednesday": 2,
+            "Thursday": 3,
+            "Friday": 4,
+            "Saturday": 5,
+        }
+        if dump_production_day not in weekdays:
+            logger(
+                logging.WARNING,
+                "ERROR: please set the day of a dump creation in bb8 config correctly, e.g. Monday",
+            )
+            return False
+        today_idx = (today.weekday() - weekdays[dump_production_day]) % 7
+        dump_date = today - timedelta(today_idx)
+        dump_dates = [dump_date.strftime("%d-%m-%Y")]
+
+    # getting structure (template) of url location of a dump
+    url_template_str = config_get(
+        "bb8",
+        "dump_url_template",
+        raise_exception=False,
+    )
+    url_template = Template(url_template_str)
+
+    # populating url template
+    urls = []
+    for d in dump_dates:
+        url = url_template.substitute({"date": d, "rse": rse, "vo": vo})
+        urls.append(url)
+    return urls
+
+
+def _list_rebalance_rule_candidates_dump(
+        rse_id: str,
+        mode: Optional[str] = None,
+        logger: "LoggerFunction" = logging.log
+) -> list[tuple]:
+    """
+    Download dump to temporary directory
+    :param rse_id:                     RSE of the source.
+    :param mode:                       Rebalancing mode.
+    :param logger:                     Logger.
+    """
+
+    # fetching the dump
+    candidates = []
+    rules = {}
+    rse_dump_urls = __dump_url(rse_id=rse_id)
+    resp = None
+    if not rse_dump_urls:
+        logger(logging.DEBUG, "URL of the dump was not built from template.")
+        return candidates
+    rse_dump_urls.reverse()
+    success = False
+    while not success and len(rse_dump_urls):
+        url = rse_dump_urls.pop()
+        resp = get(url, stream=True)
+        if resp:
+            success = True
+    if not resp or resp is None:
+        logger(logging.WARNING, "RSE dump not available")
+        return candidates
+
+    # looping over the dump and selecting the rules
+    for line in resp.iter_lines():
+        if line:
+            _, _, rule_id, rse_expression, account, file_size, state = line.split("\t")
+            if rule_id not in rules:
+                rule_info = {}
+                try:
+                    rule_info = get_rule(rule_id=rule_id)
+                except Exception as err:
+                    rules[rule_id] = {"state": "DELETED"}
+                    logger(logging.ERROR, str(err))
+                    continue
+                if rule_info["child_rule_id"]:
+                    rules[rule_id] = {"state": "DELETED"}
+                    continue
+                rules[rule_id] = {
+                    "scope": rule_info["scope"],
+                    "name": rule_info["name"],
+                    "rse_expression": rse_expression,
+                    "subscription_id": rule_info["subscription_id"],
+                    "length": 1,
+                    "state": "ACTIVE",
+                    "bytes": int(file_size),
+                }
+            elif rules[rule_id]["state"] == "ACTIVE":
+                rules[rule_id]["length"] += 1
+                rules[rule_id]["bytes"] += int(file_size)
+
+    # looping over agragated rules collected from dump
+    for r_id in rules:
+        if mode == "decommission":  # other modes can be added later
+            if rules[r_id]["state"] == "DELETED":
+                continue
+            if int(rules[r_id]["length"]) == 0:
+                continue
+            candidates.append(
+                (
+                    rules[r_id]["scope"],
+                    rules[r_id]["name"],
+                    r_id,
+                    rules[r_id]["rse_expression"],
+                    rules[r_id]["subscription_id"],
+                    rules[r_id]["bytes"],
+                    rules[r_id]["length"],
+                    int(rules[r_id]["bytes"] / rules[r_id]["length"]),
+                )
+            )
+    return candidates
+
+
+@transactional_session
+def list_rebalance_rule_candidates(
+    rse_id: str,
+    mode: Optional[str] = None,
+    *,
+    session: Optional[Session] = None
+) -> Union[list[tuple], list["Row[tuple]"]]:
+    """
+    List the rebalance rule candidates based on the agreed on specification
+    :param rse_id:       RSE of the source.
+    :param mode:         Rebalancing mode.
+    :param session:      DB Session.
+    """
+
+    vo = get_rse_vo(rse_id=rse_id)
+
+    # dumps can be applied only for decommission since the dumps doesn't contain info from dids
+    if mode == "decommission":
+        return _list_rebalance_rule_candidates_dump(rse_id, mode)
+
+    # If no decommissioning use SQLAlchemy
+
+    # Rules constraints. By default only moves rules in state OK that have no children and have only one copy
+    # Additional constraints can be imposed by setting specific configuration
+    rule_clause = [
+        models.ReplicationRule.state == RuleState.OK,
+        models.ReplicationRule.child_rule_id.is_(None),
+        models.ReplicationRule.copies == 1,
+    ]
+
+    # Only move rules w/o expiration date, or rules with expiration_date > >min_expires_date_in_days> days
+    expiration_clause = models.ReplicationRule.expires_at.is_(None)
+    min_expires_date_in_days = config_get_int(
+        section="bb8",
+        option="min_expires_date_in_days",
+        raise_exception=False,
+        default=-1,
+        expiration_time=3600,
+    )
+    if min_expires_date_in_days > 0:
+        min_expires_date_in_days = datetime.utcnow() + timedelta(
+            days=min_expires_date_in_days
+        )
+        expiration_clause = or_(models.ReplicationRule.expires_at > min_expires_date_in_days,
+                                models.ReplicationRule.expires_at.is_(None))
+    rule_clause.append(expiration_clause)
+
+    # Only move rules which were created more than <min_created_days> days ago
+    min_created_days = config_get_int(
+        section="bb8",
+        option="min_created_days",
+        raise_exception=False,
+        default=-1,
+        expiration_time=3600,
+    )
+    if min_created_days > 0:
+        min_created_days = datetime.utcnow() - timedelta(days=min_created_days)
+        rule_clause.append(models.ReplicationRule.created_at < min_created_days)
+
+    # Only move rules which are owned by <allowed_accounts> (coma separated accounts, e.g. panda,root,ddmadmin,jdoe)
+    allowed_accounts = config_get(
+        section="bb8",
+        option="allowed_accounts",
+        raise_exception=False,
+        default=None,
+        expiration_time=3600,
+    )
+    if allowed_accounts:
+        allowed_accounts = [
+            InternalAccount(acc.strip(" "), vo=vo)
+            for acc in allowed_accounts.split(",")
+        ]
+        rule_clause.append(models.ReplicationRule.account.in_(allowed_accounts))
+
+    # Only move rules which with scope <allowed_scopes> (coma separated scopes, e.g. mc16_13TeV,data18_13TeV)
+    allowed_scopes = config_get(
+        section="bb8",
+        option="allowed_scopes",
+        raise_exception=False,
+        default=None,
+        expiration_time=3600,
+    )
+    if allowed_scopes:
+        allowed_scopes = [
+            InternalScope(scope.strip(" "), vo=vo)
+            for scope in allowed_scopes.split(",")
+        ]
+        rule_clause.append(models.ReplicationRule.scope.in_(allowed_scopes))
+
+    # Only move rules that have a certain grouping <allowed_grouping> (accepted values : all, dataset, none)
+    rule_grouping_mapping = {
+        "all": RuleGrouping.ALL,
+        "dataset": RuleGrouping.DATASET,
+        "none": RuleGrouping.NONE,
+    }
+    allowed_grouping = config_get(
+        section="bb8",
+        option="allowed_grouping",
+        raise_exception=False,
+        default=None,
+        expiration_time=3600,
+    )
+    if allowed_grouping:
+        rule_clause.append(
+            models.ReplicationRule.grouping
+            == rule_grouping_mapping.get(allowed_grouping)
+        )
+
+    # DIDs constraints. By default only moves rules of DID where we can compute the size
+    # Additional constraints can be imposed by setting specific configuration
+    did_clause = [models.DataIdentifier.bytes.isnot(None)]
+
+    type_to_did_type_mapping = {
+        "all": [DIDType.CONTAINER, DIDType.DATASET, DIDType.FILE],
+        "collection": [DIDType.CONTAINER, DIDType.DATASET],
+        "container": [DIDType.CONTAINER],
+        "dataset": [DIDType.DATASET],
+        "file": [DIDType.FILE],
+    }
+
+    # Only allows to migrate rules of a certain did_type <allowed_did_type> (accepted values : all, collection, container, dataset, file)
+    allowed_did_type = config_get(
+        section="bb8",
+        option="allowed_did_type",
+        raise_exception=False,
+        default=None,
+        expiration_time=3600,
+    )
+    if allowed_did_type:
+        allowed_did_type = [
+            models.DataIdentifier.did_type == did_type
+            for did_type in type_to_did_type_mapping.get(allowed_did_type)
+        ]
+        did_clause.append(or_(*allowed_did_type))
+
+    # Only allows to migrate rules of closed DID is <only_move_closed_did> is set
+    only_move_closed_did = config_get_bool(
+        section="bb8",
+        option="only_move_closed_did",
+        raise_exception=False,
+        default=None,
+        expiration_time=3600,
+    )
+    if only_move_closed_did:
+        did_clause.append(models.DataIdentifier.is_open == false())
+
+    # Now build the query
+    external_dsl = aliased(models.DatasetLock)
+    count_locks = select(
+        func.count()
+    ).select_from(
+        models.DatasetLock
+    ).where(
+        and_(external_dsl.scope == models.DatasetLock.scope,
+             external_dsl.name == models.DatasetLock.name,
+             external_dsl.rse_id == models.DatasetLock.rse_id)
+    ).scalar_subquery()
+
+    stmt = select(
+        models.DatasetLock.scope,
+        models.DatasetLock.name,
+        models.ReplicationRule.id,
+        models.ReplicationRule.rse_expression,
+        models.ReplicationRule.subscription_id,
+        models.DataIdentifier.bytes,
+        models.DataIdentifier.length,
+        case(
+            (
+                or_(models.DatasetLock.length < 1,
+                    models.DatasetLock.length.is_(None)),
+                0
+            ),
+            else_=cast(
+                models.DatasetLock.bytes / models.DatasetLock.length, BigInteger
+            )
+        )
+    ).join(
+        models.ReplicationRule,
+        models.ReplicationRule.id == models.DatasetLock.rule_id
+    ).join(
+        models.DataIdentifier,
+        and_(models.DatasetLock.scope == models.DataIdentifier.scope,
+             models.DatasetLock.name == models.DataIdentifier.name)
+    ).where(
+        and_(models.DatasetLock.rse_id == rse_id,
+             *rule_clause,
+             *did_clause,
+             case(
+                 (
+                     or_(models.DatasetLock.length < 1,
+                         models.DatasetLock.length.is_(None)),
+                     0
+                 ),
+                 else_=cast(
+                     models.DatasetLock.bytes / models.DatasetLock.length, BigInteger
+                 )
+             ) > 1000000000,
+             count_locks == 1)
+    ).order_by(
+        case(
+            (
+                or_(models.DatasetLock.length < 1,
+                    models.DatasetLock.length.is_(None)),
+                0
+            ),
+            else_=cast(
+                models.DatasetLock.bytes / models.DatasetLock.length, BigInteger
+            )
+        ),
+        models.DatasetLock.accessed_at
+    )
+    return list(session.execute(stmt).all())  # type: ignore (session could be None)
+
+
+@read_session
+def select_target_rse(
+    parent_rule: "Mapping[str, Any]",
+    current_rse_id: str,
+    rse_expression: str,
+    subscription_id: str,
+    rse_attributes: "Mapping[str, Any]",
+    other_rses: Optional["Sequence[str]"] = None,
+    exclude_expression: Optional[str] = None,
+    force_expression: Optional[str] = None,
+    *,
+    session: Optional[Session] = None,
+) -> str:
+    """
+    Select a new target RSE for a rebalanced rule.
+    :param parent_rule           rule that is rebalanced.
+    :param current_rse_id:       RSE of the source.
+    :param rse_expression:       RSE Expression of the source rule.
+    :param subscription_id:      Subscription ID of the source rule.
+    :param rse_attributes:       The attributes of the source rse.
+    :param other_rses:           Other RSEs with existing dataset replicas.
+    :param exclude_expression:   Exclude this rse_expression from being target_rses.
+    :param force_expression:     Force a specific rse_expression as target.
+    :param session:              The DB Session.
+    :returns:                    New RSE expression.
+    """
+
+    other_rses = other_rses or []
+    current_rse = get_rse_name(rse_id=current_rse_id)
+    current_rse_expr = current_rse
+    # if parent rule has a vo, enforce it
+    vo = parent_rule["scope"].vo
+    if exclude_expression:
+        target_rse = "((%s)|(%s))" % (exclude_expression, current_rse_expr)
+    else:
+        target_rse = current_rse_expr
+    list_target_rses = [
+        rse["rse"]
+        for rse in parse_expression(
+            expression=target_rse, filter_={"vo": vo}, session=session
+        )
+    ]
+    list_target_rses.sort()
+
+    rses = parse_expression(
+        expression=rse_expression, filter_={"vo": vo}, session=session
+    )
+
+    # TODO: Implement subscription rebalancing
+
+    if force_expression is not None:
+        if parent_rule["grouping"] != RuleGrouping.NONE:
+            rses = parse_expression(
+                expression="(%s)\\%s" % (force_expression, target_rse),
+                filter_={"vo": vo, "availability_write": True},
+                session=session,
+            )
+        else:
+            # in order to avoid replication of the part of distributed dataset not present at rebalanced rse -> rses in force_expression
+            # this will be extended with development of delayed rule
+            rses = parse_expression(
+                expression="((%s)|(%s))\\%s"
+                % (force_expression, rse_expression, target_rse),
+                filter_={"vo": vo, "availability_write": True},
+                session=session,
+            )
+    else:
+        # force_expression is not set, RSEs will be selected as rse_expression\rse
+        list_rses = [rse["rse"] for rse in rses]
+        list_rses.sort()
+        if list_rses == list_target_rses:
+            raise InsufficientTargetRSEs(
+                "Not enough RSEs to rebalance rule %s" % parent_rule["id"]
+            )
+        else:
+            rses = parse_expression(
+                expression="(%s)\\%s" % (rse_expression, target_rse),
+                filter_={"vo": vo, "availability_write": True},
+                session=session,
+            )
+    rseselector = RSESelector(
+        account=InternalAccount("root", vo=vo),
+        rses=rses,
+        weight="freespace",
+        copies=1,
+        ignore_account_limit=True,
+        session=session,
+    )
+    return get_rse_name(
+        [
+            rse_id
+            for rse_id, _, _ in rseselector.select_rse(
+                size=0, preferred_rse_ids=[], blocklist=other_rses
+            )
+        ][0],
+        session=session,
+    )
+
+
+@transactional_session
+def rebalance_rse(
+    rse_id: str,
+    max_bytes: float = 1e9,
+    max_files: Optional[int] = None,
+    dry_run: bool = False,
+    exclude_expression: Optional[str] = None,
+    comment: Optional[str] = None,
+    force_expression: Optional[str] = None,
+    mode: Optional[str] = None,
+    priority: int = 3,
+    source_replica_expression: str = "*\\bb8-enabled=false",
+    *,
+    session: Optional[Session] = None,
+    logger: "LoggerFunction" = logging.log,
+) -> list[tuple]:
+    """
+    Rebalance data from an RSE
+    :param rse_id:                     RSE to rebalance data from.
+    :param max_bytes:                  Maximum amount of bytes to rebalance.
+    :param max_files:                  Maximum amount of files to rebalance.
+    :param dry_run:                    Only run in dry-run mode.
+    :param exclude_expression:         Exclude this rse_expression from being target_rses.
+    :param comment:                    Comment to set on the new rules.
+    :param force_expression:           Force a specific rse_expression as target.
+    :param mode:                       BB8 mode to execute (None=normal, 'decomission'=Decomission mode)
+    :param priority:                   Priority of the new created rules.
+    :param source_replica_expression:  Source replica expression of the new created rules.
+    :param session:                    The database session.
+    :param logger:                     Logger.
+    :returns:                          List of rebalanced datasets.
+    """
+    rebalanced_bytes = 0
+    rebalanced_files = 0
+    rebalanced_datasets = []
+
+    rse_attributes = list_rse_attributes(rse_id=rse_id, session=session)
+    src_rse = get_rse_name(rse_id=rse_id)
+
+    logger(logging.INFO, "***************************")
+    logger(logging.INFO, "BB8 - Execution Summary")
+    logger(logging.INFO, "Mode:    %s" % ("STANDARD" if mode is None else mode.upper()))
+    logger(logging.INFO, "Dry Run: %s" % (dry_run))
+    logger(logging.INFO, "***************************")
+
+    for (
+        scope,
+        name,
+        rule_id,
+        rse_expression,
+        subscription_id,
+        bytes_,
+        length,
+        fsize,
+    ) in list_rebalance_rule_candidates(rse_id=rse_id, mode=mode):
+        if force_expression is not None and subscription_id is not None:
+            continue
+
+        if rebalanced_bytes + bytes_ > max_bytes:
+            continue
+        if max_files:
+            if rebalanced_files + length > max_files:
+                continue
+
+        try:
+            rule = get_rule(rule_id=rule_id)
+            other_rses = [
+                r["rse_id"] for r in get_dataset_locks(scope, name, session=session)
+            ]
+            # Select the target RSE for this rule
+            try:
+                target_rse_exp = select_target_rse(
+                    parent_rule=rule,
+                    current_rse_id=rse_id,
+                    rse_expression=rse_expression,
+                    subscription_id=subscription_id,
+                    rse_attributes=rse_attributes,
+                    other_rses=other_rses,
+                    exclude_expression=exclude_expression,
+                    force_expression=force_expression,
+                    session=session,
+                )
+                # Rebalance this rule
+                if not dry_run:
+                    child_rule_id = rebalance_rule(
+                        parent_rule=rule,
+                        activity="Data Rebalancing",
+                        rse_expression=target_rse_exp,
+                        priority=priority,
+                        source_replica_expression=source_replica_expression,
+                        comment=comment,
+                    )
+                else:
+                    child_rule_id = ""
+            except (
+                InsufficientTargetRSEs,
+                DuplicateRule,
+                RuleNotFound,
+                InsufficientAccountLimit,
+            ) as err:
+                logger(logging.ERROR, str(err))
+                continue
+            if child_rule_id is None:
+                logger(
+                    logging.WARNING,
+                    "A rule for %s:%s already exists on %s. It cannot be rebalanced",
+                    scope,
+                    name,
+                    target_rse_exp,
+                )
+                continue
+            logger(
+                logging.INFO,
+                "Rebalancing %s:%s rule %s (%f GB) from %s to %s. New rule %s",
+                scope,
+                name,
+                str(rule_id),
+                bytes_ / 1e9,
+                rule["rse_expression"],
+                target_rse_exp,
+                child_rule_id,
+            )
+            rebalanced_bytes += bytes_
+            rebalanced_files += length
+            rebalanced_datasets.append(
+                (scope, name, bytes_, length, target_rse_exp, rule_id, child_rule_id)
+            )
+        except Exception as error:
+            logger(
+                logging.ERROR,
+                "Exception %s occurred while rebalancing %s:%s, rule_id: %s!",
+                str(error),
+                scope,
+                name,
+                str(rule_id),
+            )
+
+    logger(
+        logging.INFO,
+        "BB8 is rebalancing %d GB of data (%d rules) from %s",
+        rebalanced_bytes / 1e9,
+        len(rebalanced_datasets),
+        src_rse,
+    )
+    return rebalanced_datasets
+
+
+@read_session
+def get_active_locks(
+    *,
+    session: Optional[Session] = None
+) -> dict[str, dict[str, int]]:
+    locks_dict: dict[str, dict[str, int]] = {}
+    stmt = select(
+        models.ReplicationRule.id
+    ).where(
+        and_(or_(models.ReplicationRule.state == RuleState.REPLICATING,
+                 models.ReplicationRule.state == RuleState.STUCK),
+             models.ReplicationRule.comments == "Background rebalancing")
+    )
+    for rule_id in session.execute(stmt).scalars().all():    # type: ignore (session could be None)
+        stmt = select(
+            func.count(),
+            func.sum(models.ReplicaLock.bytes),
+            models.ReplicaLock.state,
+            models.ReplicaLock.rse_id
+        ).select_from(
+            models.ReplicaLock
+        ).where(
+            and_(models.ReplicaLock.rule_id == rule_id,
+                 models.ReplicaLock.state != LockState.OK)
+        ).group_by(
+            models.ReplicaLock.state,
+            models.ReplicaLock.rse_id
+        )
+        for cnt, size, _, rse_id in session.execute(stmt).all():  # type: ignore (session could be None)
+            if rse_id not in locks_dict:
+                locks_dict[rse_id] = {"bytes": 0, "locks": 0}
+            locks_dict[rse_id]["locks"] += cnt
+            locks_dict[rse_id]["bytes"] += size
+    return locks_dict