internal/scheduler

import "github.com/neochaotic/leoflow/internal/scheduler"

const LockID int64 = 0x4C656F466C6F77

func CanTransition(from, to domain.TaskState) bool

func CanTransitionDagRun(from, to domain.DagRunState) bool

func FinalizeRun(run RunState) (domain.DagRunState, bool)

func PoolKey(tenant, pool string) string

type Alerter interface {
    // AlertRunFailed sends every on_failure rule and reports whether all of them
    // were delivered. The caller stamps delivery only on true, so a failed send
    // stays retryable instead of being marked and lost.
    AlertRunFailed(ctx context.Context, run RunState) bool
}

type Dispatcher interface {
    Dispatch(ctx context.Context, runID, dagID, dagVersionID string, task domain.TaskSpec) (executor.Disposition, error)
}

type ExecutionReaper interface {
    // ReapOnce runs every reaper once. Implementations isolate per-reaper and
    // per-candidate failures internally and log/meter list errors, so the
    // scheduler ignores the return today; the error is kept for the seam.
    ReapOnce(ctx context.Context) error
}

type Leader struct {
    // contains filtered or unexported fields
}

func NewLeader(pool *pgxpool.Pool) *Leader

func (l *Leader) HoldsLock(ctx context.Context) (bool, error)

func (l *Leader) Release(ctx context.Context) error

func (l *Leader) TryAcquire(ctx context.Context) (bool, error)

type LeaderHealthReader struct {
    // contains filtered or unexported fields
}

func NewLeaderHealthReader(pool *pgxpool.Pool) *LeaderHealthReader

func (r *LeaderHealthReader) Heartbeat() (healthy bool, last time.Time)

type PlannedTransition struct {
    TaskID string
    To     domain.TaskState
}

func PlanRun(run RunState) []PlannedTransition

type Recorder interface {
    RecordSchedulerDecision(decisionType string)
    RecordTaskTransition(from, to, dagID string)
    // RecordUndispatchable counts tasks queued with no executor to launch them.
    RecordUndispatchable(reason string)
    // RecordSchedulerStepDown counts a leader step-down by reason — the rate
    // of this counter is the operator-facing SLI for leader-churn (#311). A
    // nil Recorder is tolerated by the scheduler so tests need not stub it.
    RecordSchedulerStepDown(reason string)
    // ObserveSchedulerReacquire records the wall-clock duration of a step-down
    // → re-acquire cycle (#311). Operators alert on P99 to spot churn that
    // starts to delay scheduling latency.
    ObserveSchedulerReacquire(d time.Duration)
    // RecordAlert counts one on-failure alert outcome by dag, channel type, and
    // result. The scheduler records result="dropped" when its dispatch semaphore
    // is saturated (#435); the dispatcher records "sent"/"failed" for deliveries.
    RecordAlert(dagID, channelType, result string)
}

type RunState struct {
    // RunID is the dag_runs primary key (a UUID). It is what dispatch, the agent
    // token and the log sink key on, so it must not be swapped for the
    // user-facing id.
    RunID string
    // DisplayRunID is the dag_runs.run_id column — the identifier an operator
    // sees in the UI and passes to the API ("manual__2026-07-30T12:00:00+00:00").
    // Alerts render this one: a UUID is not something anyone can act on.
    DisplayRunID string
    // LogicalDate is the run's logical date in RFC3339, empty for an unscheduled
    // run. Carried here so an alert can say which interval failed without a
    // second query.
    LogicalDate string
    DagID       string
    // DagVersionID is the dag_versions row this run pins (dag_runs.dag_version_id).
    // It names the warm-worker pool a task is placed onto (ADR 0058 N1b1-place):
    // dispatch threads it to the placer so an attempt reuses a worker of the exact
    // version it compiled against. Empty for runs loaded by paths that do not set
    // it; the dedicated pod path ignores it.
    DagVersionID string
    TenantID     string
    State        domain.DagRunState
    Tasks        []domain.TaskSpec
    States       map[string]domain.TaskState
    // Tries and MaxTries hold the current and maximum attempt counts per task,
    // driving retry decisions. Absent entries mean no retry budget.
    Tries    map[string]int
    MaxTries map[string]int
    // EndedAt holds the failure timestamp per task (only meaningful for tasks
    // currently in `up_for_retry`). Combined with RetryDelaySeconds + Now, it
    // gates the `up_for_retry → none` transition so retries honor user-
    // declared backoff (issue #201). Absent entries mean no cooldown applies.
    EndedAt           map[string]*time.Time
    RetryDelaySeconds map[string]int
    // RescheduleAt holds the next-poke time per task (only meaningful for tasks in
    // up_for_reschedule). The planner re-dispatches such a task once Now >=
    // reschedule_at — without consuming retry budget (#380). Absent/zero entries
    // re-dispatch immediately (preserves the test seam, like EndedAt).
    RescheduleAt map[string]*time.Time
    // NextDispatchAt holds the earliest time a `scheduled` task may be dispatched
    // again after a synchronous dispatch failure (ADR 0031 Amendment A). The
    // planner does not promote scheduled→queued while Now < next_dispatch_at.
    // Absent/zero entries dispatch immediately (the common case: never failed).
    NextDispatchAt map[string]*time.Time
    // DispatchAttempts counts consecutive synchronous dispatch failures per task.
    // It is a separate counter from try_number — a dispatch failure is infra, not
    // a task failure — and drives the dispatch_failed give-up at
    // dispatchMaxAttempts. Absent entries mean zero.
    DispatchAttempts map[string]int
    // InfraFailed marks a `failed` task whose failure was infra-caused (agent/pod/
    // dispatch lost), from the durable last_failure_kind signal. Such a task
    // re-places without consuming the task's retry budget (ADR 0051 Phase 1) — an
    // infrastructure fault is not the user's task failing.
    InfraFailed map[string]bool
    // InfraAttempts counts asynchronous infra re-placements per task — the
    // try_number-free analog of DispatchAttempts for agent/pod/dispatch-lost faults.
    // It bounds the re-place at infraMaxAttempts so a poison placement cannot loop
    // forever. Absent entries mean zero.
    InfraAttempts map[string]int
    // Now is the wall-clock value the planner compares against EndedAt. Zero
    // means "skip the cooldown gate" so legacy callers + tests that don't set
    // retry_delay get the previous (immediate-retry) behavior.
    Now time.Time
    // Alerts carries the DAG's native on-failure alerting config (#424), loaded
    // from the dag.json alongside Tasks. nil means no alerting; the scheduler
    // dispatches these rules once the run finalizes as failed.
    Alerts *domain.AlertsConfig
    // MaxActiveTasks is the DAG's per-DAG task-concurrency cap (Airflow's
    // max_active_tasks, ADR 0053 Stage 1), loaded from the spec. Zero or negative
    // means unlimited — the admission gate is a no-op, so a DAG that never sets it
    // (and all of Lite) plans exactly as before.
    MaxActiveTasks int
    // ActiveTaskCount is the DAG-wide number of currently non-terminal (queued or
    // running) task instances across every active run of this DAG, plus any this
    // tick already admitted for the same DAG in earlier sibling runs. PlanRun
    // subtracts it from MaxActiveTasks to derive this run's admission headroom, so
    // the cap holds per-DAG across runs and one tick cannot breach it. The Step
    // loop populates it; it defaults to zero, leaving the gate unlimited when
    // MaxActiveTasks is unset.
    ActiveTaskCount int
    // PoolsEnabled turns on the named-pool slot gate (ADR 0053 Stage 3, Pro only).
    // The Step loop sets it from the scheduler's edition flag. Lite and any
    // non-Pro deployment leave it false, so the pool gate is a no-op and planning
    // is byte-identical to the max_active_tasks-only path.
    PoolsEnabled bool
    // PoolBudgets is the per-pool slot cap keyed by PoolKey(TenantID, pool). A pool
    // with a non-positive or absent budget is unlimited (fail open, never
    // deadlock). The Step loop sets it from the once-per-tick PoolBudgets snapshot;
    // nil in Lite. Shared read-only across sibling runs.
    PoolBudgets map[string]int
    // PoolActive is the cross-DAG count of currently non-terminal (queued+running)
    // task instances per pool, keyed like PoolBudgets, plus any this tick already
    // admitted into the same pool by earlier runs. PlanRun adds its own within-call
    // promotions on top. The Step loop sets it and folds each run's admissions back
    // in so a single tick cannot breach a pool across runs; nil in Lite.
    PoolActive map[string]int
}

type ScheduledDAG struct {
    DagID       string
    Schedule    string
    LastLogical *time.Time
    StartDate   *time.Time
    Catchup     bool
    // MaxActiveRuns caps how many runs of this DAG may be active (queued or
    // running) at once. Zero means "unlimited" (Airflow-faithful: a missing
    // limit is unbounded). The scheduler enforces the cap in createDueRuns:
    // once the active-run count reaches this value, additional due slots are
    // skipped on this tick (issue #200).
    MaxActiveRuns int
}

type Scheduler struct {
    // contains filtered or unexported fields
}

func NewScheduler(store Store, logger *slog.Logger, interval time.Duration) *Scheduler

func (s *Scheduler) ClearSteppingDown()

func (s *Scheduler) EnablePools()

func (s *Scheduler) Heartbeat() (bool, time.Time)

func (s *Scheduler) IsLeading() bool

func (s *Scheduler) MarkSteppingDown(reason string)

func (s *Scheduler) RecordReacquireSince(stepDownAt time.Time)

func (s *Scheduler) Run(ctx context.Context) error

func (s *Scheduler) SetAlertConcurrency(n int)

func (s *Scheduler) SetAlerter(a Alerter)

func (s *Scheduler) SetDispatcher(d Dispatcher)

func (s *Scheduler) SetExecutionReaper(r ExecutionReaper)

func (s *Scheduler) SetLeading(on bool)

func (s *Scheduler) SetRecorder(r Recorder)

func (s *Scheduler) SetStepTimeout(d time.Duration)

func (s *Scheduler) Step(ctx context.Context) error

func (s *Scheduler) SteppingDown() bool

type Store interface {
    ActiveRuns(ctx context.Context) ([]RunState, error)
    // PoolBudgets returns every named pool's slot cap keyed by PoolKey(tenant,
    // pool) — the cross-DAG admission budget the pool gate enforces (ADR 0053
    // Stage 3). Called once per tick, and ONLY on the Pro path (poolsEnabled);
    // Lite never invokes it.
    PoolBudgets(ctx context.Context) (map[string]int, error)
    MaterializeTasks(ctx context.Context, runID string, tasks []domain.TaskSpec) error
    ApplyTransition(ctx context.Context, runID, taskID string, to domain.TaskState) error
    // ApplyTransitions moves every listed task of a run to the SAME target state
    // in one statement — the batched form of calling ApplyTransition once per
    // task. The scheduler groups a tick's plain state-set transitions
    // (scheduled/skipped/upstream_failed/up_for_retry) by target state and flushes
    // each group here, so R updates collapse to one per distinct state with a
    // byte-identical result. It carries no source-state guard, matching
    // ApplyTransition; the guarded none-rail resets stay on their own methods.
    ApplyTransitions(ctx context.Context, runID string, taskIDs []string, to domain.TaskState) error
    // ResetForRetry returns a task to 'none' and increments its try number so a
    // retry re-evaluates and re-runs it. It is guarded to the up_for_retry source
    // state; the bool reports whether the reset actually fired (false = the TI was
    // no longer up_for_retry, nothing reset) so the caller does not record a
    // phantom retry on a lost race.
    ResetForRetry(ctx context.Context, runID, taskID string) (bool, error)
    // ResetForInfraReplace returns a task a reaper failed as infra to 'none' for a
    // re-run, PRESERVING try_number and bumping infra_attempts instead — an infra
    // fault must not consume the retry budget (ADR 0051 Phase 1). Guarded to the
    // failed+infra source; the bool reports whether the reset fired (false = no
    // longer a failed-infra candidate) so the caller does not record a phantom
    // re-placement on a lost race.
    ResetForInfraReplace(ctx context.Context, runID, taskID string) (bool, error)
    // RedispatchReschedule returns a task parked in up_for_reschedule to 'none' for
    // re-dispatch, PRESERVING try_number (reschedule is not a retry; #380).
    RedispatchReschedule(ctx context.Context, runID, taskID string) error
    // RecordDispatchFailure backs off a scheduled task after a synchronous dispatch
    // failure: it increments dispatch_attempts and sets next_dispatch_at so the
    // planner does not re-attempt until the backoff elapses (ADR 0031 Amendment A).
    RecordDispatchFailure(ctx context.Context, runID, taskID string, nextAt time.Time) error
    // RecordDispatchBackpressure backs off a scheduled task after a retriable-
    // forever cluster-backpressure dispatch failure (quota 403 / APF 429): it sets
    // next_dispatch_at WITHOUT incrementing dispatch_attempts, so the task is
    // re-offered once the backoff elapses but never accumulates toward the
    // dispatch_failed cap (ADR 0053).
    RecordDispatchBackpressure(ctx context.Context, runID, taskID string, nextAt time.Time) error
    // FailDispatchExhausted fails a scheduled task as dispatch_failed once its
    // dispatch-attempt budget is spent, so the run finalizes instead of looping.
    FailDispatchExhausted(ctx context.Context, runID, taskID, reason string) error
    SetRunState(ctx context.Context, runID string, state domain.DagRunState) error
    // ClaimAlertAttempt atomically claims ONE on-failure send attempt, reporting
    // true iff this call won it. It refuses when the episode was already
    // delivered, when the attempt budget is spent, or when the backoff has not
    // elapsed — so a dead endpoint stops being retried instead of being hit once
    // per tick for the life of the run (#431, and the delivery split below).
    // It returns the attempt number won, or 0 when the claim was refused.
    ClaimAlertAttempt(ctx context.Context, runID string, maxAttempts int, backoff time.Duration) (int, error)
    // MarkRunAlertDelivered stamps the episode as paged, for the attempt the
    // caller won. Claiming and stamping were once the same operation, which meant
    // every failed send was a page lost with no retry. The attempt is part of the
    // call because the send is detached from the tick: an operator clear can start
    // a new episode mid-send, and a stamp from the superseded one must not land.
    MarkRunAlertDelivered(ctx context.Context, runID string, attempt int) error
    ScheduledDAGs(ctx context.Context) ([]ScheduledDAG, error)
    CreateScheduledRun(ctx context.Context, dagID string, logical time.Time) error
    // SetTaskNote attaches operational context to a task instance (shown in the
    // UI), e.g. why it is queued but not running.
    SetTaskNote(ctx context.Context, runID, taskID, note string) error
}

type TriggerDecision int

const (
    // DecisionWait means the upstreams are not yet settled; re-evaluate later.
    DecisionWait TriggerDecision = iota
    // DecisionSchedule means dependencies are satisfied; move the task to scheduled.
    DecisionSchedule
    // DecisionSkip means the trigger rule can no longer be satisfied; skip the task.
    DecisionSkip
    // DecisionUpstreamFailed means a required upstream failed; propagate the failure.
    DecisionUpstreamFailed
)

func EvaluateTriggerRule(rule domain.TriggerRule, upstreams []domain.TaskState) TriggerDecision