From 55558922d8c26185e6cea9759d0b2c931dd16d9c Mon Sep 17 00:00:00 2001
From: Brian Caswell <bmc@shmoo.com>
Date: Thu, 22 Jul 2021 15:09:55 -0400
Subject: [PATCH] allow pools & scalesets set to `shutdown` to `halt`

---
 .../__app__/onefuzzlib/workers/pools.py       | 26 ++++++++++---------
 .../__app__/onefuzzlib/workers/scalesets.py   | 15 +++--------
 2 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/src/api-service/__app__/onefuzzlib/workers/pools.py b/src/api-service/__app__/onefuzzlib/workers/pools.py
index 409b9f8331..dd1f83a157 100644
--- a/src/api-service/__app__/onefuzzlib/workers/pools.py
+++ b/src/api-service/__app__/onefuzzlib/workers/pools.py
@@ -135,8 +135,7 @@ def get_pool_queue(self) -> str:
     def init(self) -> None:
         create_queue(self.get_pool_queue(), StorageType.corpus)
         ShrinkQueue(self.pool_id).create()
-        self.state = PoolState.running
-        self.save()
+        self.set_state(PoolState.running)
 
     def schedule_workset(self, work_set: WorkSet) -> bool:
         # Don't schedule work for pools that can't and won't do work.
@@ -183,15 +182,10 @@ def search_states(cls, *, states: Optional[List[PoolState]] = None) -> List["Poo
         return cls.search(query=query)
 
     def set_shutdown(self, now: bool) -> None:
-        if self.state in [PoolState.halt, PoolState.shutdown]:
-            return
-
         if now:
-            self.state = PoolState.halt
+            self.set_state(PoolState.halt)
         else:
-            self.state = PoolState.shutdown
-
-        self.save()
+            self.set_state(PoolState.shutdown)
 
     def shutdown(self) -> None:
         """shutdown allows nodes to finish current work then delete"""
@@ -202,8 +196,6 @@ def shutdown(self) -> None:
         nodes = Node.search(query={"pool_name": [self.name]})
         if not scalesets and not nodes:
             logging.info("pool stopped, deleting: %s", self.name)
-
-            self.state = PoolState.halt
             self.delete()
             return
 
@@ -227,7 +219,6 @@ def halt(self) -> None:
             delete_queue(self.get_pool_queue(), StorageType.corpus)
             ShrinkQueue(self.pool_id).delete()
             logging.info("pool stopped, deleting: %s", self.name)
-            self.state = PoolState.halt
             self.delete()
             return
 
@@ -239,6 +230,17 @@ def halt(self) -> None:
 
         self.save()
 
+    def set_state(self, state: PoolState) -> None:
+        if self.state == state:
+            return
+
+        # scalesets should never leave the `halt` state
+        if self.state == PoolState.halt:
+            return
+
+        self.state = state
+        self.save()
+
     @classmethod
     def key_fields(cls) -> Tuple[str, str]:
         return ("name", "pool_id")
diff --git a/src/api-service/__app__/onefuzzlib/workers/scalesets.py b/src/api-service/__app__/onefuzzlib/workers/scalesets.py
index 4aa20a5751..2e1921bcef 100644
--- a/src/api-service/__app__/onefuzzlib/workers/scalesets.py
+++ b/src/api-service/__app__/onefuzzlib/workers/scalesets.py
@@ -635,22 +635,11 @@ def reimage_nodes(self, nodes: List[Node]) -> None:
                 node.delete()
 
     def set_shutdown(self, now: bool) -> None:
-        if self.state in [ScalesetState.halt, ScalesetState.shutdown]:
-            return
-
-        logging.info(
-            SCALESET_LOG_PREFIX + "scaleset set_shutdown: scaleset_id:%s now:%s",
-            self.scaleset_id,
-            now,
-        )
-
         if now:
             self.set_state(ScalesetState.halt)
         else:
             self.set_state(ScalesetState.shutdown)
 
-        self.save()
-
     def shutdown(self) -> None:
         size = get_vmss_size(self.scaleset_id)
         if size is None:
@@ -807,6 +796,10 @@ def set_state(self, state: ScalesetState) -> None:
         if self.state == state:
             return
 
+        # scalesets should never leave the `halt` state
+        if self.state == ScalesetState.halt:
+            return
+
         self.state = state
         self.save()
         if self.state == ScalesetState.resize: