expose the ability manually override node reset (#201)

src/api-service/__app__/agent_events/__init__.py

@@ -8,7 +8,13 @@ from typing import Optional, cast
 from uuid import UUID
 
 import azure.functions as func
-from onefuzztypes.enums import ErrorCode, NodeState, NodeTaskState, TaskState
+from onefuzztypes.enums import (
+    ErrorCode,
+    NodeState,
+    NodeTaskState,
+    TaskDebugFlag,
+    TaskState,
+)
 from onefuzztypes.models import (
     Error,
     NodeDoneEventData,
@@ -173,8 +179,21 @@ def on_worker_event(machine_id: UUID, event: WorkerEvent) -> None:
                     ],
                 )
             )
+            if task.config.debug and (
+                TaskDebugFlag.keep_node_on_failure in task.config.debug
+                or TaskDebugFlag.keep_node_on_completion in task.config.debug
+            ):
+                node.debug_keep_node = True
+                node.save()
+
         else:
             task.mark_stopping()
+            if (
+                task.config.debug
+                and TaskDebugFlag.keep_node_on_completion in task.config.debug
+            ):
+                node.debug_keep_node = True
+                node.save()
 
         node.to_reimage(done=True)
     else:

src/api-service/__app__/node/__init__.py

@@ -6,7 +6,7 @@
 import azure.functions as func
 from onefuzztypes.enums import ErrorCode
 from onefuzztypes.models import Error
-from onefuzztypes.requests import NodeGet, NodeSearch
+from onefuzztypes.requests import NodeGet, NodeSearch, NodeUpdate
 from onefuzztypes.responses import BoolResult
 
 from ..onefuzzlib.pools import Node, NodeTasks
@@ -42,6 +42,24 @@ def get(req: func.HttpRequest) -> func.HttpResponse:
     return ok(nodes)
 
 
+def post(req: func.HttpRequest) -> func.HttpResponse:
+    request = parse_request(NodeUpdate, req)
+    if isinstance(request, Error):
+        return not_ok(request, context="NodeUpdate")
+
+    node = Node.get_by_machine_id(request.machine_id)
+    if not node:
+        return not_ok(
+            Error(code=ErrorCode.UNABLE_TO_FIND, errors=["unable to find node"]),
+            context=request.machine_id,
+        )
+    if request.debug_keep_node is not None:
+        node.debug_keep_node = request.debug_keep_node
+
+    node.save()
+    return ok(BoolResult(result=True))
+
+
 def delete(req: func.HttpRequest) -> func.HttpResponse:
     request = parse_request(NodeGet, req)
     if isinstance(request, Error):
@@ -55,6 +73,9 @@ def delete(req: func.HttpRequest) -> func.HttpResponse:
         )
 
     node.set_halt()
+    if node.debug_keep_node:
+        node.debug_keep_node = False
+        node.save()
 
     return ok(BoolResult(result=True))
 
@@ -72,6 +93,9 @@ def patch(req: func.HttpRequest) -> func.HttpResponse:
         )
 
     node.stop()
+    if node.debug_keep_node:
+        node.debug_keep_node = False
+        node.save()
     return ok(BoolResult(result=True))
 
 

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

@@ -847,7 +847,16 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             logging.debug("scaleset delete will delete node: %s", self.scaleset_id)
             return
 
-        machine_ids = [x.machine_id for x in nodes]
+        machine_ids = []
+        for node in nodes:
+            if node.debug_keep_node:
+                logging.warning(
+                    "delete manually overridden %s:%s",
+                    self.scaleset_id,
+                    node.machine_id,
+                )
+            else:
+                machine_ids.append(node.machine_id)
 
         logging.info("deleting %s:%s", self.scaleset_id, machine_ids)
         delete_vmss_nodes(self.scaleset_id, machine_ids)
@@ -865,7 +874,16 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             logging.debug("scaleset delete will delete node: %s", self.scaleset_id)
             return
 
-        machine_ids = [x.machine_id for x in nodes]
+        machine_ids = []
+        for node in nodes:
+            if node.debug_keep_node:
+                logging.warning(
+                    "reimage manually overridden %s:%s",
+                    self.scaleset_id,
+                    node.machine_id,
+                )
+            else:
+                machine_ids.append(node.machine_id)
 
         result = reimage_vmss_nodes(self.scaleset_id, machine_ids)
         if isinstance(result, Error):