From 8d20cc7d7e498dda11d6b18845eb70b8ce975629 Mon Sep 17 00:00:00 2001
From: Brian Caswell <bmc@shmoo.com>
Date: Sun, 21 Feb 2021 12:23:54 -0500
Subject: [PATCH] don't schedule work to a node if the scaleset or pool is
 shutting down

---
 .../__app__/onefuzzlib/workers/nodes.py       | 41 ++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/api-service/__app__/onefuzzlib/workers/nodes.py b/src/api-service/__app__/onefuzzlib/workers/nodes.py
index 411c009c4e..ce4695500e 100644
--- a/src/api-service/__app__/onefuzzlib/workers/nodes.py
+++ b/src/api-service/__app__/onefuzzlib/workers/nodes.py
@@ -8,7 +8,7 @@
 from typing import List, Optional, Tuple
 from uuid import UUID
 
-from onefuzztypes.enums import ErrorCode, NodeState, TaskState
+from onefuzztypes.enums import ErrorCode, NodeState, PoolState, ScalesetState, TaskState
 from onefuzztypes.events import (
     EventNodeCreated,
     EventNodeDeleted,
@@ -241,6 +241,9 @@ def could_shrink_scaleset(self) -> bool:
         return False
 
     def can_process_new_work(self) -> bool:
+        from .pools import Pool
+        from .scalesets import Scaleset
+
         if self.is_outdated():
             logging.info(
                 "can_schedule agent and service versions differ, stopping node. "
@@ -279,6 +282,42 @@ def can_process_new_work(self) -> bool:
             logging.info("node scheduled to shrink.  machine_id:%s", self.machine_id)
             return False
 
+        if self.scaleset_id:
+            scaleset = Scaleset.get_by_id(self.scaleset_id)
+            if isinstance(scaleset, Error):
+                logging.info(
+                    "can_schedule - invalid scaleset.  scaleset_id:%s machine_id:%s",
+                    self.scaleset_id,
+                    self.machine_id,
+                )
+                return False
+
+            if scaleset.state not in ScalesetState.available():
+                logging.info(
+                    "can_schedule - scaleset not available for work. "
+                    "scaleset_id:%s machine_id:%s",
+                    self.scaleset_id,
+                    self.machine_id,
+                )
+                return False
+
+        pool = Pool.get_by_name(self.pool_name)
+        if isinstance(pool, Error):
+            logging.info(
+                "can_schedule - invalid pool. " "pool_name:%s machine_id:%s",
+                self.pool_name,
+                self.machine_id,
+            )
+            return False
+        if pool.state not in PoolState.available():
+            logging.info(
+                "can_schedule - pool is not available for work. "
+                "pool_name:%s machine_id:%s",
+                self.pool_name,
+                self.machine_id,
+            )
+            return False
+
         return True
 
     def is_outdated(self) -> bool: