allow pools & scalesets set to shutdown to halt

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")

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: