retoor
9cfaddfc40
The creation and run quotas were checked and then acted on, so two concurrent create_task calls or two schedulers could both pass the check and overshoot the limit. Both are now a single conditional INSERT decided on the driver rowcount: reserve_run takes a run slot after the claim and releases the claim by deferring when the quota is spent, and insert_task_within_quota does the same for the task row itself. Racing twelve and sixteen processes now yields exactly the limit. The atomic insert names its columns, and dataset skips a None valued key when it creates a table lazily, so the store declares the full task column set up front. Both the column and index ensures now tolerate a concurrent duplicate, since several processes build a store at once and SQLite DDL is not idempotent. Adds the quota, task-run context, guard, store and scheduler test suites, and documents the chokepoints and the unhackable task-run flag.
254 lines
7.7 KiB
Python
254 lines
7.7 KiB
Python
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
import dataset
|
|
import sqlalchemy
|
|
|
|
from .guards import (
|
|
AutomationDenied,
|
|
REASON_CREATE_QUOTA,
|
|
REASON_NESTED,
|
|
REASON_NOT_A_USER,
|
|
automation_allowed,
|
|
creation_denial,
|
|
nesting_allowed,
|
|
)
|
|
from .limits import create_quota, insert_task_within_quota
|
|
from .schedule import now_utc
|
|
|
|
logger = logging.getLogger("devii.tasks.store")
|
|
|
|
TABLE = "devii_tasks"
|
|
INDEXED_COLUMNS = (
|
|
["owner_kind", "owner_id"],
|
|
["uid"],
|
|
["enabled"],
|
|
["next_run_at"],
|
|
["status"],
|
|
)
|
|
TASK_COLUMNS = (
|
|
("uid", ""),
|
|
("owner_kind", ""),
|
|
("owner_id", ""),
|
|
("label", ""),
|
|
("prompt", ""),
|
|
("enabled", True),
|
|
("status", ""),
|
|
("created_at", ""),
|
|
("next_run_at", ""),
|
|
("last_run_at", ""),
|
|
("run_count", 0),
|
|
("last_result", ""),
|
|
("last_error", ""),
|
|
("kind", ""),
|
|
("run_at", ""),
|
|
("delay_seconds", 0),
|
|
("every_seconds", 0),
|
|
("start_at", ""),
|
|
("cron", ""),
|
|
("max_runs", 0),
|
|
("deleted_at", ""),
|
|
("deleted_by", ""),
|
|
("notify", 0),
|
|
("tz", ""),
|
|
("expires_at", ""),
|
|
("failure_count", 0),
|
|
)
|
|
DUE_SQL = (
|
|
"SELECT * FROM devii_tasks WHERE enabled = 1 AND status = 'pending' "
|
|
"AND deleted_at IS NULL AND next_run_at IS NOT NULL AND next_run_at <= :now "
|
|
"ORDER BY next_run_at LIMIT :limit"
|
|
)
|
|
CLAIM_SQL = (
|
|
"UPDATE devii_tasks SET status = 'running' WHERE uid = :uid "
|
|
"AND status = 'pending' AND enabled = 1 AND deleted_at IS NULL"
|
|
)
|
|
|
|
|
|
def memory_db() -> Any:
|
|
return dataset.connect("sqlite:///:memory:")
|
|
|
|
|
|
ACTIVE_STATUSES = ("pending", "running")
|
|
|
|
|
|
def due_rows(db: Any, now_iso: str, limit: int) -> list[dict[str, Any]]:
|
|
if TABLE not in db.tables:
|
|
return []
|
|
return list(db.query(DUE_SQL, now=now_iso, limit=limit))
|
|
|
|
|
|
def claim(db: Any, uid: str) -> bool:
|
|
statement = sqlalchemy.text(CLAIM_SQL)
|
|
with db:
|
|
result = db.executable.execute(statement, {"uid": uid})
|
|
return result.rowcount == 1
|
|
|
|
|
|
class TaskStore:
|
|
def __init__(
|
|
self, db: Any, owner_kind: str, owner_id: str, operator: bool = False
|
|
) -> None:
|
|
self._db = db
|
|
self._owner_kind = owner_kind
|
|
self._owner_id = owner_id
|
|
self._operator = operator
|
|
self._ensure_indexes()
|
|
|
|
def _ensure_indexes(self) -> None:
|
|
if TABLE not in self._db.tables:
|
|
return
|
|
self._ensure_schema()
|
|
for columns in INDEXED_COLUMNS:
|
|
try:
|
|
self._table.create_index(columns)
|
|
except sqlalchemy.exc.OperationalError as exc:
|
|
if "already exists" not in str(exc).lower():
|
|
raise
|
|
logger.debug("Index on %s was created by another process", columns)
|
|
|
|
def _ensure_schema(self) -> None:
|
|
table = self._db[TABLE]
|
|
for column, example in TASK_COLUMNS:
|
|
if table.has_column(column):
|
|
continue
|
|
try:
|
|
table.create_column_by_example(column, example)
|
|
except sqlalchemy.exc.OperationalError as exc:
|
|
if "duplicate column" not in str(exc).lower():
|
|
raise
|
|
logger.debug("Column %s was added by another process", column)
|
|
table._reflect_table()
|
|
|
|
@property
|
|
def db(self) -> Any:
|
|
return self._db
|
|
|
|
def automation_allowed(self) -> bool:
|
|
if self._operator:
|
|
return True
|
|
return automation_allowed(self._owner_kind, self._owner_id)
|
|
|
|
def require_automation(self) -> None:
|
|
if not self.automation_allowed():
|
|
raise AutomationDenied(REASON_NOT_A_USER)
|
|
|
|
def require_scheduling_allowed(self) -> None:
|
|
self.require_automation()
|
|
if self._operator:
|
|
return
|
|
if not nesting_allowed(self._owner_id):
|
|
raise AutomationDenied(REASON_NESTED)
|
|
|
|
def require_creation_allowed(self) -> None:
|
|
if self._operator:
|
|
return
|
|
denial = creation_denial(self._db, self._owner_kind, self._owner_id, now_utc())
|
|
if denial is not None:
|
|
raise denial
|
|
|
|
def count_active(self) -> int:
|
|
rows = self._table.find(enabled=True, deleted_at=None, **self._scope)
|
|
return sum(1 for row in rows if row.get("status") in ACTIVE_STATUSES)
|
|
|
|
@property
|
|
def _table(self) -> Any:
|
|
return self._db[TABLE]
|
|
|
|
@property
|
|
def _scope(self) -> dict[str, str]:
|
|
return {"owner_kind": self._owner_kind, "owner_id": self._owner_id}
|
|
|
|
def create(self, record: dict[str, Any]) -> None:
|
|
self.require_creation_allowed()
|
|
self._ensure_schema()
|
|
row = {
|
|
"deleted_at": None,
|
|
"deleted_by": None,
|
|
"failure_count": 0,
|
|
**record,
|
|
**self._scope,
|
|
}
|
|
if self._operator:
|
|
self._table.insert(row)
|
|
elif not insert_task_within_quota(
|
|
self._db, row, self._owner_kind, self._owner_id, now_utc()
|
|
):
|
|
quota = create_quota(self._db, self._owner_kind, self._owner_id, now_utc())
|
|
raise AutomationDenied(REASON_CREATE_QUOTA, quota.free_at)
|
|
logger.info(
|
|
"Task created uid=%s owner=%s/%s",
|
|
record.get("uid"),
|
|
self._owner_kind,
|
|
self._owner_id,
|
|
)
|
|
|
|
def get(self, uid: str) -> dict[str, Any] | None:
|
|
return self._table.find_one(uid=uid, deleted_at=None, **self._scope)
|
|
|
|
def list(
|
|
self, enabled_only: bool = False, status: str | None = None
|
|
) -> list[dict[str, Any]]:
|
|
criteria: dict[str, Any] = dict(self._scope)
|
|
criteria["deleted_at"] = None
|
|
if enabled_only:
|
|
criteria["enabled"] = True
|
|
if status:
|
|
criteria["status"] = status
|
|
return list(self._table.find(**criteria))
|
|
|
|
def update(self, uid: str, changes: dict[str, Any]) -> None:
|
|
if changes.get("enabled"):
|
|
self.require_scheduling_allowed()
|
|
changes = {**changes, "uid": uid, **self._scope}
|
|
self._table.update(changes, ["uid", "owner_kind", "owner_id"])
|
|
logger.debug("Task updated uid=%s changes=%s", uid, list(changes))
|
|
|
|
def delete(self, uid: str) -> bool:
|
|
row = self._table.find_one(uid=uid, deleted_at=None, **self._scope)
|
|
if not row:
|
|
return False
|
|
self._table.update(
|
|
{
|
|
"uid": uid,
|
|
"deleted_at": datetime.now(timezone.utc).isoformat(),
|
|
"deleted_by": f"{self._owner_kind}:{self._owner_id}",
|
|
**self._scope,
|
|
},
|
|
["uid", "owner_kind", "owner_id"],
|
|
)
|
|
logger.info("Task soft-deleted uid=%s", uid)
|
|
return True
|
|
|
|
def recover_running(self) -> int:
|
|
if TABLE not in self._db.tables:
|
|
return 0
|
|
stuck = list(self._table.find(status="running", deleted_at=None, **self._scope))
|
|
for row in stuck:
|
|
self._table.update(
|
|
{"uid": row["uid"], "status": "pending", **self._scope},
|
|
["uid", "owner_kind", "owner_id"],
|
|
)
|
|
if stuck:
|
|
logger.info("Recovered %d task(s) stuck in running", len(stuck))
|
|
return len(stuck)
|
|
|
|
def has_pending(self) -> bool:
|
|
rows = self._table.find(enabled=True, deleted_at=None, **self._scope)
|
|
return any(row.get("status") in ACTIVE_STATUSES for row in rows)
|
|
|
|
def due(self, now_iso: str) -> list[dict[str, Any]]:
|
|
rows = self._table.find(
|
|
enabled=True, status="pending", deleted_at=None, **self._scope
|
|
)
|
|
return [
|
|
row
|
|
for row in rows
|
|
if row.get("next_run_at") and row["next_run_at"] <= now_iso
|
|
]
|