assemblyline-core 4.4.0.57__tar.gz → 4.4.0.59__tar.gz

{assemblyline-core-4.4.0.57 → assemblyline-core-4.4.0.59}/PKG-INFO

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

assemblyline-core-4.4.0.59/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.4.0.59

{assemblyline-core-4.4.0.57 → assemblyline-core-4.4.0.59}/assemblyline_core/dispatching/client.py

@@ -164,6 +164,9 @@ class DispatchClient:
             dispatcher_id = queue.pop(timeout=5)
             listed_dispatchers.discard(dispatcher_id)
 
+    def queued_submissions(self) -> list[dict]:
+        return self.submission_queue.content()
+
     def outstanding_services(self, sid) -> Optional[dict[str, int]]:
         """
         List outstanding services for a given submission and the number of file each

{assemblyline-core-4.4.0.57 → assemblyline-core-4.4.0.59}/assemblyline_core/dispatching/dispatcher.py

@@ -274,6 +274,10 @@ class Dispatcher(ThreadedCoreBase):
     def instance_assignment_size(persistent_redis, instance_id):
         return Hash(DISPATCH_TASK_ASSIGNMENT + instance_id, host=persistent_redis).length()
 
+    @staticmethod
+    def instance_assignment(persistent_redis, instance_id) -> list[str]:
+        return Hash(DISPATCH_TASK_ASSIGNMENT + instance_id, host=persistent_redis).keys()
+
     @staticmethod
     def all_queue_lengths(redis, instance_id):
         return {

{assemblyline-core-4.4.0.57 → assemblyline-core-4.4.0.59}/assemblyline_core/ingester/ingester.py

@@ -17,8 +17,8 @@ from random import random
 from typing import Any, Iterable, List, Optional, Tuple
 
 import elasticapm
-from assemblyline.common.postprocess import ActionWorker
 
+from assemblyline.common.postprocess import ActionWorker
 from assemblyline_core.server_base import ThreadedCoreBase
 from assemblyline.common.metrics import MetricsFactory
 from assemblyline.common.str_utils import dotdump, safe_str
@@ -41,6 +41,7 @@ from assemblyline.odm.models.submission import SubmissionParams, Submission as D
 from assemblyline.odm.models.alert import EXTENDED_SCAN_VALUES
 from assemblyline.odm.messages.submission import Submission as MessageSubmission, SubmissionMessage
 
+from assemblyline_core.dispatching.dispatcher import Dispatcher
 from assemblyline_core.submission_client import SubmissionClient
 from .constants import INGEST_QUEUE_NAME, drop_chance, COMPLETE_QUEUE_NAME
 
@@ -189,7 +190,8 @@ class Ingester(ThreadedCoreBase):
     def try_run(self):
         threads_to_maintain = {
             'Retries': self.handle_retries,
-            'Timeouts': self.handle_timeouts
+            'Timeouts': self.handle_timeouts,
+            'Missing': self.handle_missing,
         }
         threads_to_maintain.update({f'Complete_{n}': self.handle_complete for n in range(COMPLETE_THREADS)})
         threads_to_maintain.update({f'Ingest_{n}': self.handle_ingest for n in range(INGEST_THREADS)})
@@ -483,6 +485,64 @@ class Ingester(ThreadedCoreBase):
             self.counter.increment_execution_time('cpu_seconds', time.process_time() - cpu_mark)
             self.counter.increment_execution_time('busy_seconds', time.time() - time_mark)
 
+    def handle_missing(self) -> None:
+        """
+        Messages get dropped or only partially processed when ingester and dispatcher containers scale up and down.
+
+        This loop checks for submissions that are in two invalid states:
+         - finished but still listed as being scanned by ingester (message probably dropped by ingester)
+         - listed by ingester but unknown by dispatcher (message could have been dropped on either end)
+
+        Loading all the info needed to do these checks is a bit slow, but doing them every 5 or 15 minutes
+        per ingester shouldn't be noteworthy. While these missing messages are bound to happen from time to time
+        they should be rare. With that in mind, a warning is raised whenever this worker processes something
+        so that if a constant stream of items are falling through and getting processed here it might stand out.
+        """
+        last_round: set[str] = set()
+
+        while self.sleep(300 if last_round else 900):
+            # Get the current set of outstanding tasks
+            outstanding: dict[str, dict] = self.scanning.items()
+
+            # Get jobs being processed by dispatcher or in dispatcher queue
+            assignment: dict[str, str] = {}
+            for data in self.submit_client.dispatcher.queued_submissions():
+                assignment[data['submission']['sid']] = ''
+            for dis in Dispatcher.all_instances(self.redis_persist):
+                for key in Dispatcher.instance_assignment(self.redis_persist, dis):
+                    assignment[key] = dis
+
+            # Filter out outstanding tasks currently assigned or in queue
+            outstanding = {
+                key: doc
+                for key, doc in outstanding.items()
+                if doc["submission"]["sid"] not in assignment
+            }
+
+            unprocessed = []
+            for key, data in outstanding.items():
+                task = IngestTask(data)
+                sid = task.submission.sid
+
+                # Check if its already complete in the database
+                from_db = self.datastore.submission.get_if_exists(sid)
+                if from_db and from_db.state == "completed":
+                    self.log.warning("Completing a hanging finished submission [{}]", sid)
+                    self.completed(from_db)
+
+                # Check for items that have been in an unknown state since the last round
+                # and put it back in processing
+                elif sid in last_round:
+                    self.log.warning("Recovering a submission dispatcher hasn't processed [{}]", sid)
+                    self.submit(task)
+
+                # Otherwise defer looking at this until next iteration
+                else:
+                    unprocessed.append(sid)
+
+            # store items for next round
+            last_round = set(unprocessed)
+
     def get_groups_from_user(self, username: str) -> List[str]:
         # Reset the group cache at the top of each hour
         if time.time()//HOUR_IN_SECONDS > self._user_groups_reset:

{assemblyline-core-4.4.0.57 → assemblyline-core-4.4.0.59}/assemblyline_core/scaler/controllers/docker_ctl.py

@@ -341,6 +341,10 @@ class DockerController(ControllerInterface):
         names = list(self._profiles.keys())
         return {name: self.get_target(name) for name in names}
 
+    def get_unavailable(self) -> dict[str, int]:
+        """Get the number of containers the orchestration layer could not start."""
+        return {}
+
     def set_target(self, service_name, target):
         """Change how many instances of a service docker is trying to keep up.
 

{assemblyline-core-4.4.0.57 → assemblyline-core-4.4.0.59}/assemblyline_core/scaler/controllers/kubernetes_ctl.py

@@ -308,6 +308,7 @@ class KubernetesController(ControllerInterface):
         pod_background.start()
 
         self._deployment_targets: dict[str, int] = {}
+        self._deployment_unavailable: dict[str, int] = {}
         deployment_background = threading.Thread(target=self._loop_forever(self._monitor_deployments), daemon=True)
         deployment_background.start()
 
@@ -601,6 +602,7 @@ class KubernetesController(ControllerInterface):
         watch = TypelessWatch()
 
         self._deployment_targets = {}
+        self._deployment_unavailable = {}
         label_selector = ','.join(f'{_n}={_v}' for _n, _v in self._labels.items() if _n != 'privilege')
 
         for event in watch.stream(func=self.apps_api.list_namespaced_deployment,
@@ -613,9 +615,11 @@ class KubernetesController(ControllerInterface):
                 name = event['raw_object']['metadata']['labels'].get('component', None)
                 if name is not None:
                     self._deployment_targets[name] = event['raw_object']['spec']['replicas']
+                    self._deployment_unavailable[name] = event['raw_object']['status'].get('unavailableReplicas', 0)
             elif event['type'] == 'DELETED':
                 name = event['raw_object']['metadata']['labels'].get('component', None)
                 self._deployment_targets.pop(name, None)
+                self._deployment_unavailable.pop(name, None)
 
     def _get_pod_used_namespace_cpu(self) -> float:
         count = 0.0
@@ -875,6 +879,10 @@ class KubernetesController(ControllerInterface):
         """Get the target for running instances of all services."""
         return self._deployment_targets
 
+    def get_unavailable(self) -> dict[str, int]:
+        """Get the number of containers the orchestration layer could not start."""
+        return self._deployment_unavailable
+
     def set_target(self, service_name: str, target: int):
         """Set the target for running instances of a service."""
         for _ in range(10):

{assemblyline-core-4.4.0.57 → assemblyline-core-4.4.0.59}/assemblyline_core/scaler/scaler_server.py

@@ -272,7 +272,13 @@ class ScalerServer(ThreadedCoreBase):
         # function properly.
         for secret in re.findall(r'\${\w+}', open('/etc/assemblyline/config.yml', 'r').read()) + ['UI_SERVER']:
             env_name = secret.strip("${}")
-            core_env[env_name] = os.environ[env_name]
+            try:
+                core_env[env_name] = os.environ[env_name]
+            except KeyError:
+                # Don't pass through variables that scaler doesn't have
+                # they are likely specific to other components and shouldn't
+                # be shared with privileged services.
+                pass
 
         labels = {
             'app': 'assemblyline',
@@ -737,14 +743,9 @@ class ScalerServer(ThreadedCoreBase):
                 #       it is one big one, and let the orchestration layer sort out the details.
                 #
 
-                # Recalculate the amount of free resources expanding the total quantity by the overallocation
-                free_cpu, total_cpu = self.controller.cpu_info()
-                used_cpu = total_cpu - free_cpu
-                free_cpu = total_cpu * self.get_cpu_overallocation() - used_cpu
-
-                free_memory, total_memory = self.controller.memory_info()
-                used_memory = total_memory - free_memory
-                free_memory = total_memory * self.get_memory_overallocation() - used_memory
+                # Get the processed resource numbers
+                free_cpu, _ = self.get_cpu_info(overallocation=True)
+                free_memory, _ = self.get_memory_info(overallocation=True)
 
                 # Make adjustments to the targets until everything is satisified
                 # or we don't have the resouces to make more adjustments
@@ -785,6 +786,44 @@ class ScalerServer(ThreadedCoreBase):
                                 self.log.info(f"Scaling service {name}: {old} -> {value}")
                                 pool.call(self.controller.set_target, name, value)
 
+    def get_cpu_info(self, overallocation: bool) -> tuple[float, float]:
+        # Get the raw used resource numbers
+        free_cpu, total_cpu = self.controller.cpu_info()
+
+        # Recalculate the amount of free resources expanding the total quantity by the overallocation
+        if overallocation:
+            used_cpu = total_cpu - free_cpu
+            free_cpu = total_cpu * self.get_cpu_overallocation() - used_cpu
+
+        # Include the service containers not counted in the raw numbers because they are pending
+        for name, pending in self.controller.get_unavailable().items():
+            profile = self.profiles.get(name)
+            if not profile or not pending:
+                continue
+
+            free_cpu = free_cpu - profile.container_config.cpu_cores * pending
+
+        return (free_cpu, total_cpu)
+
+    def get_memory_info(self, overallocation: bool) -> tuple[float, float]:
+        # Get the raw used resource numbers
+        free_memory, total_memory = self.controller.memory_info()
+
+        # Recalculate the amount of free resources expanding the total quantity by the overallocation
+        if overallocation:
+            used_memory = total_memory - free_memory
+            free_memory = total_memory * self.get_memory_overallocation() - used_memory
+
+        # Include the service containers not counted in the raw numbers because they are pending
+        for name, pending in self.controller.get_unavailable().items():
+            profile = self.profiles.get(name)
+            if not profile or not pending:
+                continue
+
+            free_memory = free_memory - profile.container_config.ram_mb * pending
+
+        return (free_memory, total_memory)
+
     def sync_metrics(self):
         """Check if there are any pub-sub messages we need."""
         while self.sleep(METRIC_SYNC_INTERVAL):
@@ -889,8 +928,8 @@ class ScalerServer(ThreadedCoreBase):
                     export_metrics_once(service_name, Status, metrics, host=HOSTNAME,
                                         counter_type='scaler_status', config=self.config, redis=self.redis)
 
-                memory, memory_total = self.controller.memory_info()
-                cpu, cpu_total = self.controller.cpu_info()
+                memory, memory_total = self.get_memory_info(overallocation=False)
+                cpu, cpu_total = self.get_cpu_info(overallocation=False)
                 metrics = {
                     'memory_total': memory_total,
                     'cpu_total': cpu_total,

{assemblyline-core-4.4.0.57 → assemblyline-core-4.4.0.59}/assemblyline_core.egg-info/PKG-INFO

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

assemblyline-core-4.4.0.57/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.4.0.57