canic-core 0.93.27

Canic — a canister orchestration and management toolkit for the Internet Computer
Documentation
//! Module: workflow::pool::scheduler
//!
//! Responsibility: schedule and execute background pool reset batches.
//! Does not own: pool admission policy, admin APIs, or stable pool schemas.
//! Boundary: workflow runtime helper coordinating timers, admission checks, reset, and metrics.

use crate::{
    InternalError,
    domain::policy::pure::pool::PoolPolicyError,
    log,
    log::Topic,
    ops::{
        runtime::{
            metrics::{
                pool::{
                    PoolMetricOperation as MetricOperation, PoolMetricOutcome as MetricOutcome,
                    PoolMetricReason as MetricReason,
                },
                recording::PoolMetricEvent as MetricEvent,
            },
            timer::TimerId,
        },
        storage::pool::PoolOps,
    },
    view::pool::PoolPendingResetCursor,
    workflow::{
        config::{WORKFLOW_INIT_DELAY, WORKFLOW_POOL_CHECK_INTERVAL},
        pool::{PoolWorkflow, admissibility::check_can_enter_pool, metric_reason_for_policy},
        runtime::timer::TimerWorkflow,
    },
};
use std::{
    cell::{Cell, RefCell},
    time::Duration,
};

/// Default batch size for resetting pending pool entries.
pub const POOL_RESET_BATCH_SIZE: usize = 10;

// -----------------------------------------------------------------------------
// Timer State
// -----------------------------------------------------------------------------

thread_local! {
    static TIMER: RefCell<Option<TimerId>> = const { RefCell::new(None) };
}

// -----------------------------------------------------------------------------
// Internal Scheduler State
// -----------------------------------------------------------------------------

thread_local! {
    static RESET_IN_PROGRESS: RefCell<bool> = const { RefCell::new(false) };
    static RESET_RESCHEDULE: RefCell<bool> = const { RefCell::new(false) };
    static RESET_TIMER: RefCell<Option<TimerId>> = const { RefCell::new(None) };
    static RESET_CURSOR: Cell<Option<PoolPendingResetCursor>> = const { Cell::new(None) };
}

///
/// PoolSchedulerWorkflow
///

pub struct PoolSchedulerWorkflow;

impl PoolSchedulerWorkflow {
    /// Start pool background scheduling.
    ///
    /// Safe to call multiple times.
    pub fn start() {
        let _ = TimerWorkflow::set_guarded_interval(
            &TIMER,
            WORKFLOW_INIT_DELAY,
            "pool:init",
            || async {
                Self::schedule_if_pending();
            },
            WORKFLOW_POOL_CHECK_INTERVAL,
            "pool:interval",
            || async {
                Self::schedule_if_pending();
            },
        );
    }

    /// Schedule a reset worker only when pending reset work exists.
    pub fn schedule_if_pending() {
        if Self::has_pending_reset() {
            Self::schedule();
        } else {
            MetricEvent::skipped(MetricOperation::Scheduler, MetricReason::Empty);
        }
    }

    /// Schedule a reset worker run.
    ///
    /// This is idempotent and guarded against concurrent execution.
    pub fn schedule() {
        MetricEvent::record(
            MetricOperation::Scheduler,
            MetricOutcome::Scheduled,
            MetricReason::Ok,
        );
        let _ = TimerWorkflow::set_guarded(&RESET_TIMER, Duration::ZERO, "pool:pending", async {
            RESET_TIMER.with_borrow_mut(|slot| *slot = None);
            let _ = Self::run_worker(POOL_RESET_BATCH_SIZE).await;
        });
    }

    fn maybe_reschedule(has_more: bool) {
        let reschedule = RESET_RESCHEDULE.with_borrow_mut(|flag| {
            let v = *flag;
            *flag = false;
            v
        });

        if should_reschedule(reschedule, has_more) {
            Self::schedule();
        }
    }

    async fn run_worker(limit: usize) -> Result<(), InternalError> {
        if limit == 0 {
            MetricEvent::skipped(MetricOperation::Scheduler, MetricReason::Empty);
            return Ok(());
        }

        let should_run = RESET_IN_PROGRESS.with_borrow_mut(|flag| {
            if *flag {
                RESET_RESCHEDULE.with_borrow_mut(|r| *r = true);
                false
            } else {
                *flag = true;
                true
            }
        });

        if !should_run {
            MetricEvent::skipped(MetricOperation::Scheduler, MetricReason::InProgress);
            return Ok(());
        }

        MetricEvent::started(MetricOperation::Scheduler);
        let result = Self::run_batch(limit).await;

        RESET_IN_PROGRESS.with_borrow_mut(|flag| *flag = false);

        let next_cursor = result.as_ref().ok().copied().flatten();
        RESET_CURSOR.set(next_cursor);
        Self::maybe_reschedule(next_cursor.is_some());

        match &result {
            Ok(_) => MetricEvent::completed(MetricOperation::Scheduler, MetricReason::Ok),
            Err(err) => MetricEvent::failed(MetricOperation::Scheduler, err),
        }

        result.map(|_| ())
    }

    async fn run_batch(limit: usize) -> Result<Option<PoolPendingResetCursor>, InternalError> {
        let after = RESET_CURSOR.get();
        let page = PoolOps::pending_reset_page(after.as_ref(), limit);

        for pid in page.pids {
            match check_can_enter_pool(pid).await {
                Ok(()) => {}

                Err(PoolPolicyError::RegisteredInSubnet(_)) => {
                    if let Err(err) = PoolWorkflow::abort_pending_pool_import_intent(pid) {
                        log!(
                            Topic::CanisterPool,
                            Warn,
                            "pool reset rejection could not abort import intent for {pid}: {err}"
                        );
                        return Err(err);
                    }
                    PoolOps::remove(&pid);
                    MetricEvent::skipped(MetricOperation::Reset, MetricReason::RegisteredInSubnet);
                    continue;
                }

                Err(PoolPolicyError::NonImportableOnLocal { .. }) => {
                    // The authoritative pending record remains queued.
                    MetricEvent::record(
                        MetricOperation::Reset,
                        MetricOutcome::Requeued,
                        MetricReason::NonImportableLocal,
                    );
                    continue;
                }

                Err(err) => {
                    MetricEvent::record(
                        MetricOperation::Reset,
                        MetricOutcome::Requeued,
                        metric_reason_for_policy(&err),
                    );
                    continue;
                }
            }

            match PoolWorkflow::reset_into_pool(pid).await {
                Ok(cycles) => {
                    PoolWorkflow::commit_pending_pool_import_intent(pid)?;
                    PoolWorkflow::mark_ready(pid, cycles);
                }
                Err(err) => {
                    log!(
                        Topic::CanisterPool,
                        Warn,
                        "pool reset failed for {pid}: {err}"
                    );
                    if let Err(abort_err) = PoolWorkflow::abort_pending_pool_import_intent(pid) {
                        return Err(err.with_diagnostic_context(format!(
                            "pool import intent abort failed for {pid}: {abort_err}"
                        )));
                    }
                    PoolWorkflow::mark_failed(pid, &err);
                }
            }
        }

        Ok(page.next_cursor)
    }

    fn has_pending_reset() -> bool {
        PoolOps::has_pending_reset()
    }
}

#[must_use]
const fn should_reschedule(requested: bool, has_more: bool) -> bool {
    requested || has_more
}

#[cfg(test)]
mod tests {
    use super::should_reschedule;

    #[test]
    fn reset_worker_only_reschedules_for_more_sweep_work_or_a_new_request() {
        assert!(!should_reschedule(false, false));
        assert!(should_reschedule(false, true));
        assert!(should_reschedule(true, false));
        assert!(should_reschedule(true, true));
    }
}