feat(cloud-api): run lifecycle workers
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
@@ -19,7 +20,9 @@ from cloud.internal_api.api import create_internal_router
|
||||
from cloud.plugins import PluginRegistry
|
||||
from cloud.pool import DevicePool
|
||||
from cloud.scheduler import TaskScheduler
|
||||
from cloud.schema import require_current_schema
|
||||
from cloud.sdk.api import create_cloud_router
|
||||
from core.models import utc_now
|
||||
|
||||
|
||||
DatabaseFactory = Callable[[CloudControlConfig], CloudDatabase]
|
||||
@@ -85,13 +88,43 @@ def create_app(
|
||||
app.state.cloud_config = control_config
|
||||
app.state.database = database
|
||||
app.state.cloud_services = services
|
||||
stop_workers = asyncio.Event()
|
||||
worker_tasks: list[asyncio.Task[None]] = []
|
||||
try:
|
||||
repository.health_check()
|
||||
if control_config.environment == "production":
|
||||
require_current_schema(control_config.database_url)
|
||||
worker_tasks = [
|
||||
asyncio.create_task(
|
||||
_run_scheduler_loop(
|
||||
services,
|
||||
stop_workers,
|
||||
interval_seconds=control_config.scheduler_interval_seconds,
|
||||
),
|
||||
name="cloud-scheduler",
|
||||
),
|
||||
asyncio.create_task(
|
||||
_run_lease_reaper_loop(
|
||||
services,
|
||||
stop_workers,
|
||||
interval_seconds=(control_config.lease_reaper_interval_seconds),
|
||||
max_attempts=control_config.max_task_attempts,
|
||||
),
|
||||
name="cloud-lease-reaper",
|
||||
),
|
||||
]
|
||||
app.state.worker_tasks = tuple(worker_tasks)
|
||||
yield
|
||||
finally:
|
||||
stop_workers.set()
|
||||
try:
|
||||
database.close()
|
||||
if worker_tasks:
|
||||
await asyncio.gather(*worker_tasks)
|
||||
finally:
|
||||
repository.unbind()
|
||||
try:
|
||||
database.close()
|
||||
finally:
|
||||
repository.unbind()
|
||||
|
||||
app = FastAPI(title="Device Cloud API", lifespan=lifespan)
|
||||
app.include_router(
|
||||
@@ -117,3 +150,39 @@ def _default_database_factory(config: CloudControlConfig) -> CloudDatabase:
|
||||
config.database_url,
|
||||
create_schema=config.environment != "production",
|
||||
)
|
||||
|
||||
|
||||
async def _run_scheduler_loop(
|
||||
services: CloudApplicationServices,
|
||||
stop: asyncio.Event,
|
||||
*,
|
||||
interval_seconds: float,
|
||||
) -> None:
|
||||
while not stop.is_set():
|
||||
services.scheduler.assign()
|
||||
if await _wait_for_stop(stop, interval_seconds):
|
||||
return
|
||||
|
||||
|
||||
async def _run_lease_reaper_loop(
|
||||
services: CloudApplicationServices,
|
||||
stop: asyncio.Event,
|
||||
*,
|
||||
interval_seconds: float,
|
||||
max_attempts: int,
|
||||
) -> None:
|
||||
while not stop.is_set():
|
||||
services.repository.reap_expired_leases(
|
||||
now=utc_now(),
|
||||
max_attempts=max_attempts,
|
||||
)
|
||||
if await _wait_for_stop(stop, interval_seconds):
|
||||
return
|
||||
|
||||
|
||||
async def _wait_for_stop(stop: asyncio.Event, interval_seconds: float) -> bool:
|
||||
try:
|
||||
await asyncio.wait_for(stop.wait(), timeout=interval_seconds)
|
||||
except TimeoutError:
|
||||
return False
|
||||
return True
|
||||
|
||||
Reference in New Issue
Block a user