Repository: rucio/rucio

Issue:
add "thread" support to lock and replica retrieval functions
Motivation
----------
needed to address scalability problems of add_rule


Modification
------------
/lib/rucio/core/lock.py
add total_threads/thread_id support to get_files_and_replica_locks_of_dataset
/lib/rucio/core/replica.py
add total_threads/thread_id support to get_and_lock_file_replicas_for_dataset
add total_threads/thread_id support to get_source_replicas_for_dataset
 




Discussion:

Resolving patch:
diff --git a/lib/rucio/core/lock.py b/lib/rucio/core/lock.py
--- a/lib/rucio/core/lock.py
+++ b/lib/rucio/core/lock.py
@@ -11,6 +11,7 @@
 # - Vincent Garonne, <vincent.garonne@cern.ch>, 2014-2018
 # - Cedric Serfon, <cedric.serfon@cern.ch>, 2014-2018
 # - Thomas Beermann, <thomas.beermann@cern.ch>, 2014-2018
+# - Luc Goossens <luc.goossens@cern.ch>, 2020
 #
 # PY3K COMPATIBLE
 
@@ -29,7 +30,7 @@
 from rucio.common.exception import RSENotFound
 from rucio.core.lifetime_exception import define_eol
 from rucio.core.rse import get_rse_name, get_rse_id
-from rucio.db.sqla import models
+from rucio.db.sqla import models, filter_thread_work
 from rucio.db.sqla.constants import LockState, RuleState, RuleGrouping, DIDType, RuleNotification
 from rucio.db.sqla.session import read_session, transactional_session, stream_session
 
@@ -182,7 +183,9 @@ def get_replica_locks_for_rule_id_per_rse(rule_id, session=None):
 
 
 @read_session
-def get_files_and_replica_locks_of_dataset(scope, name, nowait=False, restrict_rses=None, only_stuck=False, session=None):
+def get_files_and_replica_locks_of_dataset(scope, name, nowait=False, restrict_rses=None, only_stuck=False,
+                                           total_threads=None, thread_id=None,
+                                           session=None):
     """
     Get all the files of a dataset and, if existing, all locks of the file.
 
@@ -191,6 +194,8 @@ def get_files_and_replica_locks_of_dataset(scope, name, nowait=False, restrict_r
     :param nowait:         Nowait parameter for the FOR UPDATE statement
     :param restrict_rses:  Possible RSE_ids to filter on.
     :param only_stuck:     If true, only get STUCK locks.
+    :param total_threads:  Total threads
+    :param thread_id:      This thread
     :param session:        The db session.
     :return:               Dictionary with keys: (scope, name)
                            and as value: [LockObject]
@@ -206,6 +211,10 @@ def get_files_and_replica_locks_of_dataset(scope, name, nowait=False, restrict_r
             filter(models.DataIdentifierAssociation.scope == scope,
                    models.DataIdentifierAssociation.name == name)
 
+        if total_threads and total_threads > 1:
+            content_query = filter_thread_work(session=session, query=content_query, total_threads=total_threads,
+                                               thread_id=thread_id, hash_variable='child_name')
+
         for child_scope, child_name in content_query.yield_per(1000):
             locks[(child_scope, child_name)] = []
 
@@ -263,6 +272,10 @@ def get_files_and_replica_locks_of_dataset(scope, name, nowait=False, restrict_r
     if only_stuck:
         query = query.filter(models.ReplicaLock.state == LockState.STUCK)
 
+    if total_threads and total_threads > 1:
+        query = filter_thread_work(session=session, query=query, total_threads=total_threads,
+                                   thread_id=thread_id, hash_variable='child_name')
+
     query = query.with_for_update(nowait=nowait, of=models.ReplicaLock.state)
 
     for child_scope, child_name, lock in query:
diff --git a/lib/rucio/core/replica.py b/lib/rucio/core/replica.py
--- a/lib/rucio/core/replica.py
+++ b/lib/rucio/core/replica.py
@@ -1961,6 +1961,7 @@ def get_source_replicas(scope, name, source_rses=None, session=None):
 
 @transactional_session
 def get_and_lock_file_replicas_for_dataset(scope, name, nowait=False, restrict_rses=None,
+                                           total_threads=None, thread_id=None,
                                            session=None):
     """
     Get file replicas for all files of a dataset.
@@ -1969,6 +1970,8 @@ def get_and_lock_file_replicas_for_dataset(scope, name, nowait=False, restrict_r
     :param name:           The name of the dataset.
     :param nowait:         Nowait parameter for the FOR UPDATE statement
     :param restrict_rses:  Possible RSE_ids to filter on.
+    :param total_threads:  Total threads
+    :param thread_id:      This thread
     :param session:        The db session in use.
     :returns:              (files in dataset, replicas in dataset)
     """
@@ -1986,6 +1989,10 @@ def get_and_lock_file_replicas_for_dataset(scope, name, nowait=False, restrict_r
             filter(models.DataIdentifierAssociation.scope == scope,
                    models.DataIdentifierAssociation.name == name)
 
+        if total_threads and total_threads > 1:
+            content_query = filter_thread_work(session=session, query=content_query, total_threads=total_threads,
+                                               thread_id=thread_id, hash_variable='child_name')
+
         for child_scope, child_name, bytes, md5, adler32 in content_query.yield_per(1000):
             files[(child_scope, child_name)] = {'scope': child_scope,
                                                 'name': child_name,
@@ -2072,6 +2079,10 @@ def get_and_lock_file_replicas_for_dataset(scope, name, nowait=False, restrict_r
                                    .filter(models.DataIdentifierAssociation.scope == scope,
                                            models.DataIdentifierAssociation.name == name)
 
+    if total_threads and total_threads > 1:
+        query = filter_thread_work(session=session, query=query, total_threads=total_threads,
+                                   thread_id=thread_id, hash_variable='child_name')
+
     query = query.with_for_update(nowait=nowait, of=models.RSEFileAssociation.lock_cnt)
 
     for child_scope, child_name, bytes, md5, adler32, replica in query.yield_per(1000):
@@ -2094,13 +2105,17 @@ def get_and_lock_file_replicas_for_dataset(scope, name, nowait=False, restrict_r
 
 
 @transactional_session
-def get_source_replicas_for_dataset(scope, name, source_rses=None, session=None):
+def get_source_replicas_for_dataset(scope, name, source_rses=None,
+                                    total_threads=None, thread_id=None,
+                                    session=None):
     """
     Get file replicas for all files of a dataset.
 
     :param scope:          The scope of the dataset.
     :param name:           The name of the dataset.
     :param source_rses:    Possible source RSE_ids to filter on.
+    :param total_threads:  Total threads
+    :param thread_id:      This thread
     :param session:        The db session in use.
     :returns:              (files in dataset, replicas in dataset)
     """
@@ -2132,6 +2147,10 @@ def get_source_replicas_for_dataset(scope, name, source_rses=None, session=None)
                                .filter(models.DataIdentifierAssociation.scope == scope,
                                        models.DataIdentifierAssociation.name == name)
 
+    if total_threads and total_threads > 1:
+        query = filter_thread_work(session=session, query=query, total_threads=total_threads,
+                                   thread_id=thread_id, hash_variable='child_name')
+
     replicas = {}
 
     for child_scope, child_name, rse_id in query:


Test patch: