assemblyline-core 4.5.0.22__tar.gz → 4.5.0.24__tar.gz

{assemblyline-core-4.5.0.22 → assemblyline-core-4.5.0.24}/PKG-INFO

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

assemblyline-core-4.5.0.24/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.5.0.24

{assemblyline-core-4.5.0.22 → assemblyline-core-4.5.0.24}/assemblyline_core/expiry/run_expiry.py

@@ -1,16 +1,19 @@
 #!/usr/bin/env python
 from __future__ import annotations
+
 import concurrent.futures
-from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, Future, as_completed
-from concurrent.futures.process import BrokenProcessPool
+import threading
 import functools
-from typing import Callable, Optional, Union, TYPE_CHECKING
 import elasticapm
 import time
 
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, Future, as_completed
+from concurrent.futures.process import BrokenProcessPool
 from datemath import dm
+from typing import Callable, Optional, TYPE_CHECKING
 
-from assemblyline.common.isotime import epoch_to_iso, now_as_iso, iso_to_epoch
+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
@@ -150,19 +153,19 @@ class ExpiryManager(ServerBase):
             bulk.add_delete_operation(sha256)
 
         if len(file_list) > 0:
-            self.log.info(f'    Deleted associated files from the '
+            self.log.info(f'[{collection.name}] Deleted associated files from the '
                           f'{"cachestore" if "cache" in collection.name else "filestore"}...')
             collection.bulk(bulk)
             self.counter.increment(f'{collection.name}', increment_by=len(file_list))
-            self.log.info(f"    Deleted {len(file_list)} items from the datastore...")
+            self.log.info(f"[{collection.name}] Deleted {len(file_list)} items from the datastore...")
         else:
-            self.log.warning('    Expiry unable to clean up any of the files in filestore.')
+            self.log.warning(f'[{collection.name}] Expiry unable to clean up any of the files in filestore.')
 
     def _simple_delete(self, collection, delete_query, number_to_delete):
         self.heartbeat()
         collection.delete_by_query(delete_query)
         self.counter.increment(f'{collection.name}', increment_by=number_to_delete)
-        self.log.info(f"    Deleted {number_to_delete} items from the datastore...")
+        self.log.info(f"[{collection.name}] Deleted {number_to_delete} items from the datastore...")
 
     def _cleanup_canceled_submission(self, sid):
         # Allowing us at minimum 5 minutes to cleanup the submission
@@ -171,7 +174,7 @@ class ExpiryManager(ServerBase):
             self.apm_client.begin_transaction("Delete canceled submissions")
 
         # Cleaning up the submission
-        self.log.info(f"Deleting incomplete submission {sid}...")
+        self.log.info(f"[submission] Deleting incomplete submission {sid}...")
         self.datastore.delete_submission_tree_bulk(sid, self.classification, transport=self.filestore)
         self.redis_bad_sids.remove(sid)
 
@@ -181,131 +184,151 @@ class ExpiryManager(ServerBase):
         if self.apm_client:
             self.apm_client.end_transaction("canceled_submissions", 'deleted')
 
-    def run_expiry_once(self, pool: ThreadPoolExecutor):
-        now = now_as_iso()
-        reached_max = False
+    def _process_chunk(self, collection: ESCollection, start, end, final_date, number_to_delete):
+        # We assume that no records are ever inserted such that their expiry_ts is in the past.
+        # We also assume that the `end` dates are also in the past.
+        # As long as these two things are true, the set returned by this query should be consistent.
+        # The one race condition is that a record might be refreshed while the file
+        # blob would be deleted anyway, leaving a file record with no filestore object
+        delete_query = f"expiry_ts:{{{start} TO {end}]"
+
+        # check if we are dealing with an index that needs file cleanup
+        if self.config.core.expiry.delete_storage and collection.name in self.fs_hashmap:
+            # Delete associated files
+            delete_objects: list[str] = []
+            for item in collection.stream_search(delete_query, fl='id', as_obj=False):
+                self.heartbeat()
+                delete_objects.append(item['id'])
 
-        # 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)
+            # 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(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]
 
-        # Expire data
-        for collection in self.expirable_collections:
-            self.heartbeat()
+            delete_tasks = self.fs_hashmap[collection.name](delete_objects, final_date)
 
-            # Start of expiry transaction
-            if self.apm_client:
-                self.apm_client.begin_transaction("Delete expired documents")
+            # Proceed with deletion, but only after all the scheduled deletes for this
+            self.log.info(f"[{collection.name}] Scheduled {len(delete_objects)}/{number_to_delete} files to be removed")
+            self._finish_delete(collection, delete_tasks, expire_only)
 
-            if self.config.core.expiry.batch_delete:
-                final_date = dm(f"{now}||-{self.config.core.expiry.delay}h/d").float_timestamp
-            else:
-                final_date = dm(f"{now}||-{self.config.core.expiry.delay}h").float_timestamp
-            final_date_string = epoch_to_iso(final_date)
+        else:
+            # Proceed with deletion
+            self._simple_delete(collection, delete_query, number_to_delete)
 
-            # Break down the expiry window into smaller chunks of data
-            unchecked_chunks: list[tuple[float, float]] = [(self._find_expiry_start(collection), final_date)]
-            ready_chunks: dict[tuple[float, float], int] = {}
-            while unchecked_chunks and len(ready_chunks) < self.config.core.expiry.iteration_max_tasks:
-                self.heartbeat()
-                start, end = unchecked_chunks.pop()
-                chunk_size = self._count_expired(collection, start, end)
+    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()
 
-                # Empty chunks are fine
-                if chunk_size == 0:
-                    continue
+        # Start of expiry transaction
+        if self.apm_client:
+            self.apm_client.begin_transaction("Delete expired documents")
 
-                # We found a small enough chunk to
-                #  run on
-                if chunk_size < self.expiry_size:
-                    ready_chunks[(start, end)] = chunk_size
-                    continue
+        final_date = self._get_final_date()
 
-                # Break this chunk into parts
-                middle = (end + start)/2
-                unchecked_chunks.append((middle, end))
-                unchecked_chunks.append((start, middle))
+        # Break down the expiry window into smaller chunks of data
+        while len(jobs) < self.config.core.expiry.iteration_max_tasks:
 
-            # If there are still chunks we haven't checked, then we know there is more data
-            if unchecked_chunks:
-                reached_max = True
+            # Get the next chunk
+            end, number_to_delete = self._get_next_chunk(collection, start, final_date)
 
-            for (start, end), number_to_delete in ready_chunks.items():
-                self.heartbeat()
-                # We assume that no records are ever inserted such that their expiry_ts is in the past.
-                # We also assume that the `end` dates are also in the past.
-                # As long as these two things are true, the set returned by this query should be consistent.
-                # The one race condition is that a record might be refreshed while the file
-                # blob would be deleted anyway, leaving a file record with no filestore object
-                self.log.info(f"Processing collection: {collection.name}")
-                delete_query = f"expiry_ts:[{epoch_to_iso(start) if start > 0 else '*'} TO {epoch_to_iso(end)}}}"
-
-                # check if we are dealing with an index that needs file cleanup
-                if self.config.core.expiry.delete_storage and collection.name in self.fs_hashmap:
-                    # Delete associated files
-                    delete_objects: list[str] = []
-                    for item in collection.stream_search(delete_query, fl='id', as_obj=False):
-                        self.heartbeat()
-                        delete_objects.append(item['id'])
-
-                    # 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)
-                        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]
-
-                    delete_tasks = self.fs_hashmap[collection.name](delete_objects, final_date_string)
-
-                    # Proceed with deletion, but only after all the scheduled deletes for this
-                    self.log.info(f"Scheduled {len(delete_objects)}/{number_to_delete} "
-                                  f"files to be removed for: {collection.name}")
-                    pool.submit(self.log_errors(self._finish_delete), collection, delete_tasks, expire_only)
-
-                else:
-                    # Proceed with deletion
-                    pool.submit(self.log_errors(self._simple_delete),
-                                collection, delete_query, number_to_delete)
-
-            # End of expiry transaction
-            if self.apm_client:
-                self.apm_client.end_transaction(collection.name, 'deleted')
-
-        return reached_max
-
-    def _find_expiry_start(self, container: ESCollection):
-        """Find earliest expiring item in this container."""
-        rows = container.search(f"expiry_ts: [* TO {epoch_to_iso(time.time())}]",
-                                rows=1, sort='expiry_ts asc', as_obj=False, fl='expiry_ts')
+            # Check if we got anything
+            if number_to_delete == 0:
+                break
+
+            # 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
+
+        # End of expiry transaction
+        if self.apm_client:
+            self.apm_client.end_transaction(collection.name, 'deleted')
+
+        return start, number_to_delete < self.expiry_size
+
+    def _get_final_date(self):
+        now = now_as_iso()
+        if self.config.core.expiry.batch_delete:
+            final_date = dm(f"{now}||-{self.config.core.expiry.delay}h/d").float_timestamp
+        else:
+            final_date = dm(f"{now}||-{self.config.core.expiry.delay}h").float_timestamp
+        return epoch_to_iso(final_date)
+
+    def _get_next_chunk(self, collection: ESCollection, start, final_date):
+        """Find date of item at chunk size and the number of items that
+           will be affected in between start date and the date found"""
+        rows = collection.search(f"expiry_ts: {{{start} TO {final_date}]", rows=1,
+                                 offset=self.expiry_size - 1, sort='expiry_ts asc',
+                                 as_obj=False, fl='expiry_ts')
         if rows['items']:
-            return iso_to_epoch(rows['items'][0]['expiry_ts'])
-        return time.time()
-
-    def _count_expired(self, container: ESCollection, start: Union[float, str], end: float) -> int:
-        """Count how many items need to be erased in the given window."""
-        if start == 0:
-            start = '*'
-        if isinstance(start, (float, int)):
-            start = epoch_to_iso(start)
-        query = f'expiry_ts:[{start} TO {epoch_to_iso(end)}}}'
-        return container.search(query, rows=0, as_obj=False, track_total_hits=self.expiry_size)['total']
+            return rows['items'][0]['expiry_ts'], self.expiry_size
+        return final_date, rows['total']
 
     def try_run(self):
+        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)
+
+        # Wait for all the threads to exit
+        for thread in main_threads:
+            thread.join()
+
+    def clean_deleted_submissions(self, pool):
+        """Delete canceled submissions"""
         while self.running:
-            try:
-                expiry_maxed_out = False
+            # 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] = []
 
-                with ThreadPoolExecutor(self.config.core.expiry.workers) as pool:
-                    try:
-                        expiry_maxed_out = self.run_expiry_once(pool)
-                    except Exception as e:
-                        self.log.exception(str(e))
+        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 not expiry_maxed_out:
+                # 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.0.22 → assemblyline-core-4.5.0.24}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline-core-4.5.0.22 → assemblyline-core-4.5.0.24}/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.0.22/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.5.0.22