fix: claim update slot atomically

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()

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 = ""

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,

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