assemblyline-core 4.5.1.dev72__tar.gz → 4.5.1.dev76__tar.gz

{assemblyline-core-4.5.1.dev72 → assemblyline-core-4.5.1.dev76}/PKG-INFO

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

assemblyline-core-4.5.1.dev76/assemblyline_core/VERSION

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

{assemblyline-core-4.5.1.dev72 → assemblyline-core-4.5.1.dev76}/assemblyline_core/metrics/heartbeat_formatter.py

@@ -7,6 +7,7 @@ from assemblyline.common.archiving import ARCHIVE_QUEUE_NAME
 from assemblyline.common.constants import DISPATCH_TASK_HASH, SUBMISSION_QUEUE, \
     SERVICE_STATE_HASH, ServiceStatus
 from assemblyline.datastore.exceptions import SearchException
+from assemblyline.odm.messages.retrohunt_heartbeat import RetrohuntMessage
 from assemblyline.odm.messages.scaler_heartbeat import ScalerMessage
 from assemblyline.odm.messages.scaler_status_heartbeat import ScalerStatusMessage
 from assemblyline.odm.messages.alerter_heartbeat import AlerterMessage
@@ -15,6 +16,7 @@ from assemblyline.odm.messages.dispatcher_heartbeat import DispatcherMessage
 from assemblyline.odm.messages.expiry_heartbeat import ExpiryMessage
 from assemblyline.odm.messages.ingest_heartbeat import IngestMessage
 from assemblyline.odm.messages.service_heartbeat import ServiceMessage
+from assemblyline.odm.messages.shard_heartbeat import ShardMessage
 from assemblyline.odm.messages.vacuum_heartbeat import VacuumMessage
 from assemblyline.remote.datatypes import get_client
 from assemblyline.remote.datatypes.hash import Hash, ExpiringHash
@@ -289,5 +291,51 @@ class HeartbeatFormatter(object):
             except Exception:
                 self.log.exception("An exception occurred while generating VacuumMessage")
 
+        elif m_type == "elastic_shards":
+            try:
+                msg = {
+                    "sender": self.sender,
+                    "msg": {
+                        "instances": instances,
+                        "request_time": m_data['request_time'],
+                        "shard_sizes": m_data['shard_sizes'],
+                    }
+                }
+                self.status_queue.publish(ShardMessage(msg).as_primitives())
+                self.log.info(f"Sent elastic shard heartbeat: {msg['msg']}")
+            except Exception:
+                self.log.exception("An exception occurred while generating ShardMessage")
+
+        elif m_type == "retrohunt":
+            try:
+
+                status = m_data.get('status', {})
+                fetcher = status.get('fetcher', {})
+                resources = status.get('resources', {})
+                storage = status.get('storage', {})
+                last_minute_cpu = sum(report.get('cpu_load_1m', 0) for report in resources.values())
+                memory = sum(report.get('memory', 0) for report in resources.values())
+                free_storage = [report.get('high_water', 0) - report.get('used', 0)
+                                for report in storage.values()]
+
+                msg = {
+                    "sender": self.sender,
+                    "msg": {
+                        'instances': instances,
+                        'request_time': m_data['request_time'],
+                        'pending_files': fetcher.get('pending_files'),
+                        'ingested_last_minute': fetcher.get('last_minute_throughput'),
+                        'worker_storage_available': min(free_storage),
+                        'total_storage_available': sum(free_storage),
+                        'active_searches': status.get('active_searches'),
+                        'last_minute_cpu': last_minute_cpu,
+                        'total_memory_used': memory,
+                    }
+                }
+                self.status_queue.publish(RetrohuntMessage(msg).as_primitives())
+                self.log.info(f"Sent retrohunt heartbeat: {msg['msg']}")
+            except Exception:
+                self.log.exception("An exception occurred while generating RetrohuntMessage")
+
         else:
             self.log.warning(f"Skipping unknown counter: {m_name} [{m_type}] ==> {m_data}")

{assemblyline-core-4.5.1.dev72 → assemblyline-core-4.5.1.dev76}/assemblyline_core/metrics/metrics_server.py

@@ -3,6 +3,7 @@
 import tempfile
 import sys
 import time
+import threading
 from collections import Counter
 from threading import Lock, Thread
 from os import environ, path
@@ -259,6 +260,7 @@ class HeartbeatManager(ServerBase):
         self.metrics_queue = CommsQueue(METRICS_QUEUE)
         self.scheduler = BackgroundScheduler(daemon=True)
         self.hm = HeartbeatFormatter("heartbeat_manager", self.log, config=self.config)
+        self.hauntedhouse_client = forge.get_hauntedhouse_client(self.config)
 
         self.counters_lock = Lock()
         self.counters = {}
@@ -281,6 +283,10 @@ class HeartbeatManager(ServerBase):
         self.scheduler.add_job(self._export_hearbeats, 'interval', seconds=self.config.core.metrics.export_interval)
         self.scheduler.start()
 
+        threading.Thread(target=self._call_interval, args=('es_shards', self._fetch_shards), daemon=True).start()
+        if self.config.retrohunt.enabled:
+            threading.Thread(target=self._call_interval, args=('retrohunt', self._fetch_retrohunt), daemon=True).start()
+
         while self.running:
             for msg in self.metrics_queue.listen():
                 # APM Transaction start
@@ -316,6 +322,64 @@ class HeartbeatManager(ServerBase):
                 if self.apm_client:
                     self.apm_client.end_transaction('process_message', 'success')
 
+    def _call_interval(self, name, method):
+        while self.running:
+            # Run the heartbeat
+            send_time = time.time()
+            try:
+                method()
+            except Exception:
+                self.log.exception('Error running metric fetcher %s', name)
+
+            # Wait until we are inline with the heartbeat interval
+            elapsed = time.time() - send_time
+            remaining = self.config.core.metrics.export_interval - elapsed
+            if remaining > 0:
+                self.sleep(remaining)     
+
+    def _fetch_shards(self):
+        request_time = None
+        sizes = {}
+        nodes = []
+        try:
+            # Pull shard data from elastisearch
+            start_time = time.time()
+            response = self.datastore.ds.client.cat.shards(bytes='b', format='json', master_timeout='5s')
+            for shard in response.body:
+                index = shard['index']
+                sizes.setdefault(index, 0)
+                sizes[index] = max(sizes[index], int(shard['store']))
+                nodes.append(shard['node'])
+            request_time = time.time() - start_time
+
+        finally:
+            # Export metrics heartbeat, send None and empty if an error occurs while getting the data
+            # the error will be logged in the outer function
+            metrics = {
+                'shard_sizes': sizes,
+                'request_time': request_time,
+            }
+            self.hm.send_heartbeat('elastic_shards', 'datastore', metrics, len(set(nodes)))
+
+    def _fetch_retrohunt(self):
+        status = {}
+        request_time = None
+        try:
+            # Pull status message from retrohunt
+            start_time = time.time()
+            status = self.hauntedhouse_client.status()
+            request_time = time.time() - start_time
+
+        finally:    
+            metrics = {
+                'request_time': request_time,
+                'status': status,
+            }
+            instances = len(status.get('storage', {}))
+
+            # Export heartbeat
+            self.hm.send_heartbeat('retrohount', 'hauntedhouse', metrics, instances)
+
     def _export_hearbeats(self):
         try:
             self.heartbeat()

{assemblyline-core-4.5.1.dev72 → assemblyline-core-4.5.1.dev76}/assemblyline_core/scaler/controllers/docker_ctl.py

@@ -121,7 +121,10 @@ class DockerController(ControllerInterface):
                                                  aliases=['service-server'])
 
                 # As long as the current service server is still running, just block its exit code in this thread
-                self.service_server.wait()
+                try:
+                    self.service_server.wait()
+                except docker.errors.NotFound:
+                    pass
 
                 # If it does return, find the new service server
                 self.service_server = self.find_service_server()

{assemblyline-core-4.5.1.dev72 → assemblyline-core-4.5.1.dev76}/assemblyline_core.egg-info/PKG-INFO

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

assemblyline-core-4.5.1.dev72/assemblyline_core/VERSION

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