From c0284b0fa1170cb72f8a75d509617af03aca9e2b Mon Sep 17 00:00:00 2001
From: Mike McLean <mikem@redhat.com>
Date: Jan 19 2024 06:40:57 +0000
Subject: [PATCH 1/2] assign overrides


---

diff --git a/kojihub/kojihub.py b/kojihub/kojihub.py
index 54d8b68..c31d5e6 100644
--- a/kojihub/kojihub.py
+++ b/kojihub/kojihub.py
@@ -11972,15 +11972,7 @@ class RootExports(object):
                 raise koji.ActionNotAllowed('Cannot cancel build, not owner')
         return cancel_build(build['id'])
 
-    def assignTask(self, task_id, host, force=False):
-        """Assign a task to a host
-
-        Specify force=True to assign a non-free task
-        """
-        context.session.assertPerm('admin')
-        task = Task(task_id)
-        host = get_host(host, strict=True)
-        return task.assign(host['id'], force)
+    assignTask = staticmethod(scheduler.do_assign)
 
     def freeTask(self, task_id):
         """Free a task"""
diff --git a/kojihub/scheduler.py b/kojihub/scheduler.py
index 7718da6..ec8166f 100644
--- a/kojihub/scheduler.py
+++ b/kojihub/scheduler.py
@@ -325,7 +325,7 @@ class TaskScheduler(object):
                 if (host['capacity'] - host['_load'] > min_avail and
                         host['_ntasks'] < self.maxjobs):
                     # add run entry
-                    self.add_run(task, host)
+                    self.assign(task, host)
                     # update our totals and rank
                     host['_load'] += task['weight']
                     host['_ntasks'] += 1
@@ -357,9 +357,11 @@ class TaskScheduler(object):
 
             taskruns = runs.get(task['task_id'], [])
             if not taskruns:
-                log_both('Assigned task with no active run entry', task_id=task['task_id'],
-                         host_id=host['id'], level=logging.ERROR)
-                kojihub.Task(task['task_id']).free()
+                # a task that is assigned without a run entry is treated as an override
+                # we simply leave these alone
+                # TODO track overrides more explicitly
+                logger.debug('Override task assignment: task %i, host %s',
+                             task['task_id'], host['name'])
                 continue
 
             if len(taskruns) > 1:
@@ -565,8 +567,17 @@ class TaskScheduler(object):
 
         return hosts
 
-    def add_run(self, task, host):
-        log_both('Assigning task', task_id=task['task_id'], host_id=host['id'])
+    def assign(self, task, host, force=False, override=False):
+        # mark the task assigned
+        success = kojihub.Task(task['task_id']).assign(host['id'], force=force)
+        if not success:
+            log_both('Assignment failed', task_id=task['task_id'], host_id=host['id'])
+            return False
+
+        if override:
+            log_both('Assigning task (override)', task_id=task['task_id'], host_id=host['id'])
+        else:
+            log_both('Assigning task', task_id=task['task_id'], host_id=host['id'])
 
         # mark any older runs inactive
         update = UpdateProcessor(
@@ -577,14 +588,31 @@ class TaskScheduler(object):
         )
         update.execute()
 
-        # add the new run
-        insert = InsertProcessor('scheduler_task_runs')
-        insert.set(task_id=task['task_id'], host_id=host['id'])
-        insert.execute()
+        if not override:
+            # add the new run
+            insert = InsertProcessor('scheduler_task_runs')
+            insert.set(task_id=task['task_id'], host_id=host['id'])
+            insert.execute()
+        # In the override case, we omit the run entry
 
-        # mark the task assigned
-        task = kojihub.Task(task['task_id'])
-        task.assign(host['id'])
+        return True
+
+
+# exported as assignTask in kojihub
+def do_assign(task_id, host, force=False, override=False):
+    """Assign a task to a host
+
+    Specify force=True to assign a non-free task
+    Specify override=True to prevent the scheduler from reassigning later
+    """
+    task_id = kojihub.convert_value(task_id, cast=int)
+    host = kojihub.get_host(host, strict=True)
+    force = kojihub.convert_value(force, cast=bool)
+    override = kojihub.convert_value(override, cast=bool)
+    context.session.assertPerm('admin')
+    task = {'task_id': task_id}  # the assign call just needs the task_id field
+    db_lock('scheduler', wait=True)
+    return TaskScheduler().assign(task, host, force=force, override=override)
 
 
 class SchedulerExports:

From feb5d0e8ab121644eec3ad6732cbf5440189e3e4 Mon Sep 17 00:00:00 2001
From: Mike McLean <mikem@redhat.com>
Date: Jan 19 2024 06:41:02 +0000
Subject: [PATCH 2/2] fix ordering in get_tasks_for_host


---

diff --git a/kojihub/scheduler.py b/kojihub/scheduler.py
index ec8166f..815b0f1 100644
--- a/kojihub/scheduler.py
+++ b/kojihub/scheduler.py
@@ -76,6 +76,7 @@ def get_tasks_for_host(hostID, retry=True):
         columns=fields, aliases=aliases, tables=['task'],
         clauses=['host_id = %(hostID)s', 'state=%(assigned)s'],
         values={'hostID': hostID, 'assigned': koji.TASK_STATES['ASSIGNED']},
+        opts={'order': 'priority,create_ts'},
     )
 
     tasks = query.execute()