From cd17119cb1234f45cb59e91279189926b4f917d0 Mon Sep 17 00:00:00 2001 From: Hardy-Cooper Date: Thu, 18 Mar 2021 15:34:55 -0400 Subject: [PATCH 1/5] Adding queueing by tag --- cuckoo/core/scheduler.py | 2 +- cuckoo/machinery/kvm.py | 21 ++++++++++++++++++++- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/cuckoo/core/scheduler.py b/cuckoo/core/scheduler.py index edd6214e41..6e79438dde 100644 --- a/cuckoo/core/scheduler.py +++ b/cuckoo/core/scheduler.py @@ -162,7 +162,7 @@ def acquire_machine(self): # In some cases it's possible that we enter this loop without # having any available machines. We should make sure this is not # such case, or the analysis task will fail completely. - if not machinery.availables(): + if not machinery.availables(tags=self.task.tags): machine_lock.release() time.sleep(1) continue diff --git a/cuckoo/machinery/kvm.py b/cuckoo/machinery/kvm.py index 9dd84bd0b8..037e415cb6 100644 --- a/cuckoo/machinery/kvm.py +++ b/cuckoo/machinery/kvm.py @@ -5,6 +5,7 @@ from cuckoo.common.abstracts import LibVirtMachinery from cuckoo.common.exceptions import CuckooCriticalError from cuckoo.common.exceptions import CuckooMachineError +from cuckoo.core.database import Machine try: import libvirt @@ -33,4 +34,22 @@ def _connect(self): def _disconnect(self, conn): """Disconnect, ignore request to disconnect.""" - pass \ No newline at end of file + pass + + def availables(self, label=None, platform=None, tags=None): + if all(param is None for param in [label, platform, tags]): + return super(KVM, self).availables() + else: + return self._get_specific_availables(label=label, platform=platform, tags=tags) + + def _get_specific_availables(self, label=None, platform=None, tags=None): + session = self.db.Session() + machines = session.query(Machine) + if label: + machines = machines.filter_by(label=label) + elif platform: + machines = machines.filter_by(platform=platform) + elif tags: + for tag in tags: + machines = machines.filter(Machine.tags.any(name=tag)) + return machines.count() From a7ff05d27f6885f54a7330d049534622313cb407 Mon Sep 17 00:00:00 2001 From: Hardy-Cooper Date: Thu, 18 Mar 2021 16:03:06 -0400 Subject: [PATCH 2/5] Safety first! Also adding platform keyword arg --- cuckoo/core/scheduler.py | 2 +- cuckoo/machinery/kvm.py | 33 ++++++++++++++++++++------------- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/cuckoo/core/scheduler.py b/cuckoo/core/scheduler.py index 6e79438dde..6803b32a18 100644 --- a/cuckoo/core/scheduler.py +++ b/cuckoo/core/scheduler.py @@ -162,7 +162,7 @@ def acquire_machine(self): # In some cases it's possible that we enter this loop without # having any available machines. We should make sure this is not # such case, or the analysis task will fail completely. - if not machinery.availables(tags=self.task.tags): + if not machinery.availables(platform=self.task.platform, tags=self.task.tags): machine_lock.release() time.sleep(1) continue diff --git a/cuckoo/machinery/kvm.py b/cuckoo/machinery/kvm.py index 037e415cb6..1a32a6cee5 100644 --- a/cuckoo/machinery/kvm.py +++ b/cuckoo/machinery/kvm.py @@ -2,10 +2,13 @@ # Copyright (C) 2014-2016 Cuckoo Foundation. # This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org # See the file 'docs/LICENSE' for copying permission. +import logging from cuckoo.common.abstracts import LibVirtMachinery from cuckoo.common.exceptions import CuckooCriticalError from cuckoo.common.exceptions import CuckooMachineError from cuckoo.core.database import Machine +from sqlalchemy.exc import SQLAlchemyError +log = logging.getLogger(__name__) try: import libvirt @@ -36,20 +39,24 @@ def _disconnect(self, conn): """Disconnect, ignore request to disconnect.""" pass - def availables(self, label=None, platform=None, tags=None): - if all(param is None for param in [label, platform, tags]): + def availables(self, platform=None, tags=None): + if all(param is None for param in [platform, tags]): return super(KVM, self).availables() else: - return self._get_specific_availables(label=label, platform=platform, tags=tags) + return self._get_specific_availables(platform=platform, tags=tags) - def _get_specific_availables(self, label=None, platform=None, tags=None): + def _get_specific_availables(self, platform=None, tags=None): session = self.db.Session() - machines = session.query(Machine) - if label: - machines = machines.filter_by(label=label) - elif platform: - machines = machines.filter_by(platform=platform) - elif tags: - for tag in tags: - machines = machines.filter(Machine.tags.any(name=tag)) - return machines.count() + try: + machines = session.query(Machine) + if platform: + machines = machines.filter_by(platform=platform) + elif tags: + for tag in tags: + machines = machines.filter(Machine.tags.any(name=tag)) + return machines.count() + except SQLAlchemyError as e: + log.exception("Database error getting specific available machines: {0}".format(e)) + return 0 + finally: + session.close() From a60c78da181e6e37cbc808f613bdf8b42a79e1bf Mon Sep 17 00:00:00 2001 From: Hardy-Cooper Date: Fri, 19 Mar 2021 11:22:51 -0400 Subject: [PATCH 3/5] Updating to support machine names --- cuckoo/core/scheduler.py | 36 +++++++++++++++++++++++++++++++++--- cuckoo/machinery/kvm.py | 17 ++++++++++------- 2 files changed, 43 insertions(+), 10 deletions(-) diff --git a/cuckoo/core/scheduler.py b/cuckoo/core/scheduler.py index 6803b32a18..fadd0e772a 100644 --- a/cuckoo/core/scheduler.py +++ b/cuckoo/core/scheduler.py @@ -20,7 +20,7 @@ ) from cuckoo.common.objects import File from cuckoo.common.files import Folders -from cuckoo.core.database import Database, TASK_COMPLETED, TASK_REPORTED +from cuckoo.core.database import Database, TASK_COMPLETED, TASK_REPORTED, TASK_PENDING, TASK_RUNNING from cuckoo.core.guest import GuestManager from cuckoo.core.plugins import RunAuxiliary, RunProcessing from cuckoo.core.plugins import RunSignatures, RunReporting @@ -1076,8 +1076,38 @@ def start(self): if task: break - if machine.is_analysis(): - available = True + # Since there are no tasks meant for a specific AVAILABLE machine, + # check if any tasks meant for a specific LOCKED machine, specific platform or specific tag + if not task: + tasks = self.db.list_tasks(status=TASK_PENDING) + if not tasks: + continue + + # If so, dispatch to an available machine + for task in tasks: + task_details = self.db.view_task(task.id) + # Note that label > platform > tags + if task_details.machine: + if machinery.availables(label=task_details.machine): + available = True + break + elif task_details.platform: + if machinery.availables(platform=task_details.platform): + available = True + break + elif task_details.tags: + tag_names = [tag.name for tag in task_details.tags] + if machinery.availables(tags=tag_names): + available = True + break + else: + available = True + break + + if task and available: + self.db.set_status(task.id, TASK_RUNNING) + else: + continue # We only fetch a new task if at least one of the available # machines is not a "service" machine (again, please refer to the diff --git a/cuckoo/machinery/kvm.py b/cuckoo/machinery/kvm.py index 1a32a6cee5..5e9517f3e1 100644 --- a/cuckoo/machinery/kvm.py +++ b/cuckoo/machinery/kvm.py @@ -39,21 +39,24 @@ def _disconnect(self, conn): """Disconnect, ignore request to disconnect.""" pass - def availables(self, platform=None, tags=None): - if all(param is None for param in [platform, tags]): + def availables(self, label=None, platform=None, tags=None): + if all(param is None for param in [label, platform, tags]): return super(KVM, self).availables() else: - return self._get_specific_availables(platform=platform, tags=tags) + return self._get_specific_availables(label=label, platform=platform, tags=tags) - def _get_specific_availables(self, platform=None, tags=None): + def _get_specific_availables(self, label=None, platform=None, tags=None): session = self.db.Session() try: machines = session.query(Machine) - if platform: - machines = machines.filter_by(platform=platform) + # Note that label > platform > tags + if label: + machines = machines.filter_by(locked=False).filter_by(label=label) + elif platform: + machines = machines.filter_by(locked=False).filter_by(platform=platform) elif tags: for tag in tags: - machines = machines.filter(Machine.tags.any(name=tag)) + machines = machines.filter_by(locked=False).filter(Machine.tags.any(name=tag)) return machines.count() except SQLAlchemyError as e: log.exception("Database error getting specific available machines: {0}".format(e)) From c2aaba2822e5e81f192a71ed6a583fdc9d86e7e4 Mon Sep 17 00:00:00 2001 From: Hardy-Cooper Date: Fri, 19 Mar 2021 11:40:05 -0400 Subject: [PATCH 4/5] Adding check for machine name during acquire() --- cuckoo/core/scheduler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuckoo/core/scheduler.py b/cuckoo/core/scheduler.py index fadd0e772a..228e49b756 100644 --- a/cuckoo/core/scheduler.py +++ b/cuckoo/core/scheduler.py @@ -162,7 +162,7 @@ def acquire_machine(self): # In some cases it's possible that we enter this loop without # having any available machines. We should make sure this is not # such case, or the analysis task will fail completely. - if not machinery.availables(platform=self.task.platform, tags=self.task.tags): + if not machinery.availables(label=self.task.machine, platform=self.task.platform, tags=self.task.tags): machine_lock.release() time.sleep(1) continue From 81e403de58c998cb97664474e837882a3024f779 Mon Sep 17 00:00:00 2001 From: Kevin Hardy-Cooper Date: Wed, 5 May 2021 12:07:15 -0400 Subject: [PATCH 5/5] Optimizing speed at which tasks are taken off queue and assigned to machines --- cuckoo/core/scheduler.py | 116 +++++++++++++++++---------------------- 1 file changed, 50 insertions(+), 66 deletions(-) diff --git a/cuckoo/core/scheduler.py b/cuckoo/core/scheduler.py index 228e49b756..3d14bb0adb 100644 --- a/cuckoo/core/scheduler.py +++ b/cuckoo/core/scheduler.py @@ -20,7 +20,7 @@ ) from cuckoo.common.objects import File from cuckoo.common.files import Folders -from cuckoo.core.database import Database, TASK_COMPLETED, TASK_REPORTED, TASK_PENDING, TASK_RUNNING +from cuckoo.core.database import Database, TASK_COMPLETED, TASK_REPORTED, TASK_RUNNING, TASK_PENDING from cuckoo.core.guest import GuestManager from cuckoo.core.plugins import RunAuxiliary, RunProcessing from cuckoo.core.plugins import RunSignatures, RunReporting @@ -965,11 +965,18 @@ def _cleanup_managers(self): cleaned.add(am) return cleaned + def _thr_periodic_log(self): + log.debug("# Tasks: %d; # Available Machines: %d; # Locked Machines: %d; # Total Machines: %d;", + self.db.count_tasks(status=TASK_PENDING), self.db.count_machines_available(), + len(self.db.list_machines(locked=True)), len(self.db.list_machines())) + threading.Timer(10, self._thr_periodic_log).start() + def start(self): """Start scheduler.""" self.initialize() log.info("Waiting for analysis tasks.") + self._thr_periodic_log() # Message queue with threads to transmit exceptions (used as IPC). errors = Queue.Queue() @@ -978,27 +985,12 @@ def start(self): if self.maxcount is None: self.maxcount = self.cfg.cuckoo.max_analysis_count + launched_analysis = True # This loop runs forever. while self.running: - time.sleep(1) - - # Run cleanup on finished analysis managers and untrack them - for am in self._cleanup_managers(): - self.analysis_managers.discard(am) - - # Wait until the machine lock is not locked. This is only the case - # when all machines are fully running, rather that about to start - # or still busy starting. This way we won't have race conditions - # with finding out there are no available machines in the analysis - # manager or having two analyses pick the same machine. - if not machine_lock.acquire(False): - logger( - "Could not acquire machine lock", - action="scheduler.machine_lock", status="busy" - ) - continue - - machine_lock.release() + if not launched_analysis: + time.sleep(1) + launched_analysis = False # If not enough free disk space is available, then we print an # error message and wait another round (this check is ignored @@ -1064,58 +1056,50 @@ def start(self): ) continue - # Fetch a pending analysis task. - # TODO This fixes only submissions by --machine, need to add - # other attributes (tags etc). - # TODO We should probably move the entire "acquire machine" logic - # from the Analysis Manager to the Scheduler and then pass the - # selected machine onto the Analysis Manager instance. - task, available = None, False - for machine in self.db.get_available_machines(): - task = self.db.fetch(machine=machine.name) - if task: - break - - # Since there are no tasks meant for a specific AVAILABLE machine, - # check if any tasks meant for a specific LOCKED machine, specific platform or specific tag - if not task: - tasks = self.db.list_tasks(status=TASK_PENDING) - if not tasks: + # Get all tasks in the queue + tasks = self.db.list_tasks(status=TASK_PENDING, details=True) + if not tasks: + continue + + for task in tasks: + # Run cleanup on finished analysis managers and untrack them + for am in self._cleanup_managers(): + self.analysis_managers.discard(am) + + # Wait until the machine lock is not locked. This is only the case + # when all machines are fully running, rather that about to start + # or still busy starting. This way we won't have race conditions + # with finding out there are no available machines in the analysis + # manager or having two analyses pick the same machine. + if not machine_lock.acquire(False): + logger( + "Could not acquire machine lock", + action="scheduler.machine_lock", status="busy" + ) continue - # If so, dispatch to an available machine - for task in tasks: - task_details = self.db.view_task(task.id) - # Note that label > platform > tags - if task_details.machine: - if machinery.availables(label=task_details.machine): - available = True - break - elif task_details.platform: - if machinery.availables(platform=task_details.platform): - available = True - break - elif task_details.tags: - tag_names = [tag.name for tag in task_details.tags] - if machinery.availables(tags=tag_names): - available = True - break - else: - available = True - break + machine_lock.release() - if task and available: - self.db.set_status(task.id, TASK_RUNNING) + available = False + # Note that label > platform > tags + if task.machine: + if machinery.availables(label=task.machine): + available = True + elif task.platform: + if machinery.availables(platform=task.platform): + available = True + elif task.tags: + tag_names = [tag.name for tag in task.tags] + if machinery.availables(tags=tag_names): + available = True else: + available = True + + if not available: continue - # We only fetch a new task if at least one of the available - # machines is not a "service" machine (again, please refer to the - # services auxiliary module for more information on service VMs). - if not task and available: - task = self.db.fetch(service=False) + self.db.set_status(task.id, TASK_RUNNING) - if task: log.debug("Processing task #%s", task.id) self.total_analysis_count += 1 @@ -1124,7 +1108,7 @@ def start(self): analysis.daemon = True analysis.start() self.analysis_managers.add(analysis) - + launched_analysis = True # Deal with errors. try: raise errors.get(block=False)