wowlab_sentinel/scheduler/

assignment.rs

//! Application workflow for one scheduler assignment round.

use uuid::Uuid;
use wowlab_types::sim::FastSet;

use super::{
    backlog::NodeBacklogs,
    dispatch::{
        AssignmentClock, AssignmentPublisher, ClaimTokenSource, DispatchDependencies, Dispatcher,
        StateAssignmentPublisher, SystemClock, UuidClaimTokenSource,
    },
    eligibility::EligibilityContext,
    planning::{JobPlan, build_runtime, order_for_dispatch, plan_jobs},
    repository::{AssignmentRepository, RepositoryError, SqlAssignmentRepository},
};
use crate::{state::RuntimeState, utils::filter_refresh::FilterMap};

#[derive(Debug, thiserror::Error)]
pub(super) enum AssignmentError {
    #[error(transparent)]
    Repository(#[from] RepositoryError),
}

#[derive(Debug, Default, Eq, PartialEq)]
pub(super) struct AssignmentRound {
    pub jobs: usize,
    pub planned: usize,
    pub started: usize,
    pub assigned: usize,
    pub invalid_configs: usize,
    pub no_eligible_node: usize,
    pub no_work: usize,
    pub rejected_publications: usize,
    pub no_online_nodes: bool,
}

pub(super) struct AssignmentWorkflow<'a, R, P, C, T> {
    runtime: &'a RuntimeState,
    filters: &'a FilterMap,
    repository: R,
    publisher: P,
    clock: C,
    claim_tokens: T,
}

#[cfg(test)]
pub(super) struct AssignmentDependencies<R, P, C, T> {
    pub repository: R,
    pub publisher: P,
    pub clock: C,
    pub claim_tokens: T,
}

pub(super) fn assignment_workflow(
    state: &crate::state::ServerState,
) -> AssignmentWorkflow<
    '_,
    SqlAssignmentRepository<'_>,
    StateAssignmentPublisher<'_>,
    SystemClock,
    UuidClaimTokenSource,
> {
    AssignmentWorkflow {
        runtime: &state.runtime,
        filters: &state.filters,
        repository: SqlAssignmentRepository::new(
            state.dbs.get::<super::SchedulerDb>(),
            state.config.scheduler_batch_size,
        ),
        publisher: StateAssignmentPublisher::new(state),
        clock: SystemClock,
        claim_tokens: UuidClaimTokenSource,
    }
}

impl<R, P, C, T> AssignmentWorkflow<'_, R, P, C, T>
where
    R: AssignmentRepository,
    P: AssignmentPublisher,
    C: AssignmentClock,
    T: ClaimTokenSource,
{
    #[cfg(test)]
    pub(super) fn new<'a>(
        runtime: &'a RuntimeState,
        filters: &'a FilterMap,
        dependencies: AssignmentDependencies<R, P, C, T>,
    ) -> AssignmentWorkflow<'a, R, P, C, T> {
        AssignmentWorkflow {
            runtime,
            filters,
            repository: dependencies.repository,
            publisher: dependencies.publisher,
            clock: dependencies.clock,
            claim_tokens: dependencies.claim_tokens,
        }
    }

    pub(super) async fn run(&self) -> Result<AssignmentRound, AssignmentError> {
        let jobs = self.repository.pending_jobs().await?;

        if jobs.is_empty() {
            tracing::debug!("No pending jobs found");
            metrics::gauge!(crate::telemetry::CHUNKS_PENDING).set(0.0);

            return Ok(AssignmentRound::default());
        }

        let pending: Vec<_> = jobs
            .iter()
            .map(|job| (job.id, job.status.as_str()))
            .collect();

        tracing::info!(?pending, "Pending jobs");
        let job_count = jobs.len();
        let mut plan = plan_jobs(jobs);
        let started = self.initialize_runtimes(&plan).await?;

        log_plan(&plan, started);

        let nodes = self.repository.online_nodes().await?;

        if nodes.is_empty() {
            tracing::info!("No online nodes available for assignment");

            return Ok(AssignmentRound {
                jobs: job_count,
                planned: plan.jobs.len(),
                started,
                invalid_configs: plan.invalid_configs,
                no_online_nodes: true,
                ..AssignmentRound::default()
            });
        }

        let public_keys = nodes
            .iter()
            .map(|node| node.public_key.clone())
            .collect::<Vec<_>>();
        let permissions = self.repository.permissions(&public_keys).await?;
        let mut backlogs = {
            let runtimes = self.runtime.jobs.read().await;

            NodeBacklogs::from_runtimes(&runtimes)
        };

        tracing::debug!(?backlogs, "Node backlogs");
        let user_ids = unique_user_ids(&plan);
        let user_discord_ids = self.repository.user_discord_ids(&user_ids).await?;
        let friend_memberships = self.repository.friend_memberships(&user_ids).await?;

        order_for_dispatch(&mut plan.jobs);
        let eligibility = EligibilityContext {
            permissions: &permissions,
            filters: self.filters,
            friend_memberships: &friend_memberships,
            user_discord_ids: &user_discord_ids,
        };
        let dispatch = Dispatcher::new(
            self.runtime,
            DispatchDependencies {
                publisher: &self.publisher,
                clock: &self.clock,
                claim_tokens: &self.claim_tokens,
            },
        )
        .dispatch(&plan.jobs, &nodes, &eligibility, &mut backlogs)
        .await;

        log_dispatch(&dispatch, plan.invalid_configs);
        record_assignment_metrics(dispatch.assigned);
        tracing::info!(
            assigned = dispatch.assigned,
            jobs = job_count,
            "Assignment round complete"
        );

        Ok(AssignmentRound {
            jobs: job_count,
            planned: plan.jobs.len(),
            started,
            assigned: dispatch.assigned,
            invalid_configs: plan.invalid_configs,
            no_eligible_node: dispatch.no_eligible_node,
            no_work: dispatch.no_work,
            rejected_publications: dispatch.rejected_publications,
            no_online_nodes: false,
        })
    }

    async fn initialize_runtimes(&self, plan: &JobPlan) -> Result<usize, AssignmentError> {
        let mut started = 0usize;

        for planned in &plan.jobs {
            let already_tracked = self.runtime.jobs.read().await.contains(&planned.job.id);

            if already_tracked {
                continue;
            }

            self.runtime
                .jobs
                .write()
                .await
                .insert(build_runtime(&planned.job, &planned.config));
            let rows = self.repository.mark_running(planned.job.id).await?;

            started += usize::try_from(rows).unwrap_or(usize::MAX);
        }

        Ok(started)
    }
}

impl AssignmentError {
    pub(super) const fn is_pending_fetch(&self) -> bool {
        matches!(self, Self::Repository(RepositoryError::FetchPending(_)))
    }
}

#[async_trait::async_trait]
impl<T> AssignmentPublisher for &T
where
    T: AssignmentPublisher + ?Sized,
{
    async fn job_running(&self, job_id: Uuid) -> super::dispatch::PublishOutcome {
        (**self).job_running(job_id).await
    }

    async fn chunk(
        &self,
        public_key: &wowlab_common::NodePublicKey,
        payload: &wowlab_common::RuntimeChunkPayload,
    ) -> super::dispatch::PublishOutcome {
        (**self).chunk(public_key, payload).await
    }
}

impl<T> AssignmentClock for &T
where
    T: AssignmentClock + ?Sized,
{
    fn now_ms(&self) -> u64 {
        (**self).now_ms()
    }
}

impl<T> ClaimTokenSource for &T
where
    T: ClaimTokenSource + ?Sized,
{
    fn next(&self) -> wowlab_common::ClaimToken {
        (**self).next()
    }
}

fn unique_user_ids(plan: &JobPlan) -> Vec<Uuid> {
    let unique: FastSet<Uuid> = plan
        .jobs
        .iter()
        .map(|planned| planned.job.user_id)
        .collect();

    unique.into_iter().collect()
}

fn log_plan(plan: &JobPlan, started: usize) {
    if plan.invalid_configs > 0 {
        tracing::error!(
            invalid_configs = plan.invalid_configs,
            "Rejected jobs with invalid sentinel config"
        );
    }

    tracing::info!(
        planned = plan.jobs.len(),
        started_jobs = started,
        "Planned assignment round"
    );
}

fn log_dispatch(report: &super::dispatch::DispatchReport, invalid_configs: usize) {
    tracing::info!(
        assigned_count = report.assigned,
        no_eligible_node = report.no_eligible_node,
        no_work = report.no_work,
        invalid_configs,
        rejected_publications = report.rejected_publications,
        "Assignment round distribution"
    );
}

fn record_assignment_metrics(assigned: usize) {
    if assigned == 0 {
        return;
    }

    metrics::counter!(crate::telemetry::CHUNKS_ASSIGNED).increment(assigned as u64);
    #[expect(
        clippy::cast_precision_loss,
        reason = "assignment counts remain exactly representable by operational metrics"
    )]
    metrics::gauge!(crate::telemetry::CHUNKS_RUNNING).increment(assigned as f64);
}

#[cfg(test)]
mod tests;