Initial commit: VMware vSphere API simulator scaffold.

Add the FastAPI app, PostgreSQL migrations, Docker/Helm packaging, API
contracts, docs, client examples, and the unit/integration/compatibility
test suite for local client and tooling labs without a real vCenter.
This commit is contained in:
2026-07-18 04:42:11 +03:00
commit f8d3cbdd59
422 changed files with 361335 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
"""Durable asynchronous task engine."""
+86
View File
@@ -0,0 +1,86 @@
"""Worker semantics for backup/vzdump tasks."""
from __future__ import annotations
import json
from typing import Any
from app.simulation.clock import Clock
from app.simulation.seed import stable_id
from app.tasks.repository import Task, TaskRepository
from app.tasks.worker import TaskHandler
def backup_handler(repository: TaskRepository, clock: Clock) -> TaskHandler:
async def execute(task: Task) -> dict[str, Any]:
if task.task_type == "aptupdate":
node = str(task.payload.get("node", "unknown"))
await repository.append_log(task.id, f"starting apt update on {node}")
await clock.sleep(1.0)
async with repository.pool.acquire() as connection:
metadata = await connection.fetchval(
"SELECT metadata FROM nodes WHERE name=$1",
node,
)
if metadata is not None:
payload = json.loads(metadata) if isinstance(metadata, str) else dict(metadata)
ops = payload.setdefault("ops", {})
apt = ops.setdefault("apt", {})
packages = list(apt.get("packages") or [])
for package in packages:
if isinstance(package, dict) and package.get("Status") == "upgradable":
package["Status"] = "installed"
if package.get("Version"):
package["OldVersion"] = package["Version"]
apt["packages"] = packages
apt["update"] = {"status": "stopped", "exitstatus": "OK"}
payload["ops"] = ops
await connection.execute(
"UPDATE nodes SET metadata=$2::jsonb, updated_at=now() WHERE name=$1",
node,
json.dumps(payload, sort_keys=True),
)
await repository.append_log(task.id, "apt update finished")
return {"status": "OK"}
node = str(task.payload["node"])
vmids = [str(item) for item in task.payload.get("vmids", [])]
storage_id = str(task.payload.get("storage") or "nfs-backup")
await repository.append_log(task.id, f"starting vzdump on {node} for {len(vmids)} guests")
async with repository.pool.acquire() as connection:
storage_resource_id = await connection.fetchval(
"SELECT resource_id FROM storages WHERE storage_id=$1",
storage_id,
)
if storage_resource_id is None:
raise ValueError(f"storage {storage_id} does not exist")
created = 0
for index, vmid in enumerate(vmids):
resource_id = await connection.fetchval(
"""SELECT r.id FROM resources r JOIN nodes n ON n.id=r.node_id
WHERE n.name=$1 AND r.kind='qemu' AND r.external_id=$2""",
node,
vmid,
)
volume_id = f"backup/vzdump-qemu-{vmid}-{task.id.hex[:8]}-{index:04d}.vma.zst"
await connection.execute(
"""INSERT INTO backups(
id, resource_id, storage_resource_id, volume_id, size_bytes, metadata
) VALUES($1, $2, $3, $4, $5, $6::jsonb)
ON CONFLICT (storage_resource_id, volume_id) DO NOTHING""",
stable_id(f"backup-task:{task.id}:{vmid}"),
resource_id,
storage_resource_id,
volume_id,
(8 + index) * 1024**3,
json.dumps(
{"mode": task.payload.get("mode", "snapshot"), "type": "vzdump"},
sort_keys=True,
),
)
created += 1
await repository.append_log(task.id, f"backup archive created: {volume_id}")
await repository.append_log(task.id, f"vzdump finished ({created} archives)")
return {"created": created}
return execute
+245
View File
@@ -0,0 +1,245 @@
"""Worker semantics for asynchronous LXC transitions."""
from __future__ import annotations
import json
import uuid
from collections.abc import Mapping
from typing import Any, cast
from app.simulation.clock import Clock
from app.simulation.transitions import VmState, plan_transition
from app.tasks.repository import Task, TaskRepository
from app.tasks.worker import TaskHandler
def lxc_handler(repository: TaskRepository, clock: Clock) -> TaskHandler:
async def execute(task: Task) -> dict[str, Any]:
operation = task.task_type.removeprefix("lxc-")
if operation == "create":
return await _create(repository, task, clock)
if operation == "clone":
return await _clone(repository, task)
resource_id = uuid.UUID(str(task.payload["resource_id"]))
if operation == "delete":
return await _delete(repository, task, resource_id)
if operation.startswith("snapshot-"):
return await _snapshot(
repository, task, resource_id, operation.removeprefix("snapshot-")
)
if operation == "migrate" or operation == "remote-migrate":
return await _migrate(repository, task, resource_id, clock)
async with repository.pool.acquire() as connection:
row = await connection.fetchrow("SELECT state FROM resources WHERE id=$1", resource_id)
if row is None:
raise ValueError("resource disappeared")
state = _object(row["state"])
transition = plan_transition(VmState(str(state["status"])), operation)
state["status"] = transition.intermediate
await connection.execute(
"UPDATE resources SET state=$2::jsonb WHERE id=$1",
resource_id,
json.dumps(state),
)
await repository.append_log(task.id, f"container {operation} started")
await clock.sleep(1.0)
async with repository.pool.acquire() as connection:
state["status"] = transition.after
await connection.execute(
"UPDATE resources SET state=$2::jsonb WHERE id=$1",
resource_id,
json.dumps(state),
)
await repository.append_log(task.id, f"container {operation} completed")
return {"status": str(transition.after)}
return execute
async def _create(repository: TaskRepository, task: Task, clock: Clock) -> dict[str, Any]:
node, vmid = str(task.payload["node"]), int(task.payload["vmid"])
config = dict(task.payload.get("config", {}))
start = bool(task.payload.get("start", False))
resource_id = uuid.uuid4()
status = "running" if start else "stopped"
state = {"status": status, **config}
async with repository.pool.acquire() as connection:
async with connection.transaction():
node_row = await connection.fetchrow(
"SELECT id, cluster_id FROM nodes WHERE name=$1", node
)
if node_row is None:
raise ValueError("node disappeared")
await connection.execute(
"""INSERT INTO resources(
id, node_id, cluster_id, kind, external_id, state, metadata
) VALUES($1, $2, $3, 'lxc', $4, $5::jsonb, '{}'::jsonb)""",
resource_id,
node_row["id"],
node_row["cluster_id"],
str(vmid),
json.dumps(state, sort_keys=True),
)
await connection.execute(
"""INSERT INTO containers(resource_id, cluster_id, vmid, config)
VALUES($1, $2, $3, $4::jsonb)""",
resource_id,
node_row["cluster_id"],
vmid,
json.dumps(config, sort_keys=True),
)
if start:
await clock.sleep(0.5)
await repository.append_log(task.id, f"container {vmid} created")
return {"vmid": vmid, "status": status}
async def _delete(repository: TaskRepository, task: Task, resource_id: uuid.UUID) -> dict[str, Any]:
async with repository.pool.acquire() as connection:
status = await connection.execute("DELETE FROM resources WHERE id=$1", resource_id)
if status != "DELETE 1":
raise ValueError("resource disappeared")
await repository.append_log(task.id, "container deleted")
return {"deleted": True}
async def _snapshot(
repository: TaskRepository,
task: Task,
resource_id: uuid.UUID,
operation: str,
) -> dict[str, Any]:
name = str(task.payload["snapname"])
async with repository.pool.acquire() as connection:
async with connection.transaction():
if operation == "create":
row = await connection.fetchrow(
"""SELECT r.state, c.config FROM resources r
JOIN containers c ON c.resource_id=r.id WHERE r.id=$1""",
resource_id,
)
if row is None:
raise ValueError("resource disappeared")
captured = {
"resource_state": _object(row["state"]),
"config": _object(row["config"]),
}
await connection.execute(
"""INSERT INTO snapshots(id, resource_id, name, description, state)
VALUES($1, $2, $3, $4, $5::jsonb)""",
uuid.uuid4(),
resource_id,
name,
str(task.payload.get("description", "")),
json.dumps(captured, sort_keys=True),
)
elif operation == "delete":
status = await connection.execute(
"DELETE FROM snapshots WHERE resource_id=$1 AND name=$2",
resource_id,
name,
)
if status != "DELETE 1":
raise ValueError("snapshot disappeared")
elif operation == "rollback":
row = await connection.fetchrow(
"SELECT state FROM snapshots WHERE resource_id=$1 AND name=$2",
resource_id,
name,
)
if row is None:
raise ValueError("snapshot disappeared")
captured = _object(row["state"])
state = dict(cast(Mapping[str, Any], captured["resource_state"]))
await connection.execute(
"""UPDATE resources SET state=$2::jsonb, version=version+1,
updated_at=now() WHERE id=$1""",
resource_id,
json.dumps(state, sort_keys=True),
)
await connection.execute(
"UPDATE containers SET config=$2::jsonb WHERE resource_id=$1",
resource_id,
json.dumps(captured["config"], sort_keys=True),
)
else:
raise ValueError(f"unsupported snapshot operation: {operation}")
await repository.append_log(task.id, f"snapshot {name} {operation} completed")
return {"snapshot": name, "operation": operation}
async def _clone(repository: TaskRepository, task: Task) -> dict[str, Any]:
source_id = uuid.UUID(str(task.payload["source_resource_id"]))
target_id = uuid.uuid4()
node, vmid = str(task.payload["node"]), int(task.payload["vmid"])
async with repository.pool.acquire() as connection:
async with connection.transaction():
source = await connection.fetchrow(
"""SELECT r.state, c.config FROM resources r
JOIN containers c ON c.resource_id=r.id WHERE r.id=$1""",
source_id,
)
target = await connection.fetchrow(
"SELECT id, cluster_id FROM nodes WHERE name=$1", node
)
if source is None or target is None:
raise ValueError("clone source or target disappeared")
config = _object(source["config"])
if task.payload.get("name") is not None:
config["hostname"] = task.payload["name"]
state = {**_object(source["state"]), **config, "status": "stopped"}
await connection.execute(
"""INSERT INTO resources(id,node_id,cluster_id,kind,external_id,state,metadata)
VALUES($1,$2,$3,'lxc',$4,$5::jsonb,'{}'::jsonb)""",
target_id,
target["id"],
target["cluster_id"],
str(vmid),
json.dumps(state),
)
await connection.execute(
"""INSERT INTO containers(resource_id,cluster_id,vmid,config)
VALUES($1,$2,$3,$4::jsonb)""",
target_id,
target["cluster_id"],
vmid,
json.dumps(config),
)
await repository.append_log(task.id, f"container cloned to {vmid}")
return {"vmid": vmid, "node": node}
async def _migrate(
repository: TaskRepository, task: Task, resource_id: uuid.UUID, clock: Clock
) -> dict[str, Any]:
target = str(task.payload["target"])
async with repository.pool.acquire() as connection:
row = await connection.fetchrow("SELECT state FROM resources WHERE id=$1", resource_id)
if row is None:
raise ValueError("resource disappeared")
state = _object(row["state"])
transition = plan_transition(VmState(str(state["status"])), "migrate")
state["status"] = transition.intermediate
await connection.execute(
"UPDATE resources SET state=$2::jsonb WHERE id=$1", resource_id, json.dumps(state)
)
await repository.append_log(task.id, f"migration to {target} started")
await clock.sleep(1.0)
async with repository.pool.acquire() as connection:
node = await connection.fetchrow("SELECT id FROM nodes WHERE name=$1", target)
if node is None:
raise ValueError("target node disappeared")
state["status"] = transition.after
await connection.execute(
"""UPDATE resources SET node_id=$2,state=$3::jsonb,version=version+1,
updated_at=now() WHERE id=$1""",
resource_id,
node["id"],
json.dumps(state),
)
await repository.append_log(task.id, f"migration to {target} completed")
return {"node": target, "status": str(transition.after)}
def _object(value: object) -> dict[str, Any]:
return json.loads(value) if isinstance(value, str) else dict(cast(Mapping[str, Any], value))
+328
View File
@@ -0,0 +1,328 @@
"""Worker semantics for asynchronous QEMU transitions."""
from __future__ import annotations
import json
import uuid
from collections.abc import Mapping
from typing import Any, cast
from app.simulation.clock import Clock
from app.simulation.transitions import VmState, plan_transition
from app.tasks.repository import Task, TaskRepository
from app.tasks.worker import TaskHandler
def qemu_handler(repository: TaskRepository, clock: Clock) -> TaskHandler:
async def execute(task: Task) -> dict[str, Any]:
operation = task.task_type.removeprefix("qemu-")
if operation == "create":
return await _create(repository, task)
if operation == "clone":
return await _clone(repository, task)
resource_id = uuid.UUID(str(task.payload["resource_id"]))
if operation == "update":
return await _update(repository, task, resource_id)
if operation == "delete":
return await _delete(repository, task, resource_id)
if operation.startswith("snapshot-"):
return await _snapshot(
repository, task, resource_id, operation.removeprefix("snapshot-")
)
if operation == "migrate" or operation == "remote-migrate":
return await _migrate(repository, task, resource_id, clock)
if operation == "move-disk":
return await _move_disk(repository, task, resource_id)
async with repository.pool.acquire() as connection:
row = await connection.fetchrow("SELECT state FROM resources WHERE id=$1", resource_id)
if row is None:
raise ValueError("resource disappeared")
raw = row["state"]
state = json.loads(raw) if isinstance(raw, str) else dict(raw)
transition = plan_transition(VmState(str(state["status"])), operation)
state["status"] = transition.intermediate
await connection.execute(
"UPDATE resources SET state=$2::jsonb WHERE id=$1",
resource_id,
json.dumps(state),
)
await repository.append_log(task.id, f"VM {operation} started")
await clock.sleep(1.0)
async with repository.pool.acquire() as connection:
state["status"] = transition.after
await connection.execute(
"UPDATE resources SET state=$2::jsonb WHERE id=$1",
resource_id,
json.dumps(state),
)
await repository.append_log(task.id, f"VM {operation} completed")
return {"status": str(transition.after)}
return execute
async def _create(repository: TaskRepository, task: Task) -> dict[str, Any]:
node, vmid = str(task.payload["node"]), int(task.payload["vmid"])
config = dict(task.payload.get("config", {}))
resource_id = uuid.uuid4()
state = {"status": "stopped", **config}
async with repository.pool.acquire() as connection:
async with connection.transaction():
node_row = await connection.fetchrow(
"SELECT id, cluster_id FROM nodes WHERE name=$1", node
)
if node_row is None:
raise ValueError("node disappeared")
await connection.execute(
"""INSERT INTO resources(
id, node_id, cluster_id, kind, external_id, state, metadata
) VALUES($1, $2, $3, 'qemu', $4, $5::jsonb, '{}'::jsonb)""",
resource_id,
node_row["id"],
node_row["cluster_id"],
str(vmid),
json.dumps(state, sort_keys=True),
)
await connection.execute(
"""INSERT INTO virtual_machines(resource_id, cluster_id, vmid, config)
VALUES($1, $2, $3, $4::jsonb)""",
resource_id,
node_row["cluster_id"],
vmid,
json.dumps(config, sort_keys=True),
)
await repository.append_log(task.id, f"VM {vmid} created")
return {"vmid": vmid, "status": "stopped"}
async def _update(repository: TaskRepository, task: Task, resource_id: uuid.UUID) -> dict[str, Any]:
changes = dict(task.payload.get("changes", {}))
delete_keys = tuple(str(task.payload.get("delete", "")).split(","))
async with repository.pool.acquire() as connection:
async with connection.transaction():
row = await connection.fetchrow(
"""SELECT r.state, v.config FROM resources r
JOIN virtual_machines v ON v.resource_id=r.id WHERE r.id=$1""",
resource_id,
)
if row is None:
raise ValueError("resource disappeared")
state = _object(row["state"])
config = _object(row["config"])
config.update(changes)
for key in delete_keys:
if key:
config.pop(key, None)
state.pop(key, None)
state.update(changes)
await connection.execute(
"""UPDATE resources SET state=$2::jsonb, version=version+1,
updated_at=now() WHERE id=$1""",
resource_id,
json.dumps(state, sort_keys=True),
)
await connection.execute(
"UPDATE virtual_machines SET config=$2::jsonb WHERE resource_id=$1",
resource_id,
json.dumps(config, sort_keys=True),
)
await repository.append_log(task.id, "VM configuration updated")
return {"updated": sorted(changes), "deleted": sorted(key for key in delete_keys if key)}
async def _delete(repository: TaskRepository, task: Task, resource_id: uuid.UUID) -> dict[str, Any]:
async with repository.pool.acquire() as connection:
status = await connection.execute("DELETE FROM resources WHERE id=$1", resource_id)
if status != "DELETE 1":
raise ValueError("resource disappeared")
await repository.append_log(task.id, "VM deleted")
return {"deleted": True}
async def _snapshot(
repository: TaskRepository,
task: Task,
resource_id: uuid.UUID,
operation: str,
) -> dict[str, Any]:
name = str(task.payload["snapname"])
async with repository.pool.acquire() as connection:
async with connection.transaction():
if operation == "create":
row = await connection.fetchrow(
"""SELECT r.state, v.config FROM resources r
JOIN virtual_machines v ON v.resource_id=r.id WHERE r.id=$1""",
resource_id,
)
if row is None:
raise ValueError("resource disappeared")
captured = {
"resource_state": _object(row["state"]),
"config": _object(row["config"]),
"vmstate": bool(task.payload.get("vmstate", False)),
}
await connection.execute(
"""INSERT INTO snapshots(id, resource_id, name, description, state)
VALUES($1, $2, $3, $4, $5::jsonb)""",
uuid.uuid4(),
resource_id,
name,
str(task.payload.get("description", "")),
json.dumps(captured, sort_keys=True),
)
elif operation == "delete":
status = await connection.execute(
"DELETE FROM snapshots WHERE resource_id=$1 AND name=$2",
resource_id,
name,
)
if status != "DELETE 1":
raise ValueError("snapshot disappeared")
elif operation == "rollback":
row = await connection.fetchrow(
"SELECT state FROM snapshots WHERE resource_id=$1 AND name=$2",
resource_id,
name,
)
if row is None:
raise ValueError("snapshot disappeared")
captured = _object(row["state"])
state = dict(cast(Mapping[str, Any], captured["resource_state"]))
if bool(task.payload.get("start", False)):
state["status"] = "running"
await connection.execute(
"""UPDATE resources SET state=$2::jsonb, version=version+1,
updated_at=now() WHERE id=$1""",
resource_id,
json.dumps(state, sort_keys=True),
)
await connection.execute(
"UPDATE virtual_machines SET config=$2::jsonb WHERE resource_id=$1",
resource_id,
json.dumps(captured["config"], sort_keys=True),
)
else:
raise ValueError(f"unsupported snapshot operation: {operation}")
await repository.append_log(task.id, f"snapshot {name} {operation} completed")
return {"snapshot": name, "operation": operation}
async def _clone(repository: TaskRepository, task: Task) -> dict[str, Any]:
source_id = uuid.UUID(str(task.payload["source_resource_id"]))
target_id = uuid.uuid4()
node, vmid = str(task.payload["node"]), int(task.payload["vmid"])
async with repository.pool.acquire() as connection:
async with connection.transaction():
source = await connection.fetchrow(
"""SELECT r.state, v.config FROM resources r
JOIN virtual_machines v ON v.resource_id=r.id WHERE r.id=$1""",
source_id,
)
target = await connection.fetchrow(
"SELECT id, cluster_id FROM nodes WHERE name=$1", node
)
if source is None or target is None:
raise ValueError("clone source or target disappeared")
config = _object(source["config"])
if task.payload.get("name") is not None:
config["name"] = task.payload["name"]
state = {**_object(source["state"]), **config, "status": "stopped"}
await connection.execute(
"""INSERT INTO resources(id,node_id,cluster_id,kind,external_id,state,metadata)
VALUES($1,$2,$3,'qemu',$4,$5::jsonb,'{}'::jsonb)""",
target_id,
target["id"],
target["cluster_id"],
str(vmid),
json.dumps(state),
)
await connection.execute(
"""INSERT INTO virtual_machines(resource_id,cluster_id,vmid,config)
VALUES($1,$2,$3,$4::jsonb)""",
target_id,
target["cluster_id"],
vmid,
json.dumps(config),
)
await repository.append_log(task.id, f"VM cloned to {vmid}")
return {"vmid": vmid, "node": node}
async def _migrate(
repository: TaskRepository, task: Task, resource_id: uuid.UUID, clock: Clock
) -> dict[str, Any]:
target = str(task.payload["target"])
async with repository.pool.acquire() as connection:
row = await connection.fetchrow("SELECT state FROM resources WHERE id=$1", resource_id)
if row is None:
raise ValueError("resource disappeared")
state = _object(row["state"])
transition = plan_transition(VmState(str(state["status"])), "migrate")
state["status"] = transition.intermediate
await connection.execute(
"UPDATE resources SET state=$2::jsonb WHERE id=$1", resource_id, json.dumps(state)
)
await repository.append_log(task.id, f"migration to {target} started")
await clock.sleep(1.0)
async with repository.pool.acquire() as connection:
node = await connection.fetchrow("SELECT id FROM nodes WHERE name=$1", target)
if node is None:
raise ValueError("target node disappeared")
state["status"] = transition.after
await connection.execute(
"""UPDATE resources SET node_id=$2,state=$3::jsonb,version=version+1,
updated_at=now() WHERE id=$1""",
resource_id,
node["id"],
json.dumps(state),
)
await repository.append_log(task.id, f"migration to {target} completed")
return {"node": target, "status": str(transition.after)}
async def _move_disk(
repository: TaskRepository, task: Task, resource_id: uuid.UUID
) -> dict[str, Any]:
disk = str(task.payload["disk"])
target_disk = str(task.payload["target_disk"])
storage = str(task.payload["storage"])
async with repository.pool.acquire() as connection:
async with connection.transaction():
row = await connection.fetchrow(
"SELECT config FROM virtual_machines WHERE resource_id=$1", resource_id
)
if row is None:
raise ValueError("resource disappeared")
config = _object(row["config"])
if disk not in config:
raise ValueError("disk disappeared")
original = str(config[disk])
suffix = original.split(":", 1)[1] if ":" in original else original
config[target_disk] = f"{storage}:{suffix}"
if bool(task.payload.get("delete", True)) and target_disk != disk:
config.pop(disk, None)
await connection.execute(
"UPDATE virtual_machines SET config=$2::jsonb WHERE resource_id=$1",
resource_id,
json.dumps(config, sort_keys=True),
)
await connection.execute(
"""UPDATE resources SET state=state || $2::jsonb, version=version+1,
updated_at=now() WHERE id=$1""",
resource_id,
json.dumps({target_disk: config[target_disk]}, sort_keys=True),
)
await connection.execute(
"""UPDATE vm_disks SET device=$2,storage_id=$3
WHERE resource_id=$1 AND device=$4""",
resource_id,
target_disk,
storage,
disk,
)
await repository.append_log(task.id, f"disk {disk} moved to {storage}")
return {"disk": target_disk, "storage": storage}
def _object(value: object) -> dict[str, Any]:
return json.loads(value) if isinstance(value, str) else dict(cast(Mapping[str, Any], value))
+197
View File
@@ -0,0 +1,197 @@
"""PostgreSQL repository for durable leased tasks."""
from __future__ import annotations
import json
import uuid
from dataclasses import dataclass
from typing import Any
import asyncpg # type: ignore[import-untyped] # noqa: F401
from asyncpg import Pool, Record
from app.db.primitives import ConflictError, require_affected, transaction
@dataclass(frozen=True, slots=True)
class Task:
id: uuid.UUID
upid: str
task_type: str
status: str
payload: dict[str, Any]
progress: int
cancel_requested: bool
attempt: int
def _task(row: Record) -> Task:
return Task(
id=row["id"],
upid=str(row["upid"]),
task_type=str(row["task_type"]),
status=str(row["status"]),
payload=json.loads(row["payload"])
if isinstance(row["payload"], str)
else dict(row["payload"]),
progress=int(row["progress"]),
cancel_requested=bool(row["cancel_requested"]),
attempt=int(row["attempt"]),
)
@dataclass(frozen=True, slots=True)
class TaskRepository:
pool: Pool
async def create(
self,
*,
upid: str,
task_type: str,
payload: dict[str, Any],
resource_key: str | None = None,
idempotency_key: str | None = None,
) -> Task:
task_id = uuid.uuid4()
async with transaction(self.pool) as connection:
if idempotency_key is not None:
existing = await connection.fetchrow(
"SELECT * FROM tasks WHERE idempotency_key=$1", idempotency_key
)
if existing is not None:
return _task(existing)
row = await connection.fetchrow(
"""INSERT INTO tasks(id, upid, task_type, status, payload, idempotency_key)
VALUES($1,$2,$3,'queued',$4::jsonb,$5) RETURNING *""",
task_id,
upid,
task_type,
json.dumps(payload),
idempotency_key,
)
if resource_key is not None:
try:
await connection.execute(
"INSERT INTO resource_locks(resource_key, task_id) VALUES($1,$2)",
resource_key,
task_id,
)
except Exception as error:
raise ConflictError(f"resource is locked: {resource_key}") from error
await connection.execute(
"INSERT INTO task_events(task_id, kind) VALUES($1,'created')", task_id
)
if row is None:
raise RuntimeError("task insert returned no row")
return _task(row)
async def claim(self, worker_id: str, lease_seconds: float) -> Task | None:
async with transaction(self.pool) as connection:
row = await connection.fetchrow(
"""WITH candidate AS (
SELECT id FROM tasks
WHERE status='queued' OR (status='running' AND lease_expires_at < now())
ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1
) UPDATE tasks SET status='running', worker_id=$1,
lease_expires_at=now() + $2 * interval '1 second', attempt=attempt+1,
updated_at=now()
WHERE id=(SELECT id FROM candidate) RETURNING *""",
worker_id,
lease_seconds,
)
if row is None:
return None
await connection.execute(
"INSERT INTO task_events(task_id, kind, data) VALUES($1,'claimed',$2::jsonb)",
row["id"],
json.dumps({"worker": worker_id}),
)
return _task(row)
async def heartbeat(self, task_id: uuid.UUID, worker_id: str, lease_seconds: float) -> None:
status = await self.pool.execute(
"""UPDATE tasks SET lease_expires_at=now()+$3*interval '1 second', updated_at=now()
WHERE id=$1 AND worker_id=$2 AND status='running'""",
task_id,
worker_id,
lease_seconds,
)
require_affected(status)
async def progress(self, task_id: uuid.UUID, worker_id: str, value: int) -> None:
status = await self.pool.execute(
"""UPDATE tasks SET progress=$3, updated_at=now()
WHERE id=$1 AND worker_id=$2 AND status='running'""",
task_id,
worker_id,
value,
)
require_affected(status)
async def append_log(self, task_id: uuid.UUID, message: str) -> None:
await self.pool.execute(
"INSERT INTO task_logs(task_id, message) VALUES($1,$2)", task_id, message
)
async def request_cancel(self, task_id: uuid.UUID) -> None:
status = await self.pool.execute(
"""UPDATE tasks SET cancel_requested=true, updated_at=now()
WHERE id=$1 AND status IN ('queued','running')""",
task_id,
)
require_affected(status)
async def finish(
self,
task_id: uuid.UUID,
worker_id: str,
*,
status: str,
result: dict[str, Any] | None = None,
error: str | None = None,
) -> None:
if status not in {"success", "error", "cancelled"}:
raise ValueError("invalid terminal task status")
async with transaction(self.pool) as connection:
command = await connection.execute(
"""UPDATE tasks SET status=$3, result=$4::jsonb, error=$5,
progress=CASE WHEN $3='success' THEN 100 ELSE progress END,
lease_expires_at=NULL, updated_at=now()
WHERE id=$1 AND worker_id=$2 AND status='running'""",
task_id,
worker_id,
status,
json.dumps(result) if result is not None else None,
error,
)
require_affected(command)
await connection.execute("DELETE FROM resource_locks WHERE task_id=$1", task_id)
await connection.execute(
"INSERT INTO task_events(task_id, kind, data) VALUES($1,$2,$3::jsonb)",
task_id,
status,
json.dumps({"error": error} if error else {}),
)
async def get(self, task_id: uuid.UUID) -> Task | None:
row = await self.pool.fetchrow("SELECT * FROM tasks WHERE id=$1", task_id)
return _task(row) if row is not None else None
async def get_by_upid(self, upid: str) -> Task | None:
row = await self.pool.fetchrow("SELECT * FROM tasks WHERE upid=$1", upid)
return _task(row) if row is not None else None
async def list_for_node(self, node: str) -> tuple[Task, ...]:
rows = await self.pool.fetch(
"""SELECT * FROM tasks WHERE payload->>'node'=$1
ORDER BY created_at DESC LIMIT 1000""",
node,
)
return tuple(_task(row) for row in rows)
async def logs(self, task_id: uuid.UUID) -> tuple[str, ...]:
rows = await self.pool.fetch(
"SELECT message FROM task_logs WHERE task_id=$1 ORDER BY sequence", task_id
)
return tuple(str(row["message"]) for row in rows)
+75
View File
@@ -0,0 +1,75 @@
"""Proxmox-compatible unique process/task identifiers."""
from __future__ import annotations
import re
import secrets
import time
from dataclasses import dataclass
UPID_RE = re.compile(
r"^UPID:(?P<node>[A-Za-z0-9][A-Za-z0-9_-]*):"
r"(?P<pid>[0-9A-Fa-f]{8}):(?P<pstart>[0-9A-Fa-f]{8}):"
r"(?P<start>[0-9A-Fa-f]{8}):(?P<type>[A-Za-z0-9_-]+):"
r"(?P<task_id>[^:]*):(?P<user>[^:]+):$"
)
@dataclass(frozen=True, slots=True)
class Upid:
node: str
pid: int
process_start: int
start_time: int
task_type: str
task_id: str
user: str
def __post_init__(self) -> None:
for name, value in (
("pid", self.pid),
("process_start", self.process_start),
("start_time", self.start_time),
):
if not 0 <= value <= 0xFFFFFFFF:
raise ValueError(f"{name} is outside the 32-bit UPID range")
if not self.node or ":" in self.node or not self.task_type or ":" in self.task_type:
raise ValueError("invalid UPID node or task type")
if ":" in self.task_id or not self.user or ":" in self.user:
raise ValueError("invalid UPID task id or user")
def __str__(self) -> str:
return (
f"UPID:{self.node}:{self.pid:08X}:{self.process_start:08X}:"
f"{self.start_time:08X}:{self.task_type}:{self.task_id}:{self.user}:"
)
@classmethod
def parse(cls, value: str) -> Upid:
match = UPID_RE.fullmatch(value)
if match is None:
raise ValueError("invalid UPID")
values = match.groupdict()
return cls(
node=values["node"],
pid=int(values["pid"], 16),
process_start=int(values["pstart"], 16),
start_time=int(values["start"], 16),
task_type=values["type"],
task_id=values["task_id"],
user=values["user"],
)
@classmethod
def allocate(cls, node: str, task_type: str, task_id: str, user: str) -> Upid:
"""Build a collision-resistant UPID for a new task."""
return cls(
node=node,
pid=secrets.randbits(32),
process_start=secrets.randbits(32),
start_time=int(time.time()) & 0xFFFFFFFF,
task_type=task_type,
task_id=str(task_id),
user=user,
)
+99
View File
@@ -0,0 +1,99 @@
"""Bounded durable task worker with cooperative cancellation."""
from __future__ import annotations
import asyncio
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any
from app.tasks.repository import Task, TaskRepository
TaskHandler = Callable[[Task], Awaitable[dict[str, Any] | None]]
logger = logging.getLogger(__name__)
@dataclass(slots=True)
class TaskWorker:
repository: TaskRepository
worker_id: str
handlers: dict[str, TaskHandler]
concurrency: int = 2
lease_seconds: float = 30.0
poll_seconds: float = 0.1
_running: set[asyncio.Task[None]] = field(default_factory=set, init=False)
_stopping: asyncio.Event = field(default_factory=asyncio.Event, init=False)
async def run(self) -> None:
self._stopping.clear()
try:
while not self._stopping.is_set():
self._reap()
if len(self._running) >= self.concurrency:
await asyncio.sleep(self.poll_seconds)
continue
try:
task = await self.repository.claim(self.worker_id, self.lease_seconds)
except asyncio.CancelledError:
raise
except Exception:
logger.exception("task claim failed; polling will retry")
await asyncio.sleep(self.poll_seconds)
continue
if task is None:
await asyncio.sleep(self.poll_seconds)
continue
execution = asyncio.create_task(self._execute(task))
self._running.add(execution)
finally:
if self._running:
await asyncio.gather(*self._running, return_exceptions=True)
self._running.clear()
def stop(self) -> None:
self._stopping.set()
def _reap(self) -> None:
self._running = {task for task in self._running if not task.done()}
async def _execute(self, task: Task) -> None:
handler = self.handlers.get(task.task_type)
if handler is None:
await self.repository.finish(
task.id, self.worker_id, status="error", error="unsupported task type"
)
return
try:
current = await self.repository.get(task.id)
if current is not None and current.cancel_requested:
await self.repository.finish(task.id, self.worker_id, status="cancelled")
return
execution: asyncio.Future[dict[str, Any] | None] = asyncio.ensure_future(handler(task))
heartbeat = asyncio.create_task(self._heartbeat(task))
try:
while not execution.done():
await asyncio.sleep(self.poll_seconds)
current = await self.repository.get(task.id)
if current is not None and current.cancel_requested:
execution.cancel()
await asyncio.gather(execution, return_exceptions=True)
await self.repository.finish(task.id, self.worker_id, status="cancelled")
return
result = await execution
finally:
heartbeat.cancel()
await asyncio.gather(heartbeat, return_exceptions=True)
await self.repository.finish(task.id, self.worker_id, status="success", result=result)
except asyncio.CancelledError:
raise
except Exception as error: # task failures are persisted, not leaked
await self.repository.finish(
task.id, self.worker_id, status="error", error=type(error).__name__
)
async def _heartbeat(self, task: Task) -> None:
interval = max(self.lease_seconds / 3, 0.01)
while True:
await asyncio.sleep(interval)
await self.repository.heartbeat(task.id, self.worker_id, self.lease_seconds)