add EventScalesetResizeScheduled (#1047)

docs/webhook_events.md

@@ -43,6 +43,7 @@ Each event will be submitted via HTTP POST to the user provided URL.
 * [scaleset_created](#scaleset_created)
 * [scaleset_deleted](#scaleset_deleted)
 * [scaleset_failed](#scaleset_failed)
+* [scaleset_resize_scheduled](#scaleset_resize_scheduled)
 * [scaleset_state_updated](#scaleset_state_updated)
 * [task_created](#task_created)
 * [task_failed](#task_failed)
@@ -2365,6 +2366,47 @@ Each event will be submitted via HTTP POST to the user provided URL.
 }
 ```
 
+### scaleset_resize_scheduled
+
+#### Example
+
+```json
+{
+    "pool_name": "example",
+    "scaleset_id": "00000000-0000-0000-0000-000000000000",
+    "size": 0
+}
+```
+
+#### Schema
+
+```json
+{
+    "properties": {
+        "pool_name": {
+            "title": "Pool Name",
+            "type": "string"
+        },
+        "scaleset_id": {
+            "format": "uuid",
+            "title": "Scaleset Id",
+            "type": "string"
+        },
+        "size": {
+            "title": "Size",
+            "type": "integer"
+        }
+    },
+    "required": [
+        "scaleset_id",
+        "pool_name",
+        "size"
+    ],
+    "title": "EventScalesetResizeScheduled",
+    "type": "object"
+}
+```
+
 ### scaleset_state_updated
 
 #### Example
@@ -5272,6 +5314,30 @@ Each event will be submitted via HTTP POST to the user provided URL.
             "title": "EventScalesetFailed",
             "type": "object"
         },
+        "EventScalesetResizeScheduled": {
+            "properties": {
+                "pool_name": {
+                    "title": "Pool Name",
+                    "type": "string"
+                },
+                "scaleset_id": {
+                    "format": "uuid",
+                    "title": "Scaleset Id",
+                    "type": "string"
+                },
+                "size": {
+                    "title": "Size",
+                    "type": "integer"
+                }
+            },
+            "required": [
+                "scaleset_id",
+                "pool_name",
+                "size"
+            ],
+            "title": "EventScalesetResizeScheduled",
+            "type": "object"
+        },
         "EventScalesetStateUpdated": {
             "properties": {
                 "pool_name": {
@@ -5456,6 +5522,7 @@ Each event will be submitted via HTTP POST to the user provided URL.
                 "scaleset_deleted",
                 "scaleset_failed",
                 "scaleset_state_updated",
+                "scaleset_resize_scheduled",
                 "task_created",
                 "task_failed",
                 "task_state_updated",
@@ -6167,6 +6234,9 @@ Each event will be submitted via HTTP POST to the user provided URL.
                 {
                     "$ref": "#/definitions/EventScalesetStateUpdated"
                 },
+                {
+                    "$ref": "#/definitions/EventScalesetResizeScheduled"
+                },
                 {
                     "$ref": "#/definitions/EventTaskFailed"
                 },

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

@@ -35,12 +35,11 @@ def scale_up(pool: Pool, scalesets: List[Scaleset], nodes_needed: int) -> None:
             if scaleset.size < max_size:
                 current_size = scaleset.size
                 if nodes_needed <= max_size - current_size:
-                    scaleset.size = current_size + nodes_needed
+                    scaleset.set_size(current_size + nodes_needed)
                     nodes_needed = 0
                 else:
-                    scaleset.size = max_size
+                    scaleset.set_size(max_size)
                     nodes_needed = nodes_needed - (max_size - current_size)
-                scaleset.set_state(ScalesetState.resize)
 
             else:
                 continue
@@ -112,7 +111,7 @@ def scale_down(scalesets: List[Scaleset], nodes_to_remove: int) -> None:
                 continue
 
             # Resize of VMSS needed
-            scaleset.size = scaleset.size - max_nodes_remove
+            scaleset.set_size(scaleset.size - max_nodes_remove)
             nodes_to_remove = nodes_to_remove - max_nodes_remove
             scaleset.set_state(ScalesetState.resize)
 

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

@@ -13,6 +13,7 @@ from onefuzztypes.events import (
     EventScalesetCreated,
     EventScalesetDeleted,
     EventScalesetFailed,
+    EventScalesetResizeScheduled,
     EventScalesetStateUpdated,
 )
 from onefuzztypes.models import Error
@@ -472,6 +473,14 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             )
             self.set_state(ScalesetState.resize)
 
+    def set_size(self, size: int) -> None:
+        # ensure we always stay within max_size boundaries
+        size = min(size, self.max_size())
+        if self.size != size:
+            self.size = size
+            self.set_state(ScalesetState.resize)
+            self.save()
+
     def resize(self) -> None:
         # no longer needing to resize
         if self.state != ScalesetState.resize:
@@ -780,8 +789,19 @@ class Scaleset(BASE_SCALESET, ORMMixin):
 
         self.state = state
         self.save()
-        send_event(
+        if self.state == ScalesetState.resize:
-            EventScalesetStateUpdated(
+            send_event(
-                scaleset_id=self.scaleset_id, pool_name=self.pool_name, state=self.state
+                EventScalesetResizeScheduled(
+                    scaleset_id=self.scaleset_id,
+                    pool_name=self.pool_name,
+                    size=self.size,
+                )
+            )
+        else:
+            send_event(
+                EventScalesetStateUpdated(
+                    scaleset_id=self.scaleset_id,
+                    pool_name=self.pool_name,
+                    state=self.state,
+                )
             )
-        )

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

@@ -134,23 +134,21 @@ def patch(req: func.HttpRequest) -> func.HttpResponse:
     if isinstance(scaleset, Error):
         return not_ok(scaleset, context="ScalesetUpdate")
 
-    if scaleset.state not in [ScalesetState.running, ScalesetState.resize]:
+    if scaleset.state not in ScalesetState.can_update():
         return not_ok(
             Error(
                 code=ErrorCode.INVALID_REQUEST,
                 errors=[
-                    "scaleset state must be 'running' or "
+                    "scaleset must be in the following states to update: "
-                    "'resize' state to modify scaleset"
+                    + ",".join(x.name for x in ScalesetState.can_update())
                 ],
             ),
             context="ScalesetUpdate",
         )
 
     if request.size is not None:
-        scaleset.size = request.size
+        scaleset.set_size(request.size)
-        scaleset.set_state(ScalesetState.resize)
 
-    scaleset.save()
     scaleset.auth = None
     return ok(scaleset)
 

src/pytypes/extra/generate-docs.py

@@ -40,6 +40,7 @@ from onefuzztypes.events import (
     EventScalesetCreated,
     EventScalesetDeleted,
     EventScalesetFailed,
+    EventScalesetResizeScheduled,
     EventScalesetStateUpdated,
     EventTaskCreated,
     EventTaskFailed,
@@ -200,6 +201,9 @@ def main() -> None:
             pool_name=PoolName("example"),
             state=ScalesetState.init,
         ),
+        EventScalesetResizeScheduled(
+            scaleset_id=UUID(int=0), pool_name=PoolName("example"), size=0
+        ),
         EventJobCreated(
             job_id=UUID(int=0),
             config=JobConfig(

src/pytypes/onefuzztypes/enums.py

@@ -304,6 +304,13 @@ class ScalesetState(Enum):
     halt = "halt"
     creation_failed = "creation_failed"
 
+    @classmethod
+    def can_update(cls) -> List["ScalesetState"]:
+        """
+        set of states that indicate the scaleset can be updated
+        """
+        return [cls.running, cls.resize]
+
     @classmethod
     def needs_work(cls) -> List["ScalesetState"]:
         """

src/pytypes/onefuzztypes/events.py

@@ -117,6 +117,12 @@ class EventScalesetDeleted(BaseEvent):
     pool_name: PoolName
 
 
+class EventScalesetResizeScheduled(BaseEvent):
+    scaleset_id: UUID
+    pool_name: PoolName
+    size: int
+
+
 class EventPoolDeleted(BaseEvent):
     pool_name: PoolName
 
@@ -223,6 +229,7 @@ Event = Union[
     EventScalesetCreated,
     EventScalesetDeleted,
     EventScalesetStateUpdated,
+    EventScalesetResizeScheduled,
     EventTaskFailed,
     EventTaskStateUpdated,
     EventTaskCreated,
@@ -252,6 +259,7 @@ class EventType(Enum):
     scaleset_deleted = "scaleset_deleted"
     scaleset_failed = "scaleset_failed"
     scaleset_state_updated = "scaleset_state_updated"
+    scaleset_resize_scheduled = "scaleset_resize_scheduled"
     task_created = "task_created"
     task_failed = "task_failed"
     task_state_updated = "task_state_updated"
@@ -282,6 +290,7 @@ EventTypeMap = {
     EventType.scaleset_deleted: EventScalesetDeleted,
     EventType.scaleset_failed: EventScalesetFailed,
     EventType.scaleset_state_updated: EventScalesetStateUpdated,
+    EventType.scaleset_resize_scheduled: EventScalesetResizeScheduled,
     EventType.task_created: EventTaskCreated,
     EventType.task_failed: EventTaskFailed,
     EventType.task_state_updated: EventTaskStateUpdated,