Convert branch state management to jobs

Author: mxsrc

src/api/organization/project/branch/__init__.py

@@ -45,7 +45,6 @@
     deployment_status,
 )
 from .....deployment.kubernetes._util import core_v1_client
-from .....deployment.kubernetes.neonvm import PowerState as NeonVMPowerState
 from .....deployment.kubernetes.neonvm import set_virtualmachine_power_state
 from .....deployment.kubernetes.volume_clone import (
     clone_branch_database_volume,
@@ -95,6 +94,12 @@
 from ....settings import get_settings as get_api_settings
 from .api_keys import api as api_key_api
 from .auth import api as auth_api
+from .control_tasks import (
+    _CONTROL_TO_POWER_STATE,
+    _CONTROL_TRANSITION_INITIAL,
+    dispatch_control,
+    perform_control,
+)
 from .resize_tasks import dispatch_resize
 from .tasks import task_api
 
@@ -284,6 +289,12 @@ async def refresh_branch_status(branch_id: Identifier) -> BranchServiceStatus:
 
 
 async def _refresh_branch_status(branch: Branch) -> BranchServiceStatus:
+    if branch.control_task_id is not None:
+        result = perform_control.AsyncResult(str(branch.control_task_id))
+        action = (result.kwargs or {}).get("action")
+        if action in _CONTROL_TRANSITION_INITIAL:
+            return _CONTROL_TRANSITION_INITIAL[action]
+
     current_status = _parse_branch_status(branch.status)
     status = deployment_status(branch.id)
 
@@ -1908,109 +1919,58 @@ async def resize(
     404: NotFound,
 }
 
-_CONTROL_TO_AUTOSCALER_POWERSTATE: dict[str, NeonVMPowerState] = {
-    "pause": "Stopped",
-    "resume": "Running",
-    "start": "Running",
-    "stop": "Stopped",
-}
-
-_CONTROL_TRANSITION_INITIAL: dict[str, BranchServiceStatus] = {
-    "pause": BranchServiceStatus.PAUSING,
-    "resume": BranchServiceStatus.RESUMING,
-    "start": BranchServiceStatus.STARTING,
-    "stop": BranchServiceStatus.STOPPING,
-}
-
-_CONTROL_TRANSITION_FINAL: dict[str, BranchServiceStatus | None] = {
-    "pause": BranchServiceStatus.PAUSED,
-    "resume": BranchServiceStatus.STARTING,
-    "start": None,
-    "stop": BranchServiceStatus.STOPPED,
-}
-
 
 async def _set_branch_status(session: SessionDep, branch: Branch, status: BranchServiceStatus):
     branch.set_status(status)
     await session.commit()
 
 
-async def _set_final_branch_status(session: SessionDep, branch: Branch, action: str) -> None:
-    final_status = _CONTROL_TRANSITION_FINAL[action]
-    if final_status is None:
-        return
-    await _set_branch_status(session, branch, final_status)
-
-
-async def _set_autoscaler_power_state(action: str, namespace: str, name: str) -> None:
-    power_state = _CONTROL_TO_AUTOSCALER_POWERSTATE.get(action)
-    if power_state is None:
-        return
-    await set_virtualmachine_power_state(namespace, name, power_state)
-
-
-async def _apply_branch_action(
-    *,
-    action: str,
-    autoscaler_namespace: str,
-    autoscaler_vm_name: str,
-) -> None:
-    await _set_autoscaler_power_state(action, autoscaler_namespace, autoscaler_vm_name)
-
-
 @instance_api.post(
     "/pause",
     name="organizations:projects:branch:pause",
-    status_code=204,
+    status_code=202,
     responses=_control_responses,
 )
 @instance_api.post(
     "/resume",
     name="organizations:projects:branch:resume",
-    status_code=204,
+    status_code=202,
     responses=_control_responses,
 )
 @instance_api.post(
     "/start",
     name="organizations:projects:branch:start",
-    status_code=204,
+    status_code=202,
     responses=_control_responses,
 )
 @instance_api.post(
     "/stop",
     name="organizations:projects:branch:stop",
-    status_code=204,
+    status_code=202,
     responses=_control_responses,
 )
 async def control_branch(
     session: SessionDep,
     request: Request,
-    _organization: OrganizationDep,
-    _project: ProjectDep,
+    organization: OrganizationDep,
+    project: ProjectDep,
     branch: BranchDep,
 ):
     action = request.scope["route"].name.split(":")[-1]
-    assert action in _CONTROL_TO_AUTOSCALER_POWERSTATE
-    branch_in_session = await session.merge(branch)
-    branch_id = branch_in_session.id
-    autoscaler_namespace, autoscaler_vm_name = get_autoscaler_vm_identity(branch_id)
-    await _set_branch_status(session, branch_in_session, _CONTROL_TRANSITION_INITIAL[action])
-    try:
-        await _apply_branch_action(
-            action=action,
-            autoscaler_namespace=autoscaler_namespace,
-            autoscaler_vm_name=autoscaler_vm_name,
-        )
-    except ApiException as e:
-        await _set_branch_status(session, branch_in_session, BranchServiceStatus.ERROR)
-        status = 404 if e.status == 404 else 400
-        raise HTTPException(status_code=status, detail=e.body or str(e)) from e
-    except VelaKubernetesError as e:
-        await _set_branch_status(session, branch_in_session, BranchServiceStatus.ERROR)
-        raise HTTPException(status_code=500, detail=str(e)) from e
-    else:
-        await _set_final_branch_status(session, branch_in_session, action)
-    return Response(status_code=204)
+    assert action in _CONTROL_TO_POWER_STATE
+    task_id = dispatch_control(str(branch.id), action)
+    branch.control_task_id = task_id
+    await session.commit()
+
+    task_url = url_path_for(
+        request,
+        "organizations:projects:branch:tasks:detail",
+        organization_id=await organization.awaitable_attrs.id,
+        project_id=await project.awaitable_attrs.id,
+        branch_id=await branch.awaitable_attrs.id,
+        task_id=task_id,
+    )
+    return Response(status_code=202, headers={"Location": task_url})
 
 
 instance_api.include_router(auth_api, prefix="/auth")

src/api/organization/project/branch/control_tasks.py

@@ -0,0 +1,109 @@
+"""Celery task for branch start/stop/pause/resume lifecycle management.
+
+A single task patches the Kubernetes VM power state, polls until the VM
+reaches the desired phase, then updates the branch status in the database.
+"""
+
+import asyncio
+import logging
+from uuid import UUID
+
+from asgiref.sync import async_to_sync
+from ulid import ULID
+
+from .....database import AsyncSessionLocal
+from .....deployment import get_autoscaler_vm_identity
+from .....deployment.kubernetes.neonvm import Phase, PowerState, get_neon_vm, set_virtualmachine_power_state
+from .....models.branch import Branch, BranchServiceStatus
+from .....worker import app
+
+logger = logging.getLogger(__name__)
+
+_CONTROL_TO_POWER_STATE: dict[str, PowerState] = {
+    "pause": "Stopped",
+    "resume": "Running",
+    "start": "Running",
+    "stop": "Stopped",
+}
+
+_CONTROL_TRANSITION_INITIAL: dict[str, BranchServiceStatus] = {
+    "pause": BranchServiceStatus.PAUSING,
+    "resume": BranchServiceStatus.RESUMING,
+    "start": BranchServiceStatus.STARTING,
+    "stop": BranchServiceStatus.STOPPING,
+}
+
+_CONTROL_TRANSITION_FINAL: dict[str, BranchServiceStatus | None] = {
+    "pause": BranchServiceStatus.PAUSED,
+    "resume": BranchServiceStatus.STARTING,
+    "start": None,
+    "stop": BranchServiceStatus.STOPPED,
+}
+
+_DESIRED_PHASES: dict[str, set[Phase]] = {
+    "start": {Phase.running},
+    "resume": {Phase.running},
+    "stop": {Phase.stopped, Phase.succeeded},
+    "pause": {Phase.stopped, Phase.succeeded},
+}
+
+_POLL_INTERVAL_SEC = 5
+_TIMEOUT_SEC = 600  # 10 minutes
+
+
+async def _async_perform_control(branch_id: str, action: str) -> dict:
+    ulid = ULID.from_str(branch_id)
+    namespace, name = get_autoscaler_vm_identity(ulid)
+
+    power_state = _CONTROL_TO_POWER_STATE[action]
+    await set_virtualmachine_power_state(namespace, name, power_state)
+
+    desired_phases = _DESIRED_PHASES[action]
+    elapsed = 0.0
+    while True:
+        vm = await get_neon_vm(namespace, name)
+        if vm.status is not None and vm.status.phase in desired_phases:
+            break
+        if elapsed >= _TIMEOUT_SEC:
+            raise TimeoutError(f"VM {name!r} did not reach {desired_phases} within {_TIMEOUT_SEC}s")
+        await asyncio.sleep(_POLL_INTERVAL_SEC)
+        elapsed += _POLL_INTERVAL_SEC
+
+    final_status = _CONTROL_TRANSITION_FINAL[action]
+    async with AsyncSessionLocal() as session:
+        branch = await session.get(Branch, ulid)
+        if branch is None:
+            logger.error("Branch %s not found after control action %s", branch_id, action)
+            return {"action": action}
+        if final_status is not None:
+            branch.set_status(final_status)
+        branch.control_task_id = None
+        await session.commit()
+
+    return {"action": action}
+
+
+async def _async_perform_control_with_error_handling(branch_id: str, action: str) -> dict:
+    try:
+        return await _async_perform_control(branch_id, action)
+    except Exception:
+        ulid = ULID.from_str(branch_id)
+        async with AsyncSessionLocal() as session:
+            branch = await session.get(Branch, ulid)
+            if branch is not None:
+                branch.set_status(BranchServiceStatus.ERROR)
+                branch.control_task_id = None
+                await session.commit()
+        raise
+
+
+@app.task(name="simplyblock.vela.branch.control")
+def perform_control(branch_id: str, action: str) -> dict:
+    """Patch VM power state, wait for desired phase, update branch status."""
+    return async_to_sync(_async_perform_control_with_error_handling)(branch_id, action)
+
+
+def dispatch_control(branch_id: str, action: str) -> UUID:
+    """Dispatch perform_control asynchronously; return the Celery task UUID."""
+    result = perform_control.apply_async(kwargs={"branch_id": branch_id, "action": action})
+    return UUID(result.id)

src/api/organization/project/branch/tasks.py

@@ -1,6 +1,6 @@
 """Branch task list/detail endpoints.
 
-Exposes Celery task state (currently resize only) under:
+Exposes Celery task state (resize and control) under:
   GET .../branches/{branch_id}/tasks
   GET .../branches/{branch_id}/tasks/{task_id}
 """
@@ -14,6 +14,7 @@
 
 from ...._util import Forbidden, NotFound, Unauthenticated
 from ....dependencies import BranchDep, OrganizationDep, ProjectDep
+from .control_tasks import perform_control
 from .resize_tasks import finalize_resize
 
 task_api = APIRouter(tags=["branch"])
@@ -37,7 +38,7 @@ class BranchTaskPublic(BaseModel):
     date_done: datetime | None
 
 
-def _build_task_public(task_id: UUID) -> BranchTaskPublic:
+def _build_resize_task_public(task_id: UUID) -> BranchTaskPublic:
     result = finalize_resize.AsyncResult(str(task_id))
     state = result.state
     status = _CELERY_STATE_TO_STATUS.get(state, state)
@@ -53,6 +54,23 @@ def _build_task_public(task_id: UUID) -> BranchTaskPublic:
     )
 
 
+def _build_control_task_public(task_id: UUID) -> BranchTaskPublic:
+    result = perform_control.AsyncResult(str(task_id))
+    state = result.state
+    status = _CELERY_STATE_TO_STATUS.get(state, state)
+    kwargs: dict = result.kwargs or {}
+    action = kwargs.get("action", "control")
+    return BranchTaskPublic(
+        id=task_id,
+        task_type=action,
+        status=status,
+        parameters={"action": action},
+        result=result.result if state == "SUCCESS" else None,
+        error=str(result.traceback) if state == "FAILURE" and result.traceback else None,
+        date_done=result.date_done,
+    )
+
+
 @task_api.get(
     "/",
     name="organizations:projects:branch:tasks:list",
@@ -64,9 +82,12 @@ async def list_tasks(
     _project: ProjectDep,
     branch: BranchDep,
 ) -> list[BranchTaskPublic]:
-    if branch.resize_task_id is None:
-        return []
-    return [_build_task_public(branch.resize_task_id)]
+    tasks = []
+    if branch.resize_task_id is not None:
+        tasks.append(_build_resize_task_public(branch.resize_task_id))
+    if branch.control_task_id is not None:
+        tasks.append(_build_control_task_public(branch.control_task_id))
+    return tasks
 
 
 @task_api.get(
@@ -81,6 +102,8 @@ async def get_task(
     branch: BranchDep,
     task_id: UUID,
 ) -> BranchTaskPublic:
-    if branch.resize_task_id != task_id:
-        raise HTTPException(status_code=404, detail="Task not found")
-    return _build_task_public(task_id)
+    if branch.resize_task_id == task_id:
+        return _build_resize_task_public(task_id)
+    if branch.control_task_id == task_id:
+        return _build_control_task_public(task_id)
+    raise HTTPException(status_code=404, detail="Task not found")

src/models/branch.py

@@ -3,7 +3,6 @@
 from enum import StrEnum
 from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, Optional
 
-import sqlalchemy as sa
 from pydantic import BaseModel, model_validator
 from pydantic import Field as PydanticField
 from sqlalchemy import BigInteger, Boolean, Column, String, Text, UniqueConstraint, text
@@ -109,7 +108,8 @@ class Branch(AsyncAttrs, Model, table=True):
     )
     pitr_enabled: bool = Field(default=False, sa_column=Column(Boolean, nullable=False, server_default=text("false")))
     resize_task_id: uuid.UUID | None = Field(default=None, nullable=True)
-    db_port: int | None = Field(default=None, sa_column=Column(sa.Integer, nullable=True))
+    db_port: int | None = None
+    control_task_id: uuid.UUID | None = None
 
     __table_args__ = (UniqueConstraint("project_id", "name", name="unique_branch_name_per_project"),)
 

src/models/migrations/versions/b1c2d3e4f5a6_branch_control_task_id.py

@@ -0,0 +1,26 @@
+"""Add control_task_id to branch table
+
+Revision ID: b1c2d3e4f5a6
+Revises: 07822c477427
+Create Date: 2026-04-01 00:00:00.000000
+
+"""
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision: str = "b1c2d3e4f5a6"
+down_revision: Union[str, Sequence[str], None] = "07822c477427"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column("branch", sa.Column("control_task_id", sa.Uuid(), nullable=True))
+
+
+def downgrade() -> None:
+    op.drop_column("branch", "control_task_id")