thread pool executor for deletion

closeio · Sep 25, 2024 · aad3b3a · aad3b3a
1 parent 487e8ab
commit aad3b3a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/bin/remove-message-attachments.py b/bin/remove-message-attachments.py
@@ -7,6 +7,7 @@
 import enum
 import logging
 from collections.abc import Iterable
+from concurrent.futures import ThreadPoolExecutor
 
 import click
 from sqlalchemy.orm import Query
@@ -126,6 +127,8 @@ def run(
     assert batch_size > 0
     assert delete_batch_size > 0
 
+    delete_executor = ThreadPoolExecutor(max_workers=10)
+
     for repetition in range(repeat):
         blocks = find_blocks(
             limit,
@@ -167,7 +170,7 @@ def run(
                 delete_sha256s.add(block.data_sha256)
 
             if len(delete_sha256s) >= delete_batch_size:
-                delete_batch(delete_sha256s, dry_run)
+                delete_executor.submit(delete_batch, delete_sha256s.copy(), dry_run)
                 delete_sha256s.clear()
 
         delete_batch(delete_sha256s, dry_run)
@@ -177,6 +180,8 @@ def run(
 
         after_id = max_id + 1
 
+    delete_executor.shutdown(wait=True)
+
 
 if __name__ == "__main__":
     run()