assemblyline-core 4.5.1.dev134__tar.gz → 4.5.1.dev136__tar.gz

{assemblyline-core-4.5.1.dev134 → assemblyline-core-4.5.1.dev136}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: assemblyline-core
-Version: 4.5.1.dev134
+Version: 4.5.1.dev136
 Summary: Assemblyline 4 - Core components
 Home-page: https://github.com/CybercentreCanada/assemblyline-core/
 Author: CCCS Assemblyline development team

assemblyline-core-4.5.1.dev136/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.5.1.dev136

{assemblyline-core-4.5.1.dev134 → assemblyline-core-4.5.1.dev136}/assemblyline_core/expiry/run_expiry.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import concurrent.futures
+import threading
 import functools
 import elasticapm
 import time
@@ -12,6 +13,7 @@ from datemath import dm
 from typing import Callable, Optional, TYPE_CHECKING
 
 from assemblyline.common.isotime import epoch_to_iso, now_as_iso
+from assemblyline.datastore.collection import Index
 from assemblyline_core.server_base import ServerBase
 from assemblyline_core.dispatching.dispatcher import BAD_SID_HASH
 from assemblyline.common import forge
@@ -202,7 +204,7 @@ class ExpiryManager(ServerBase):
             # Filter archived documents if archive filestore is the same as the filestore
             expire_only = []
             if self.same_storage and self.config.datastore.archive.enabled and collection.name == 'file':
-                archived_files = self.datastore.file.multiexists_in_archive(delete_objects)
+                archived_files = self.datastore.file.multiexists(delete_objects, index_type=Index.ARCHIVE)
                 delete_objects = [k for k, v in archived_files.items() if not v]
                 expire_only = [k for k, v in archived_files.items() if v]
 
@@ -217,56 +219,39 @@ class ExpiryManager(ServerBase):
             # Proceed with deletion
             self._simple_delete(collection, delete_query, number_to_delete)
 
-    def run_expiry_once(self, pool: ThreadPoolExecutor):
-        busy_iteration = False
+    def feed_expiry_jobs(self, collection, start, jobs: list[concurrent.futures.Future], pool: ThreadPoolExecutor) -> tuple[str, bool]:
+        _process_chunk = self.log_errors(self._process_chunk)
+        number_to_delete = 0
+        self.heartbeat()
+        
+        # Start of expiry transaction
+        if self.apm_client:
+            self.apm_client.begin_transaction("Delete expired documents")
 
-        # Delete canceled submissions
-        # Make sure we're not dedicating more then a quarter of the pool to this operation because it is costly
-        for submission in self.datastore.submission.search(
-                "to_be_deleted:true", fl="sid", rows=max(1, int(self.config.core.expiry.workers / 4)))['items']:
-            if submission.sid not in self.current_submission_cleanup:
-                self.current_submission_cleanup.add(submission.sid)
-                pool.submit(self.log_errors(self._cleanup_canceled_submission), submission.sid)
+        final_date = self._get_final_date()
 
-        # Expire data
-        for collection in self.expirable_collections:
+        # Break down the expiry window into smaller chunks of data
+        while len(jobs) < self.config.core.expiry.iteration_max_tasks:
             self.heartbeat()
 
-            # Start of expiry transaction
-            if self.apm_client:
-                self.apm_client.begin_transaction("Delete expired documents")
-
-            final_date = self._get_final_date()
-
-            # Break down the expiry window into smaller chunks of data
-            start = "*"
-            iterations = 0
-            while iterations < self.config.core.expiry.iteration_max_tasks:
-                self.heartbeat()
-
-                # Get the next chunk
-                end, number_to_delete = self._get_next_chunk(collection, start, final_date)
-
-                # Check if we got anything
-                if number_to_delete == 0:
-                    break
+            # Get the next chunk
+            end, number_to_delete = self._get_next_chunk(collection, start, final_date)
 
-                # Tell the outer loop not to sleep between runs
-                if number_to_delete >= self.expiry_size:
-                    busy_iteration = True
+            # Check if we got anything
+            if number_to_delete == 0:
+                break
 
-                # Process the chunk in the threadpool
-                pool.submit(self.log_errors(self._process_chunk), collection, start, end, final_date, number_to_delete)
+            # Process the chunk in the threadpool
+            jobs.append(pool.submit(_process_chunk, collection, start, end, final_date, number_to_delete))
 
-                # Prepare for next chunk
-                start = end
-                iterations += 1
+            # Prepare for next chunk
+            start = end
 
-            # End of expiry transaction
-            if self.apm_client:
-                self.apm_client.end_transaction(collection.name, 'deleted')
+        # End of expiry transaction
+        if self.apm_client:
+            self.apm_client.end_transaction(collection.name, 'deleted')
 
-        return busy_iteration
+        return start, number_to_delete < self.expiry_size
 
     def _get_final_date(self):
         now = now_as_iso()
@@ -287,17 +272,65 @@ class ExpiryManager(ServerBase):
         return final_date, rows['total']
 
     def try_run(self):
-        while self.running:
-            try:
-                busy_iteration = False
+        pool = ThreadPoolExecutor(self.config.core.expiry.workers)
+        main_threads = []
+
+        # Launch a thread that will expire submissions that have been deleted
+        thread = threading.Thread(target=self.clean_deleted_submissions, args=[pool])
+        thread.start()
+        main_threads.append(thread)
+
+        # Launch threads that expire data from each collection of data
+        for collection in self.expirable_collections:
+            thread = threading.Thread(target=self.run_collection, args=[pool, collection])
+            thread.start()
+            main_threads.append(thread)
 
-                with ThreadPoolExecutor(self.config.core.expiry.workers) as pool:
-                    try:
-                        busy_iteration = self.run_expiry_once(pool)
-                    except Exception as e:
-                        self.log.exception(str(e))
+        # Wait for all the threads to exit
+        for thread in main_threads:
+            thread.join()
 
-                if not busy_iteration:
+    def clean_deleted_submissions(self, pool):
+        """Delete canceled submissions"""
+        while self.running:
+            # Make sure we're not dedicating more then a quarter of the pool to this operation because it is costly
+            for submission in self.datastore.submission.search(
+                    "to_be_deleted:true", fl="sid", rows=max(1, int(self.config.core.expiry.workers / 4)))['items']:
+                if submission.sid not in self.current_submission_cleanup:
+                    self.current_submission_cleanup.add(submission.sid)
+                    pool.submit(self.log_errors(self._cleanup_canceled_submission), submission.sid)
+            self.sleep_with_heartbeat(self.config.core.expiry.sleep_time)
+
+    def run_collection(self, pool: concurrent.futures.ThreadPoolExecutor, collection):
+        """Feed batches of jobs to delete to the thread pool for the given collection."""
+        start = "*"
+        jobs: list[concurrent.futures.Future] = []
+
+        while self.running:
+            try:
+                try:
+                    # Fill up 'jobs' with tasks that have been sent to the thread pool
+                    # 'jobs' may already have items in it, but 'start' makes sure the new
+                    # task added starts where the last finshed
+                    start, final_job_small = self.feed_expiry_jobs(collection, start, jobs, pool)
+
+                    # Wait until some of our work finishes and there is room in the queue for more work
+                    finished, _jobs = concurrent.futures.wait(jobs, return_when=concurrent.futures.FIRST_COMPLETED)
+                    jobs = list(_jobs)
+                    for job in finished:
+                        job.result()
+
+                    # If we have expired all the data reset the start pointer
+                    if len(jobs) == 0:
+                        start = '*'
+
+                except Exception as e:
+                    self.log.exception(str(e))
+                    continue
+
+                # IF the most recent job added to the jobs list is short then
+                # all the data is currently queued up to delete and we can sleep
+                if final_job_small:
                     self.sleep_with_heartbeat(self.config.core.expiry.sleep_time)
 
             except BrokenProcessPool:

{assemblyline-core-4.5.1.dev134 → assemblyline-core-4.5.1.dev136}/assemblyline_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: assemblyline-core
-Version: 4.5.1.dev134
+Version: 4.5.1.dev136
 Summary: Assemblyline 4 - Core components
 Home-page: https://github.com/CybercentreCanada/assemblyline-core/
 Author: CCCS Assemblyline development team

{assemblyline-core-4.5.1.dev134 → assemblyline-core-4.5.1.dev136}/test/test_expiry.py

@@ -63,7 +63,8 @@ def test_expire_all(ds_expiry):
     expiry.running = True
     expiry.counter = FakeCounter()
     with concurrent.futures.ThreadPoolExecutor(5) as pool:
-        expiry.run_expiry_once(pool)
+        for collection in expiry.expirable_collections:
+            expiry.feed_expiry_jobs(collection=collection, pool=pool, start='*', jobs=[])
 
     for k, v in expiry_collections_len.items():
         assert v == expiry.counter.get(k)

assemblyline-core-4.5.1.dev134/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.5.1.dev134