66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
"""Durable-friendly recurring task scheduling primitives."""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Any
|
|
|
|
|
|
class RecurringTaskError(ValueError):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class RecurringTask:
|
|
task_id: str
|
|
interval_seconds: int
|
|
next_run_at: datetime
|
|
run_count: int = 0
|
|
max_runs: int | None = None
|
|
cancelled: bool = False
|
|
|
|
def due(self, now: datetime) -> bool:
|
|
return not self.cancelled and (self.max_runs is None or self.run_count < self.max_runs) and now >= self.next_run_at
|
|
|
|
|
|
class RecurringTaskScheduler:
|
|
def __init__(self) -> None:
|
|
self.tasks: dict[str, RecurringTask] = {}
|
|
self.successors: set[str] = set()
|
|
|
|
def schedule(self, task_id: str, interval_seconds: int, *, now: datetime | None = None, max_runs: int | None = None) -> RecurringTask:
|
|
if task_id in self.tasks:
|
|
raise RecurringTaskError("task already exists")
|
|
if interval_seconds <= 0 or (max_runs is not None and max_runs <= 0):
|
|
raise RecurringTaskError("invalid recurring task bounds")
|
|
now = now or datetime.now(timezone.utc)
|
|
task = RecurringTask(task_id, interval_seconds, now, 0, max_runs, False)
|
|
self.tasks[task_id] = task
|
|
return task
|
|
|
|
def due(self, *, now: datetime | None = None) -> list[RecurringTask]:
|
|
now = now or datetime.now(timezone.utc)
|
|
return [task for task in self.tasks.values() if task.due(now)]
|
|
|
|
def mark_run(self, task_id: str, *, now: datetime | None = None) -> RecurringTask:
|
|
task = self.tasks[task_id]
|
|
now = now or datetime.now(timezone.utc)
|
|
if not task.due(now):
|
|
raise RecurringTaskError("task is not due")
|
|
task.run_count += 1
|
|
task.next_run_at = now + timedelta(seconds=task.interval_seconds)
|
|
return task
|
|
|
|
def cancel(self, task_id: str) -> None:
|
|
self.tasks[task_id].cancelled = True
|
|
|
|
def create_successor_once(self, predecessor: str, successor: str, *, interval_seconds: int, now: datetime | None = None) -> RecurringTask:
|
|
key = f"{predecessor}->{successor}"
|
|
if key in self.successors:
|
|
return self.tasks[successor]
|
|
self.successors.add(key)
|
|
if successor not in self.tasks:
|
|
return self.schedule(successor, interval_seconds, now=now)
|
|
return self.tasks[successor]
|
|
|