Merge pull request #306 from CybercentreCanada/scaler-performance

Scaler performance (dev)

assemblyline/remote/datatypes/queues/priority.py

@@ -1,4 +1,5 @@
 import json
+from typing import List
 
 from assemblyline.remote.datatypes import get_client, retry_call, decode
 
@@ -212,3 +213,15 @@ def select(*queues, **kw):
         return response
 
     return response[0].decode('utf-8'), json.loads(response[1][21:])
+
+
+def length(*queues: PriorityQueue) -> List[int]:
+    """Utility function for batch reading queue lengths."""
+    if not queues:
+        return []
+    pipeline = queues[0].c.pipeline(transaction=False)
+
+    for que in queues:
+        pipeline.zcard(que.name)
+
+    return retry_call(pipeline.execute)

test/test_remote_datatypes.py

@@ -180,7 +180,7 @@ def locked_execution(next_thread=None):
 
 # noinspection PyShadowingNames,PyUnusedLocal
 def test_priority_queue(redis_connection):
-    from assemblyline.remote.datatypes.queues.priority import PriorityQueue
+    from assemblyline.remote.datatypes.queues.priority import PriorityQueue, length
     with PriorityQueue('test-priority-queue') as pq:
         pq.delete()
 
@@ -212,6 +212,11 @@ def test_priority_queue(redis_connection):
         assert pq.dequeue_range(lower_limit=102, skip=1) == [2]  # 2 and 3 are both options, 3 has higher score, skip it
         assert pq.dequeue_range(upper_limit=100, num=10) == [5, 0]  # Take some off the other end
         assert pq.length() == 2
+
+        with PriorityQueue('second-priority-queue') as other:
+            other.push(100, 'a')
+            assert length(other, pq) == [1, 2]
+
         pq.pop(2)
 
         pq.push(50, 'first')