From 8e6d337e7c77b89d5e5b64c6757c8743ab5c6695 Mon Sep 17 00:00:00 2001 From: Tim Kambic Date: Thu, 16 Jul 2026 13:19:30 +0200 Subject: [PATCH] fix: claim update slot atomically --- goosebit/db/models.py | 1 + goosebit/device_manager.py | 39 +++++++++ goosebit/updater/controller/v1/routes.py | 17 ++-- .../unit/updater/controller/v1/test_routes.py | 79 ++++++++++++++++++- 4 files changed, 125 insertions(+), 11 deletions(-) diff --git a/goosebit/db/models.py b/goosebit/db/models.py index 2fa0243b..48ab356a 100644 --- a/goosebit/db/models.py +++ b/goosebit/db/models.py @@ -38,6 +38,7 @@ class UpdateStateEnum(IntEnum): RUNNING = 3 ERROR = 4 FINISHED = 5 + RESERVED = 6 def __str__(self) -> str: return self.name.capitalize() diff --git a/goosebit/device_manager.py b/goosebit/device_manager.py index e42729dc..885aa0f0 100644 --- a/goosebit/device_manager.py +++ b/goosebit/device_manager.py @@ -7,6 +7,8 @@ from aiocache import caches from fastapi.requests import Request +from tortoise.expressions import Subquery +from tortoise.functions import Count from goosebit.db.models import ( Device, @@ -160,6 +162,43 @@ async def update_config_data(device: Device, **kwargs: dict[str, Any]) -> None: if modified: await DeviceManager.save_device(device, update_fields=["hardware_id", "last_state", "sw_version"]) + @staticmethod + async def try_claim_update_slot(device: Device, max_concurrent: int) -> bool: + """Atomically move the device into RESERVED to claim an update slot. + + Both RESERVED and RUNNING count against the cap. The device is promoted to + RUNNING (and its log/progress reset) only once it reports progress. + """ + # a device already holding a slot keeps it; swupdate needs the link re-served + if device.last_state in (UpdateStateEnum.RESERVED, UpdateStateEnum.RUNNING): + return True + + occupied = Subquery( + Device.filter(last_state__in=[UpdateStateEnum.RESERVED, UpdateStateEnum.RUNNING]) + .annotate(count=Count("id")) + .values("count") + ) + # count occupied slots and claim in one statement: atomic on SQLite, ms-scale + # residual race on PostgreSQL READ COMMITTED + rowcount = ( + await Device.filter( + id=device.id, + last_state__not_in=[UpdateStateEnum.RESERVED, UpdateStateEnum.RUNNING], + ) + .annotate(occupied=occupied) + .filter(occupied__lt=max_concurrent) + .update(last_state=UpdateStateEnum.RESERVED) + ) + if rowcount != 1: + return False + + device.last_state = UpdateStateEnum.RESERVED + # the claim is already committed; only keep the cached copy coherent (no further + # DB write). + result = await caches.get("default").set(device.id, device, ttl=600) + assert result, "device being cached" + return True + @staticmethod async def deployment_action_start(device: Device) -> None: device.last_log = "" diff --git a/goosebit/updater/controller/v1/routes.py b/goosebit/updater/controller/v1/routes.py index 1a9db13e..c8f3c76e 100644 --- a/goosebit/updater/controller/v1/routes.py +++ b/goosebit/updater/controller/v1/routes.py @@ -58,8 +58,8 @@ async def polling(request: Request, device: Device = Depends(get_device)) -> dic # won't confirm a successful testing (might be a bug/problem in swupdate) handling_type, software = await DeviceManager.get_update(device) if handling_type != HandlingType.SKIP and software is not None: - number_of_running = await Device.filter(last_state=UpdateStateEnum.RUNNING).count() - if number_of_running < config.max_concurrent_updates or device.last_state == UpdateStateEnum.RUNNING: + # claim a slot before handing out the link + if await DeviceManager.try_claim_update_slot(device, config.max_concurrent_updates): links["deploymentBase"] = { "href": str( request.url_for( @@ -71,12 +71,11 @@ async def polling(request: Request, device: Device = Depends(get_device)) -> dic } logger.info(f"Forced: update available, device={device.id}") else: - number_of_running = await Device.filter(last_state=UpdateStateEnum.RUNNING).count() - if number_of_running < config.max_concurrent_updates or device.last_state == UpdateStateEnum.RUNNING: - plugin_sources = await DeviceManager.get_alt_src_updates(request, device) - for handling_type, _ in plugin_sources: - if handling_type == HandlingType.SKIP: - continue + plugin_sources = await DeviceManager.get_alt_src_updates(request, device) + for handling_type, _ in plugin_sources: + if handling_type == HandlingType.SKIP: + continue + if await DeviceManager.try_claim_update_slot(device, config.max_concurrent_updates): links["deploymentBase"] = { "href": str( request.url_for( @@ -86,7 +85,7 @@ async def polling(request: Request, device: Device = Depends(get_device)) -> dic ) ) } - break + break return { "config": {"polling": {"sleep": sleep}}, "_links": links, diff --git a/tests/unit/updater/controller/v1/test_routes.py b/tests/unit/updater/controller/v1/test_routes.py index 26a7790e..36cb9419 100644 --- a/tests/unit/updater/controller/v1/test_routes.py +++ b/tests/unit/updater/controller/v1/test_routes.py @@ -3,7 +3,7 @@ import pytest from httpx import AsyncClient -from goosebit.db.models import Device, Hardware, Software +from goosebit.db.models import Device, Hardware, Software, UpdateStateEnum from goosebit.device_manager import DeviceManager, get_device from goosebit.settings import config @@ -308,7 +308,8 @@ async def _assert_log_lines(async_client: AsyncClient, device: Device, expected_ assert response.status_code == 200 log = response.json()["log"] - if log is None: + # a claim resets last_log to "" (previously NULL until first feedback) + if not log: assert expected_line_count == 0 else: actual_line_count = log.count("\n") @@ -353,3 +354,77 @@ async def test_update_logs_and_progress(async_client: AsyncClient, test_data: Di # fake installation start confirmation to check clearing of logs await _feedback(async_client, device.id, software, "none", "proceeding", "Downloaded 1%") await _assert_log_lines(async_client, device, 1) + + +@pytest.mark.asyncio +async def test_concurrent_update_cap( + async_client: AsyncClient, test_data: Dict[str, Any], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(config, "max_concurrent_updates", 1) + + device1 = test_data["device_rollout"] + device2 = test_data["device_assigned"] + software = test_data["software_release"] + + # first device claims the only slot (reserved at hand-out, before the install starts) + await _poll(async_client, device1.id, software) + device_api = await _api_device_get(async_client, device1.id) + assert device_api["last_state"] == "Reserved" + + # cap exhausted: no link for the second device + await _poll(async_client, device2.id, software, expect_update=False) + + # the reserved device keeps receiving its link on subsequent polls + await _poll(async_client, device1.id, software) + + # first device finishes, freeing the slot + await _feedback(async_client, device1.id, software, "success", "closed") + + # second device now gets the link + await _poll(async_client, device2.id, software) + + +@pytest.mark.asyncio +async def test_running_device_keeps_slot_when_cap_exhausted( + async_client: AsyncClient, test_data: Dict[str, Any], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(config, "max_concurrent_updates", 1) + + device1 = test_data["device_rollout"] + device2 = test_data["device_assigned"] + software = test_data["software_release"] + + # device1 already occupies the single slot + device1 = await get_device(dev_id=device1.id) + await DeviceManager.update_device_state(device1, UpdateStateEnum.RUNNING) + + # cap is exhausted for a fresh claim + await _poll(async_client, device2.id, software, expect_update=False) + + # running device keeps receiving the link + await _poll(async_client, device1.id, software) + + +@pytest.mark.asyncio +async def test_claim_reserves_and_proceeding_resets(async_client: AsyncClient, test_data: Dict[str, Any]) -> None: + device = test_data["device_rollout"] + software = test_data["software_release"] + + # stale bookkeeping from a previous update + device.last_log = "stale log entry\n" + device.progress = 50 + await device.save(update_fields=["last_log", "progress"]) + + # claiming only reserves the slot; bookkeeping is left untouched until the install starts + await _poll(async_client, device.id, software) + await device.refresh_from_db() + assert device.last_state == UpdateStateEnum.RESERVED + assert device.last_log == "stale log entry\n" + assert device.progress == 50 + + # first proceeding feedback promotes to running and resets the stale bookkeeping + await _feedback(async_client, device.id, software, "none", "proceeding") + await device.refresh_from_db() + assert device.last_state == UpdateStateEnum.RUNNING + assert "stale log entry" not in (device.last_log or "") + assert device.progress == 0