perf: use lpop with count and don't fetch upload in Upload task (#707)

codecov · Sep 16, 2024 · 276c39f · 276c39f
1 parent 26de475
commit 276c39f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 10 deletions.
diff --git a/tasks/tests/unit/test_upload_task.py b/tasks/tests/unit/test_upload_task.py
@@ -85,14 +85,24 @@ def get(self, key):
             return res.encode()
         return res
 
-    def lpop(self, key):
+    def lpop(self, key, count=None):
         list = self.lists.get(key)
         if not list:
             return None
 
-        res = list.pop(0)
-        if list == []:
-            del self.lists[key]
+        res = None
+        if count:
+            res = []
+            for _ in range(count):
+                res.append(list.pop(0))
+                if list == []:
+                    del self.lists[key]
+                    break
+        else:
+            res = list.pop(0)
+            if list == []:
+                del self.lists[key]
+
         return res
 
     def delete(self, key):

diff --git a/tasks/upload.py b/tasks/upload.py
@@ -168,8 +168,9 @@ def arguments_list(self):
         """
         uploads_list_key = self.upload_location
         log.debug("Fetching arguments from redis %s", uploads_list_key)
-        while arguments := self.redis_connection.lpop(uploads_list_key):
-            yield loads(arguments)
+        while arguments := self.redis_connection.lpop(uploads_list_key, count=50):
+            for arg in arguments:
+                yield loads(arg)
 
     def normalize_arguments(self, commit: Commit, arguments: dict[str, Any]):
         """
@@ -495,15 +496,13 @@ def run_impl_within_lock(
         for arguments in upload_context.arguments_list():
             normalized_arguments = upload_context.normalize_arguments(commit, arguments)
             if "upload_id" in normalized_arguments:
-                upload = report_service.fetch_report_upload(
-                    commit_report, normalized_arguments["upload_id"]
-                )
+                normalized_arguments["upload_pk"] = normalized_arguments["upload_id"]
             else:
                 upload = report_service.create_report_upload(
                     normalized_arguments, commit_report
                 )
 
-            normalized_arguments["upload_pk"] = upload.id_
+                normalized_arguments["upload_pk"] = upload.id_
             argument_list.append(normalized_arguments)
 
         if argument_list: