Repository: rucio/rucio

Issue:
Various bugs affecting BB8
Motivation
----------
Various bugs affecting BB8 were identified






Discussion:

Resolving patch:
diff --git a/lib/rucio/daemons/bb8/bb8.py b/lib/rucio/daemons/bb8/bb8.py
--- a/lib/rucio/daemons/bb8/bb8.py
+++ b/lib/rucio/daemons/bb8/bb8.py
@@ -17,177 +17,298 @@
 BB8 is a daemon the re-balance data between RSEs.
 """
 
+import functools
 import logging
 import socket
 import threading
-import time
-import os
+from typing import TYPE_CHECKING
 
-from sqlalchemy import func, or_, and_
-from rucio.db.sqla.session import read_session
-
-from rucio.db.sqla import models
-from rucio.db.sqla.constants import RuleState, LockState
+from rucio.common.config import config_get_float
 from rucio.common.exception import InvalidRSEExpression
-from rucio.common.logging import formatted_logger, setup_logging
-from rucio.core import config as config_core
+from rucio.common.logging import setup_logging
 from rucio.core.rse_expression_parser import parse_expression
-from rucio.core.heartbeat import live, die, sanity_check, list_payload_counts
+from rucio.core.heartbeat import sanity_check, list_payload_counts
 from rucio.core.rse import get_rse_usage
-from rucio.daemons.bb8.common import rebalance_rse
+from rucio.daemons.bb8.common import rebalance_rse, get_active_locks
+from rucio.daemons.common import run_daemon
+
 
+if TYPE_CHECKING:
+    from rucio.daemons.common import HeartbeatHandler
 
-GRACEFUL_STOP = threading.Event()
+graceful_stop = threading.Event()
 
 
-def rule_rebalancer(rse_expression, move_subscriptions=False, use_dump=False, sleep_time=300, once=True, dry_run=False):
+def rule_rebalancer(
+    rse_expression: str,
+    move_subscriptions: bool = False,
+    use_dump: bool = False,
+    sleep_time: int = 300,
+    once: bool = True,
+    dry_run: bool = False,
+) -> None:
     """
-    Main loop to rebalancer rules automatically
+    Create a rule_rebalancer worker
+
+    :param rse_expression: The RSE expression where the rule rebalancing is applied.
+    :param move_subscription: To allow rebalancing of subscription rules. Not implemented yet.
+    :param use_dump: To use dump instead of DB query.
+    :param sleep_time: Time between two cycles.
+    :param once: Run only once.
+    :param dry_run: To run in dry run mode (i.e. rules are not created).
     """
-
+    run_daemon(
+        once=once,
+        graceful_stop=graceful_stop,
+        executable="rucio-bb8",
+        logger_prefix="rucio-bb8",
+        partition_wait_time=1,
+        sleep_time=sleep_time,
+        run_once_fnc=functools.partial(
+            run_once,
+            rse_expression=rse_expression,
+            move_subscriptions=move_subscriptions,
+            use_dump=use_dump,
+            dry_run=dry_run,
+        ),
+    )
+
+
+def run_once(
+    heartbeat_handler: "HeartbeatHandler",
+    rse_expression: str,
+    move_subscriptions: bool,
+    use_dump: bool,
+    dry_run: bool,
+    **_kwargs
+) -> bool:
+
+    must_sleep = False
     total_rebalance_volume = 0
-    executable = 'rucio-bb8'
-    hostname = socket.gethostname()
-    pid = os.getpid()
-    hb_thread = threading.current_thread()
-    heart_beat = live(executable, hostname, pid, hb_thread)
-    prepend_str = 'bb8[%i/%i] ' % (heart_beat['assign_thread'], heart_beat['nr_threads'])
-    logger = formatted_logger(logging.log, prepend_str + '%s')
-    logger(logging.DEBUG, 'rse_expression: %s', rse_expression)
-    logger(logging.INFO, 'BB8 started')
-
-    while not GRACEFUL_STOP.is_set():
-        logger(logging.INFO, 'Starting new cycle')
-        heart_beat = live(executable, hostname, pid, hb_thread)
-        start_time = time.time()
-        total_rebalance_volume = 0
-        tolerance = config_core.get('bb8', 'tolerance', default=0.05)
-        max_total_rebalance_volume = config_core.get('bb8', 'max_total_rebalance_volume', default=10 * 1E12)
-        max_rse_rebalance_volume = config_core.get('bb8', 'max_rse_rebalance_volume', default=500 * 1E9)
-        min_total = config_core.get('bb8', 'min_total', default=20 * 1E9)
-        payload_cnt = list_payload_counts(executable, older_than=600, hash_executable=None, session=None)
-        if rse_expression in payload_cnt:
-            logger(logging.WARNING, 'One BB8 instance already running with the same RSE expression. Stopping')
-            break
-        else:
-            # List the RSEs represented by rse_expression
+    worker_number, total_workers, logger = heartbeat_handler.live()
+    logger(logging.DEBUG, "Running BB8 on rse_expression: %s", rse_expression)
+    tolerance = config_get_float("bb8", "tolerance", default=0.05)
+    max_total_rebalance_volume = config_get_float(
+        "bb8", "max_total_rebalance_volume", default=10 * 1e12
+    )
+    max_rse_rebalance_volume = config_get_float(
+        "bb8", "max_rse_rebalance_volume", default=500 * 1e9
+    )
+    min_total = config_get_float("bb8", "min_total", default=20 * 1e9)
+    payload_cnt = list_payload_counts(
+        executable="rucio-bb8", older_than=600, hash_executable=None, session=None
+    )
+    if rse_expression in payload_cnt:
+        logger(
+            logging.WARNING,
+            "One BB8 instance already running with the same RSE expression. Stopping",
+        )
+        must_sleep = True
+        return must_sleep
+    else:
+        # List the RSEs represented by rse_expression
+        try:
+            rses = [rse for rse in parse_expression(rse_expression)]
+            list_rses2 = [rse["rse"] for rse in rses]
+        except InvalidRSEExpression as err:
+            logger(logging.ERROR, err)
+            return must_sleep
+        # List the RSEs represented by all the RSE expressions stored in heartbeat payload
+        list_rses1 = []
+        for rse_exp in payload_cnt:
+            if rse_exp:
+                list_rses1 = [rse["rse"] for rse in parse_expression(rse_exp)]
+        for rse in list_rses2:
+            if rse in list_rses1:
+                logger(
+                    logging.WARNING,
+                    "Overlapping RSE expressions %s vs %s. Stopping",
+                    rse_exp,
+                    rse_expression,
+                )
+                return must_sleep
+
+        logger(logging.INFO, "Will process rebalancing on %s", rse_expression)
+        worker_number, total_workers, logger = heartbeat_handler.live()
+        total_primary = 0
+        total_secondary = 0
+        total_total = 0
+        global_ratio = float(0)
+        for rse in rses:
+            logger(logging.DEBUG, "Getting RSE usage on %s", rse["rse"])
+            rse_usage = get_rse_usage(rse_id=rse["id"])
+            usage_dict = {}
+            for item in rse_usage:
+                # TODO Check last update
+                usage_dict[item["source"]] = {
+                    "used": item["used"],
+                    "free": item["free"],
+                    "total": item["total"],
+                }
+
             try:
-                rses = [rse for rse in parse_expression(rse_expression)]
-                list_rses2 = [rse['rse'] for rse in rses]
-            except InvalidRSEExpression as err:
-                logger(logging.ERROR, err)
+                rse["primary"] = (
+                    usage_dict["rucio"]["used"] - usage_dict["expired"]["used"]
+                )
+                rse["secondary"] = usage_dict["expired"]["used"]
+                rse["total"] = (
+                    usage_dict["storage"]["total"]
+                    - usage_dict["min_free_space"]["used"]
+                )
+                rse["ratio"] = float(rse["primary"]) / float(rse["total"])
+            except KeyError as err:
+                logger(
+                    logging.ERROR,
+                    "Missing source usage %s for RSE %s. Exiting",
+                    err,
+                    rse["rse"],
+                )
                 break
-            # List the RSEs represented by all the RSE expressions stored in heartbeat payload
-            list_rses1 = []
-            for rse_exp in payload_cnt:
-                if rse_exp:
-                    list_rses1 = [rse['rse'] for rse in parse_expression(rse_exp)]
-            for rse in list_rses2:
-                if rse in list_rses1:
-                    logger(logging.WARNING, 'Overlapping RSE expressions %s vs %s. Stopping', rse_exp, rse_expression)
-                    break
-
-            logger(logging.INFO, 'Will process rebalancing on %s', rse_expression)
-            heart_beat = live(executable, hostname, pid, hb_thread, older_than=max(600, sleep_time), hash_executable=None, payload=rse_expression, session=None)
-            total_primary = 0
-            total_secondary = 0
-            total_total = 0
-            global_ratio = float(0)
-            for rse in rses:
-                logger(logging.DEBUG, 'Getting RSE usage on %s', rse['rse'])
-                rse_usage = get_rse_usage(rse_id=rse['id'])
-                usage_dict = {}
-                for item in rse_usage:
-                    # TODO Check last update
-                    usage_dict[item['source']] = {'used': item['used'], 'free': item['free'], 'total': item['total']}
-
-                try:
-                    rse['primary'] = usage_dict['rucio']['used'] - usage_dict['expired']['used']
-                    rse['secondary'] = usage_dict['expired']['used']
-                    rse['total'] = usage_dict['storage']['total'] - usage_dict['min_free_space']['used']
-                    rse['ratio'] = float(rse['primary']) / float(rse['total'])
-                except KeyError as err:
-                    logger(logging.ERROR, 'Missing source usage %s for RSE %s. Exiting', err, rse['rse'])
-                    break
-                total_primary += rse['primary']
-                total_secondary += rse['secondary']
-                total_total += float(rse['total'])
-                rse['receive_volume'] = 0  # Already rebalanced volume in this run
-                global_ratio = float(total_primary) / float(total_total)
-                logger(logging.INFO, 'Global ratio: %f' % (global_ratio))
-
-            for rse in sorted(rses, key=lambda k: k['ratio']):
-                logger(logging.INFO, '%s Sec/Prim local ratio (%f) vs global %s', rse['rse'], rse['ratio'], global_ratio)
-            rses_over_ratio = sorted([rse for rse in rses if rse['ratio'] > global_ratio + global_ratio * tolerance], key=lambda k: k['ratio'], reverse=True)
-            rses_under_ratio = sorted([rse for rse in rses if rse['ratio'] < global_ratio - global_ratio * tolerance], key=lambda k: k['ratio'], reverse=False)
-
-            # Excluding RSEs
-            logger(logging.DEBUG, 'Excluding RSEs as destination which are too small by size:')
-            for des in rses_under_ratio:
-                if des['total'] < min_total:
-                    logger(logging.DEBUG, 'Excluding %s', des['rse'])
-                    rses_under_ratio.remove(des)
-            logger(logging.DEBUG, 'Excluding RSEs as sources which are too small by size:')
-            for src in rses_over_ratio:
-                if src['total'] < min_total:
-                    logger(logging.DEBUG, 'Excluding %s', src['rse'])
-                    rses_over_ratio.remove(src)
-            logger(logging.DEBUG, 'Excluding RSEs as destinations which are not available for write:')
-            for des in rses_under_ratio:
-                if des['availability'] & 2 == 0:
-                    logger(logging.DEBUG, 'Excluding %s', des['rse'])
-                    rses_under_ratio.remove(des)
-            logger(logging.DEBUG, 'Excluding RSEs as sources which are not available for read:')
-            for src in rses_over_ratio:
-                if src['availability'] & 4 == 0:
-                    logger(logging.DEBUG, 'Excluding %s', src['rse'])
-                    rses_over_ratio.remove(src)
-
-            # Gets the number of active transfers per location
-            dict_locks = get_active_locks(session=None)
-
-            # Loop over RSEs over the ratio
-            for index, source_rse in enumerate(rses_over_ratio):
-
-                # The volume that would be rebalanced, not real availability of the data:
-                available_source_rebalance_volume = int((source_rse['primary'] - global_ratio * source_rse['secondary']) / (global_ratio + 1))
-                if available_source_rebalance_volume > max_rse_rebalance_volume:
-                    available_source_rebalance_volume = max_rse_rebalance_volume
-                if available_source_rebalance_volume > max_total_rebalance_volume - total_rebalance_volume:
-                    available_source_rebalance_volume = max_total_rebalance_volume - total_rebalance_volume
-
-                # Select a target:
-                for destination_rse in rses_under_ratio:
-                    if available_source_rebalance_volume > 0:
-                        vo_str = ' on VO {}'.format(destination_rse['vo']) if destination_rse['vo'] != 'def' else ''
-                        if index == 0 and destination_rse['id'] in dict_locks:
-                            replicating_volume = dict_locks[destination_rse['id']]['bytes']
-                            logger(logging.DEBUG, 'Already %f TB replicating to %s%s', replicating_volume / 1E12, destination_rse['rse'], vo_str)
-                            destination_rse['receive_volume'] += replicating_volume
-                        if destination_rse['receive_volume'] >= max_rse_rebalance_volume:
-                            continue
-                        available_target_rebalance_volume = max_rse_rebalance_volume - destination_rse['receive_volume']
-                        if available_target_rebalance_volume >= available_source_rebalance_volume:
-                            available_target_rebalance_volume = available_source_rebalance_volume
-
-                        logger(logging.INFO, 'Rebalance %d TB from %s(%f) to %s(%f)%s', available_target_rebalance_volume / 1E12, source_rse['rse'], source_rse['ratio'], destination_rse['rse'], destination_rse['ratio'], vo_str)
-                        expr = destination_rse['rse']
-                        rebalance_rse(rse_id=source_rse['id'], max_bytes=available_target_rebalance_volume, dry_run=dry_run, comment='Background rebalancing', force_expression=expr, logger=logger)
-
-                        destination_rse['receive_volume'] += available_target_rebalance_volume
-                        total_rebalance_volume += available_target_rebalance_volume
-                        available_source_rebalance_volume -= available_target_rebalance_volume
-
-        if once:
-            break
-
-        end_time = time.time()
-        time_diff = end_time - start_time
-        if time_diff < sleep_time:
-            logger(logging.INFO, 'Sleeping for a while : %f seconds', sleep_time - time_diff)
-            GRACEFUL_STOP.wait(sleep_time - time_diff)
-
-    die(executable='rucio-bb8', hostname=hostname, pid=pid, thread=hb_thread)
+            total_primary += rse["primary"]
+            total_secondary += rse["secondary"]
+            total_total += float(rse["total"])
+            rse["receive_volume"] = 0  # Already rebalanced volume in this run
+            global_ratio = float(total_primary) / float(total_total)
+            logger(logging.INFO, "Global ratio: %f" % (global_ratio))
+
+        for rse in sorted(rses, key=lambda k: k["ratio"]):
+            logger(
+                logging.INFO,
+                "%s Sec/Prim local ratio (%f) vs global %s",
+                rse["rse"],
+                rse["ratio"],
+                global_ratio,
+            )
+        rses_over_ratio = sorted(
+            [
+                rse
+                for rse in rses
+                if rse["ratio"] > global_ratio + global_ratio * tolerance
+            ],
+            key=lambda k: k["ratio"],
+            reverse=True,
+        )
+        rses_under_ratio = sorted(
+            [
+                rse
+                for rse in rses
+                if rse["ratio"] < global_ratio - global_ratio * tolerance
+            ],
+            key=lambda k: k["ratio"],
+            reverse=False,
+        )
+
+        # Excluding RSEs
+        logger(
+            logging.DEBUG, "Excluding RSEs as destination which are too small by size:"
+        )
+        for des in rses_under_ratio:
+            if des["total"] < min_total:
+                logger(logging.DEBUG, "Excluding %s", des["rse"])
+                rses_under_ratio.remove(des)
+        logger(logging.DEBUG, "Excluding RSEs as sources which are too small by size:")
+        for src in rses_over_ratio:
+            if src["total"] < min_total:
+                logger(logging.DEBUG, "Excluding %s", src["rse"])
+                rses_over_ratio.remove(src)
+        logger(
+            logging.DEBUG,
+            "Excluding RSEs as destinations which are not available for write:",
+        )
+        for des in rses_under_ratio:
+            if des["availability"] & 2 == 0:
+                logger(logging.DEBUG, "Excluding %s", des["rse"])
+                rses_under_ratio.remove(des)
+        logger(
+            logging.DEBUG, "Excluding RSEs as sources which are not available for read:"
+        )
+        for src in rses_over_ratio:
+            if src["availability"] & 4 == 0:
+                logger(logging.DEBUG, "Excluding %s", src["rse"])
+                rses_over_ratio.remove(src)
+
+        # Gets the number of active transfers per location
+        dict_locks = get_active_locks(session=None)
+
+        # Loop over RSEs over the ratio
+        for index, source_rse in enumerate(rses_over_ratio):
+
+            # The volume that would be rebalanced, not real availability of the data:
+            available_source_rebalance_volume = int(
+                (source_rse["primary"] - global_ratio * source_rse["secondary"])
+                / (global_ratio + 1)
+            )
+            if available_source_rebalance_volume > max_rse_rebalance_volume:
+                available_source_rebalance_volume = max_rse_rebalance_volume
+            if (
+                available_source_rebalance_volume
+                > max_total_rebalance_volume - total_rebalance_volume
+            ):
+                available_source_rebalance_volume = (
+                    max_total_rebalance_volume - total_rebalance_volume
+                )
+
+            # Select a target:
+            for destination_rse in rses_under_ratio:
+                if available_source_rebalance_volume > 0:
+                    vo_str = (
+                        " on VO {}".format(destination_rse["vo"])
+                        if destination_rse["vo"] != "def"
+                        else ""
+                    )
+                    if index == 0 and destination_rse["id"] in dict_locks:
+                        replicating_volume = dict_locks[destination_rse["id"]]["bytes"]
+                        logger(
+                            logging.DEBUG,
+                            "Already %f TB replicating to %s%s",
+                            replicating_volume / 1e12,
+                            destination_rse["rse"],
+                            vo_str,
+                        )
+                        destination_rse["receive_volume"] += replicating_volume
+                    if destination_rse["receive_volume"] >= max_rse_rebalance_volume:
+                        continue
+                    available_target_rebalance_volume = (
+                        max_rse_rebalance_volume - destination_rse["receive_volume"]
+                    )
+                    if (
+                        available_target_rebalance_volume
+                        >= available_source_rebalance_volume
+                    ):
+                        available_target_rebalance_volume = (
+                            available_source_rebalance_volume
+                        )
+
+                    logger(
+                        logging.INFO,
+                        "Rebalance %d TB from %s(%f) to %s(%f)%s",
+                        available_target_rebalance_volume / 1e12,
+                        source_rse["rse"],
+                        source_rse["ratio"],
+                        destination_rse["rse"],
+                        destination_rse["ratio"],
+                        vo_str,
+                    )
+                    expr = destination_rse["rse"]
+                    rebalance_rse(
+                        rse_id=source_rse["id"],
+                        max_bytes=available_target_rebalance_volume,
+                        dry_run=dry_run,
+                        comment="Background rebalancing",
+                        force_expression=expr,
+                        logger=logger,
+                    )
+
+                    destination_rse[
+                        "receive_volume"
+                    ] += available_target_rebalance_volume
+                    total_rebalance_volume += available_target_rebalance_volume
+                    available_source_rebalance_volume -= (
+                        available_target_rebalance_volume
+                    )
+
+    must_sleep = True
+    return must_sleep
 
 
 def stop(signum=None, frame=None):
@@ -195,42 +316,39 @@ def stop(signum=None, frame=None):
     Graceful exit.
     """
 
-    GRACEFUL_STOP.set()
+    graceful_stop.set()
 
 
-def run(once, rse_expression, move_subscriptions=False, use_dump=False, sleep_time=300, threads=1, dry_run=False):
+def run(
+    once: bool,
+    rse_expression: str,
+    move_subscriptions: bool = False,
+    use_dump: bool = False,
+    sleep_time: int = 300,
+    threads: int = 1,
+    dry_run: bool = False,
+) -> None:
     """
     Starts up the BB8 rebalancing threads.
     """
 
     setup_logging()
     hostname = socket.gethostname()
-    sanity_check(executable='rucio-bb8', hostname=hostname)
-
-    if once:
-        rule_rebalancer(rse_expression=rse_expression, move_subscriptions=move_subscriptions, use_dump=use_dump, once=once)
-    else:
-        logging.info('BB8 starting %s threads', str(threads))
-        threads = [threading.Thread(target=rule_rebalancer, kwargs={'once': once, 'rse_expression': rse_expression, 'sleep_time': sleep_time, 'dry_run': dry_run}) for _ in range(0, threads)]
-        [thread.start() for thread in threads]
-        # Interruptible joins require a timeout.
-        while threads[0].is_alive():
-            [thread.join(timeout=3.14) for thread in threads]
-
-
-@read_session
-def get_active_locks(session=None):
-    locks_dict = {}
-    rule_ids = session.query(models.ReplicationRule.id).filter(or_(models.ReplicationRule.state == RuleState.REPLICATING, models.ReplicationRule.state == RuleState.STUCK),
-                                                               models.ReplicationRule.comments == 'Background rebalancing').all()
-    for row in rule_ids:
-        rule_id = row[0]
-        query = session.query(func.count(), func.sum(models.ReplicaLock.bytes), models.ReplicaLock.state, models.ReplicaLock.rse_id).\
-            filter(and_(models.ReplicaLock.rule_id == rule_id, models.ReplicaLock.state != LockState.OK)).group_by(models.ReplicaLock.state, models.ReplicaLock.rse_id)
-        for lock in query.all():
-            cnt, size, _, rse_id = lock
-            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
+    sanity_check(executable="rucio-bb8", hostname=hostname)
+    logging.info("BB8 starting %s threads", str(threads))
+    threads = [
+        threading.Thread(
+            target=rule_rebalancer,
+            kwargs={
+                "once": once,
+                "rse_expression": rse_expression,
+                "sleep_time": sleep_time,
+                "dry_run": dry_run,
+            },
+        )
+        for _ in range(0, threads)
+    ]
+    [thread.start() for thread in threads]
+    # Interruptible joins require a timeout.
+    while threads[0].is_alive():
+        [thread.join(timeout=3.14) for thread in threads]
diff --git a/lib/rucio/daemons/bb8/common.py b/lib/rucio/daemons/bb8/common.py
--- a/lib/rucio/daemons/bb8/common.py
+++ b/lib/rucio/daemons/bb8/common.py
@@ -29,18 +29,30 @@
 from rucio.core.rse import list_rse_attributes, get_rse_name, get_rse_vo
 from rucio.core.rse_selector import RSESelector
 from rucio.common.config import config_get, config_get_int, config_get_bool
-from rucio.common.exception import (InsufficientTargetRSEs, RuleNotFound, DuplicateRule,
-                                    InsufficientAccountLimit)
-from rucio.common.types import InternalAccount
+from rucio.common.exception import (
+    InsufficientTargetRSEs,
+    RuleNotFound,
+    DuplicateRule,
+    InsufficientAccountLimit,
+)
+from rucio.common.types import InternalAccount, InternalScope
 
 from rucio.db.sqla.session import transactional_session, read_session
 from rucio.db.sqla import models
-from rucio.db.sqla.constants import (DIDType, RuleState, RuleGrouping)
+from rucio.db.sqla.constants import DIDType, RuleState, RuleGrouping, LockState
 from requests import get
 
 
 @transactional_session
-def rebalance_rule(parent_rule, activity, rse_expression, priority, source_replica_expression='*\\bb8-enabled=false', comment=None, session=None):
+def rebalance_rule(
+    parent_rule,
+    activity,
+    rse_expression,
+    priority,
+    source_replica_expression="*\\bb8-enabled=false",
+    comment=None,
+    session=None,
+):
     """
     Rebalance a replication rule to a new RSE
     :param parent_rule:                Replication rule to be rebalanced.
@@ -52,50 +64,58 @@ def rebalance_rule(parent_rule, activity, rse_expression, priority, source_repli
     :returns:                          The new child rule id.
     """
 
-    if parent_rule['expires_at'] is None:
+    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'
+        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'
+        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']
+    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)
+    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
 
 
@@ -112,28 +132,45 @@ def __dump_url(rse_id, logger=logging.log):
     # 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)
+    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'))
+            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}
+        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')
+            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')]
+        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, default='http://rucio-analytix.cern.ch:8080/LOCKS/GetFileFromHDFS?date=${date}&rse=${rse}')
+    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})
+        url = url_template.substitute({"date": d, "rse": rse, "vo": vo})
         urls.append(url)
     return urls
 
@@ -153,7 +190,7 @@ def _list_rebalance_rule_candidates_dump(rse_id, mode=None, logger=logging.log):
     rse_dump_urls.reverse()
     resp = None
     if not rse_dump_urls:
-        logger(logging.DEBUG, 'URL of the dump was not built from template.')
+        logger(logging.DEBUG, "URL of the dump was not built from template.")
         return candidates
     success = False
     while not success and len(rse_dump_urls):
@@ -162,50 +199,56 @@ def _list_rebalance_rule_candidates_dump(rse_id, mode=None, logger=logging.log):
         if resp:
             success = True
     if not resp or resp is None:
-        logger(logging.WARNING, 'RSE dump not available')
+        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')
+            _, _, 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'}
+                    rules[rule_id] = {"state": "DELETED"}
                     logger(logging.ERROR, str(err))
                     continue
-                if rule_info['child_rule_id']:
-                    rules[rule_id] = {'state': 'DELETED'}
+                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)
+                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':
+        if mode == "decommission":  # other modes can be added later
+            if rules[r_id]["state"] == "DELETED":
                 continue
-            if int(rules[r_id]['length']) == 0:
+            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'])))
+            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
 
 
@@ -221,100 +264,239 @@ def list_rebalance_rule_candidates(rse_id, mode=None, session=None):
     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':
+    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]
+    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)
+    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))
+        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)
+    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.now() - 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)
+    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(',')]
+        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
+        "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)
+    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))
+        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]
+        "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)
+    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))
+        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)
+    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)  # NOQA
 
     # Now build the query
     external_dsl = aliased(models.DatasetLock)
-    count_locks = select([func.count()]).where(and_(external_dsl.scope == models.DatasetLock.scope,
-                                                    external_dsl.name == models.DatasetLock.name,
-                                                    external_dsl.rse_id == models.DatasetLock.rse_id)).as_scalar()
-    query = session.query(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)).\
-        filter(models.DatasetLock.rse_id == rse_id).\
-        filter(and_(*rule_clause)).\
-        filter(and_(*did_clause)).\
-        filter(case([(or_(models.DatasetLock.length < 1, models.DatasetLock.length.is_(None)), 0)],
-                    else_=cast(models.DatasetLock.bytes / models.DatasetLock.length, BigInteger)) > 1000000000).\
-        filter(count_locks == 1)
-    summary = query.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).all()
+    count_locks = (
+        select([func.count()])
+        .where(
+            and_(
+                external_dsl.scope == models.DatasetLock.scope,
+                external_dsl.name == models.DatasetLock.name,
+                external_dsl.rse_id == models.DatasetLock.rse_id,
+            )
+        )
+        .as_scalar()
+    )
+    query = (
+        session.query(
+            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,
+            ),
+        )
+        .filter(models.DatasetLock.rse_id == rse_id)
+        .filter(and_(*rule_clause))
+        .filter(and_(*did_clause))
+        .filter(
+            case(
+                [
+                    (
+                        or_(
+                            models.DatasetLock.length < 1,
+                            models.DatasetLock.length.is_(None),
+                        ),
+                        0,
+                    )
+                ],
+                else_=cast(
+                    models.DatasetLock.bytes / models.DatasetLock.length, BigInteger
+                ),
+            )
+            > 1000000000
+        )
+        .filter(count_locks == 1)
+    )
+    summary = query.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,
+    ).all()
     return summary
 
 
 @read_session
-def select_target_rse(parent_rule, current_rse_id, rse_expression, subscription_id, rse_attributes, other_rses=[], exclude_expression=None, force_expression=None, session=None):
+def select_target_rse(
+    parent_rule,
+    current_rse_id,
+    rse_expression,
+    subscription_id,
+    rse_attributes,
+    other_rses=[],
+    exclude_expression=None,
+    force_expression=None,
+    session=None,
+):
     """
     Select a new target RSE for a rebalanced rule.
     :param parent_rule           rule that is rebalanced.
@@ -332,39 +514,89 @@ def select_target_rse(parent_rule, current_rse_id, rse_expression, subscription_
     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
+    vo = parent_rule["scope"].vo
     if exclude_expression:
-        target_rse = '((%s)|(%s))' % (exclude_expression, current_rse_expr)
+        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 = [
+        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)
+    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)
+        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)
+            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 = [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'])
+            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)
+            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, max_bytes=1E9, max_files=None, dry_run=False, exclude_expression=None, comment=None, force_expression=None, mode=None, priority=3, source_replica_expression='*\\bb8-enabled=false', session=None, logger=logging.log):
+def rebalance_rse(
+    rse_id,
+    max_bytes=1e9,
+    max_files=None,
+    dry_run=False,
+    exclude_expression=None,
+    comment=None,
+    force_expression=None,
+    mode=None,
+    priority=3,
+    source_replica_expression="*\\bb8-enabled=false",
+    session=None,
+    logger=logging.log,
+):
     """
     Rebalance data from an RSE
     :param rse_id:                     RSE to rebalance data from.
@@ -388,13 +620,22 @@ def rebalance_rse(rse_id, max_bytes=1E9, max_files=None, dry_run=False, exclude_
     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):
+    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
 
@@ -406,40 +647,122 @@ def rebalance_rse(rse_id, max_bytes=1E9, max_files=None, dry_run=False, exclude_
 
         try:
             rule = get_rule(rule_id=rule_id)
-            other_rses = [r['rse_id'] for r in get_dataset_locks(scope, name, session=session)]
+            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)
+                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)
+                    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:
+                    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)
+                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)
+            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))
+            rebalanced_datasets.append(
+                (scope, name, bytes_, length, target_rse_exp, rule_id, child_rule_id)
+            )
         except Exception as error:
-            logger(logging.ERROR, 'Exception %s occured 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)
+            logger(
+                logging.ERROR,
+                "Exception %s occured 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=None):
+    locks_dict = {}
+    rule_ids = (
+        session.query(models.ReplicationRule.id)
+        .filter(
+            or_(
+                models.ReplicationRule.state == RuleState.REPLICATING,
+                models.ReplicationRule.state == RuleState.STUCK,
+            ),
+            models.ReplicationRule.comments == "Background rebalancing",
+        )
+        .all()
+    )
+    for row in rule_ids:
+        rule_id = row[0]
+        query = (
+            session.query(
+                func.count(),
+                func.sum(models.ReplicaLock.bytes),
+                models.ReplicaLock.state,
+                models.ReplicaLock.rse_id,
+            )
+            .filter(
+                and_(
+                    models.ReplicaLock.rule_id == rule_id,
+                    models.ReplicaLock.state != LockState.OK,
+                )
+            )
+            .group_by(models.ReplicaLock.state, models.ReplicaLock.rse_id)
+        )
+        for lock in query.all():
+            cnt, size, _, rse_id = lock
+            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


Test patch: