Add StopIfFree node command to tell free nodes to stop asking for new work (#866)

src/api-service/__app__/onefuzzlib/versions.py

@@ -6,6 +6,7 @@
 import os
 from typing import Dict
 
+import semver
 from memoization import cached
 from onefuzztypes.responses import Version
 
@@ -29,3 +30,8 @@ def versions() -> Dict[str, Version]:
         version=__version__,
     )
     return {"onefuzz": entry}
+
+
+def is_minimum_version(*, version: str, minimum: str) -> bool:
+    # check if version is at least (or higher) than minimum
+    return bool(semver.VersionInfo.parse(version).compare(minimum) >= 0)

src/api-service/__app__/onefuzzlib/workers/nodes.py

@@ -16,7 +16,12 @@ from onefuzztypes.events import (
 )
 from onefuzztypes.models import Error
 from onefuzztypes.models import Node as BASE_NODE
-from onefuzztypes.models import NodeAssignment, NodeCommand, NodeCommandAddSshKey
+from onefuzztypes.models import (
+    NodeAssignment,
+    NodeCommand,
+    NodeCommandAddSshKey,
+    NodeCommandStopIfFree,
+)
 from onefuzztypes.models import NodeTasks as BASE_NODE_TASK
 from onefuzztypes.models import Result, StopNodeCommand, StopTaskNodeCommand
 from onefuzztypes.primitives import PoolName
@@ -26,6 +31,7 @@ from ..__version__ import __version__
 from ..azure.vmss import get_instance_id
 from ..events import send_event
 from ..orm import MappingIntStrAny, ORMMixin, QueryFilter
+from ..versions import is_minimum_version
 
 NODE_EXPIRATION_TIME: datetime.timedelta = datetime.timedelta(hours=1)
 NODE_REIMAGE_TIME: datetime.timedelta = datetime.timedelta(days=7)
@@ -338,6 +344,11 @@ class Node(BASE_NODE, ORMMixin):
         if not self.reimage_requested and not self.delete_requested:
             logging.info("setting reimage_requested: %s", self.machine_id)
             self.reimage_requested = True
+
+        # if we're going to reimage, make sure the node doesn't pick up new work
+        # too.
+        self.send_stop_if_free()
+
         self.save()
 
     def add_ssh_public_key(self, public_key: str) -> Result[None]:
@@ -355,6 +366,10 @@ class Node(BASE_NODE, ORMMixin):
         )
         return None
 
+    def send_stop_if_free(self) -> None:
+        if is_minimum_version(version=self.version, minimum="2.16.1"):
+            self.send_message(NodeCommand(stop_if_free=NodeCommandStopIfFree()))
+
     def stop(self, done: bool = False) -> None:
         self.to_reimage(done=done)
         self.send_message(NodeCommand(stop=StopNodeCommand()))

src/api-service/__app__/onefuzzlib/workers/scalesets.py

@@ -440,10 +440,19 @@ class Scaleset(BASE_SCALESET, ORMMixin):
         return
 
     def _resize_shrink(self, to_remove: int) -> None:
+        logging.info(
+            SCALESET_LOG_PREFIX + "shrinking scaleset. scaleset_id:%s to_remove:%d",
+            self.scaleset_id,
+            to_remove,
+        )
         queue = ScalesetShrinkQueue(self.scaleset_id)
         for _ in range(to_remove):
             queue.add_entry()
 
+        nodes = Node.search_states(scaleset_id=self.scaleset_id)
+        for node in nodes:
+            node.send_stop_if_free()
+
     def resize(self) -> None:
         # no longer needing to resize
         if self.state != ScalesetState.resize: