From 55366e751a844e4fb283a019d7be1a1383cc3a87 Mon Sep 17 00:00:00 2001
From: bmc-msft <41130664+bmc-msft@users.noreply.github.com>
Date: Fri, 23 Jul 2021 09:14:47 -0400
Subject: [PATCH] allow pools & scalesets set to `shutdown` to `halt` (#1104)

Currently, if a pool or scaleset is set to `shutdown`, it cannot be set to `halt`.

While moving from `halt` to `shutdown` would cause issues, moving from `shutdown` to `halt` is fine.
---
 .../__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 409b9f833..dd1f83a15 100644
--- a/src/api-service/__app__/onefuzzlib/workers/pools.py
+++ b/src/api-service/__app__/onefuzzlib/workers/pools.py
@@ -135,8 +135,7 @@ class Pool(BASE_POOL, ORMMixin):
     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 @@ class Pool(BASE_POOL, ORMMixin):
         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 @@ class Pool(BASE_POOL, ORMMixin):
         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 @@ class Pool(BASE_POOL, ORMMixin):
             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 @@ class Pool(BASE_POOL, ORMMixin):
 
         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 4aa20a575..2e1921bce 100644
--- a/src/api-service/__app__/onefuzzlib/workers/scalesets.py
+++ b/src/api-service/__app__/onefuzzlib/workers/scalesets.py
@@ -635,22 +635,11 @@ class Scaleset(BASE_SCALESET, ORMMixin):
                 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 @@ class Scaleset(BASE_SCALESET, ORMMixin):
         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: