ax-task 0.8.1

//! Runtime-backed soft and explicitly hard kernel timer registration.

use crate::{
    runtime::{
        context::{
            RuntimeIrqGuard, runtime_current_cpu_mut, runtime_task_system, validate_task_context,
        },
        cpu::{CpuLocal, SchedulerDeadlineUpdate},
        task_runtime,
    },
    sched::system::DeadlineBaseGuardSource,
    thread::TaskError,
    time::{
        MonotonicDeadline,
        hard_timer::HardKernelTimerHandle,
        queue::{
            HardKernelTimerCallback, KernelTimerCallback, KernelTimerCancelOutcome,
            KernelTimerEntry, KernelTimerHandle, RestartableKernelTimerCallback, TaskDeadlineError,
        },
    },
};

enum KernelTimerRegistrationResult {
    Registered(KernelTimerHandle, Option<SchedulerDeadlineUpdate>),
    Rejected(TaskError, KernelTimerEntry),
}

struct KernelTimerCancellationResult {
    outcome: Result<KernelTimerCancelOutcome, TaskError>,
    removed: Option<KernelTimerEntry>,
}

/// Registers a task-context callback on the calling CPU's monotonic clock base.
///
/// Callback ownership is allocated before IRQs are excluded. Hard IRQ only
/// promotes the entry into the existing `ktimers/%u` service; the callback is
/// invoked later without the deadline lock or an IRQ guard held.
///
/// # Errors
///
/// Returns [`TaskError::UnsafeContext`] outside ordinary task context,
/// [`TaskError::TimerCapacity`] when the per-CPU callback base is full, or a
/// runtime/clockevent error without leaving a hidden registration behind.
pub fn register_kernel_timer(
    deadline: MonotonicDeadline,
    callback: KernelTimerCallback,
) -> Result<KernelTimerHandle, TaskError> {
    validate_task_context()?;
    let entry = KernelTimerEntry::new(deadline, callback).map_err(kernel_timer_error)?;
    register_kernel_timer_entry(entry)
}

/// Registers a stable callback that may rearm the same timer identity.
///
/// The callback runs in the owner CPU's `ktimers/%u` task and returns either
/// [`KernelTimerAction::Complete`] or an absolute deadline for the same entry.
/// Cancellation remains non-blocking; if it races with an executing callback,
/// the callback may finish but cannot rearm the cancelled registration.
pub fn register_restartable_kernel_timer(
    deadline: MonotonicDeadline,
    callback: RestartableKernelTimerCallback,
) -> Result<KernelTimerHandle, TaskError> {
    validate_task_context()?;
    let entry =
        KernelTimerEntry::new_restartable(deadline, callback).map_err(kernel_timer_error)?;
    register_kernel_timer_entry(entry)
}

/// Registers a stable callback with explicit hard-IRQ expiry semantics.
///
/// The callback capability carries the caller's proof that invocation is
/// bounded and hard-IRQ-safe. Completion or cancellation drops its payload in
/// task context; returning [`crate::time::hard_timer::HardKernelTimerAction::Rearm`] preserves the same
/// timer identity and physical clockevent owner.
pub fn register_hard_restartable_kernel_timer(
    deadline: MonotonicDeadline,
    callback: HardKernelTimerCallback,
) -> Result<HardKernelTimerHandle, TaskError> {
    validate_task_context()?;
    let entry =
        KernelTimerEntry::new_hard_restartable(deadline, callback).map_err(kernel_timer_error)?;
    register_kernel_timer_entry(entry).map(HardKernelTimerHandle::new)
}

/// Arms an inactive hard timer or requests its next arm during execution.
///
/// A request during execution takes precedence over the callback's return
/// action. An already queued timer must be disarmed before it can be armed.
/// A registration with accepted cancellation cannot be rearmed.
///
/// The registration identity and callback allocation are reused. A caller
/// that moves the consumer to another CPU must destroy the old registration
/// and create a new owner-local one rather than remotely programming a
/// physical comparator.
pub fn arm_hard_kernel_timer(
    handle: HardKernelTimerHandle,
    deadline: MonotonicDeadline,
) -> Result<(), TaskError> {
    validate_task_context()?;
    let update = {
        let mut irq = RuntimeIrqGuard::enter();
        let cpu = runtime_current_cpu_mut(&mut irq)?;
        if cpu.owner() != handle.owner() {
            return Err(TaskError::CpuOwnerMismatch {
                expected: handle.owner().as_u32(),
                actual: cpu.owner().as_u32(),
            });
        }
        let mut deadline_base = cpu
            .remote()
            .lock_deadline_activity(DeadlineBaseGuardSource::Registration);
        let non_timer = deadline_base.non_timer;
        if !deadline_base
            .kernel_timers
            .arm_hard(handle.into(), deadline)
        {
            return Err(TaskError::InvalidConfiguration);
        }
        match CpuLocal::update_scheduler_deadline_registration_publication_if_changed(
            &mut deadline_base,
            non_timer,
        ) {
            Ok(update) => update,
            Err(error) => {
                assert_eq!(
                    deadline_base.kernel_timers.disarm_hard(handle.into()),
                    Some(Some(deadline)),
                    "failed hard-timer arm publication must restore inactivity"
                );
                return Err(error);
            }
        }
    };
    if let Some(update) = update {
        task_runtime::publish_scheduler_deadline(update);
    }
    Ok(())
}

/// Disarms one stable hard timer without destroying its callback payload.
///
/// Remote disarm only changes the logical owner base. Any already programmed
/// edge remains conservative and is reconciled by that CPU's firing
/// transaction; this operation never writes another CPU's comparator.
pub fn disarm_hard_kernel_timer(handle: HardKernelTimerHandle) -> Result<(), TaskError> {
    validate_task_context()?;
    let update = {
        let mut irq = RuntimeIrqGuard::enter();
        let current = runtime_current_cpu_mut(&mut irq)?;
        let system = runtime_task_system()?;
        let remote = system
            .cpu_remote(handle.owner())
            .ok_or(TaskError::InvalidConfiguration)?;
        let local_owner = current.owner() == handle.owner();
        let mut deadline_base =
            remote.lock_deadline_activity(DeadlineBaseGuardSource::Registration);
        let non_timer = local_owner.then_some(deadline_base.non_timer);
        let transition = deadline_base
            .kernel_timers
            .disarm_hard(handle.into())
            .ok_or(TaskError::InvalidConfiguration)?;
        let (Some(non_timer), Some(previous_deadline)) = (non_timer, transition) else {
            return Ok(());
        };
        match CpuLocal::update_scheduler_deadline_registration_publication_if_changed(
            &mut deadline_base,
            non_timer,
        ) {
            Ok(update) => update,
            Err(error) => {
                assert!(
                    deadline_base
                        .kernel_timers
                        .arm_hard(handle.into(), previous_deadline),
                    "failed hard-timer disarm publication must restore the active entry"
                );
                return Err(error);
            }
        }
    };
    if let Some(update) = update {
        task_runtime::publish_scheduler_deadline(update);
    }
    Ok(())
}

fn register_kernel_timer_entry(entry: KernelTimerEntry) -> Result<KernelTimerHandle, TaskError> {
    let result = {
        let mut irq = RuntimeIrqGuard::enter();
        let cpu = runtime_current_cpu_mut(&mut irq)?;
        let owner = cpu.owner();
        let mut deadline_base = cpu
            .remote()
            .lock_deadline_activity(DeadlineBaseGuardSource::Registration);
        let non_timer = deadline_base.non_timer;
        let inserted = deadline_base.kernel_timers.insert(owner, entry);
        match inserted {
            Ok(handle) => {
                match CpuLocal::update_scheduler_deadline_registration_publication_if_changed(
                    &mut deadline_base,
                    non_timer,
                ) {
                    Ok(update) => KernelTimerRegistrationResult::Registered(handle, update),
                    Err(error) => {
                        let removed = deadline_base
                            .kernel_timers
                            .cancel(handle)
                            .1
                            .expect("failed timer publication must roll back its new entry");
                        KernelTimerRegistrationResult::Rejected(error, removed)
                    }
                }
            }
            Err(entry) => KernelTimerRegistrationResult::Rejected(TaskError::TimerCapacity, entry),
        }
    };
    finish_kernel_timer_registration(result)
}

/// Cancels a registration without waiting for a callback already claimed.
///
/// `CancellationDeferred` accepts destruction and suppresses callback restart,
/// but is not a callback-completion or payload-reclamation barrier.
///
/// A remote cancellation mutates only the original owner base. It may leave a
/// conservative stale hardware edge; only the owner CPU may reprogram its
/// physical comparator.
pub fn cancel_kernel_timer(
    handle: KernelTimerHandle,
) -> Result<KernelTimerCancelOutcome, TaskError> {
    validate_task_context()?;
    let system = runtime_task_system()?;
    let result = {
        let mut irq = RuntimeIrqGuard::enter();
        let current = runtime_current_cpu_mut(&mut irq)?;
        let remote = system
            .cpu_remote(handle.owner())
            .ok_or(TaskError::InvalidConfiguration)?;
        let local_owner = current.owner() == handle.owner();
        let mut deadline_base =
            remote.lock_deadline_activity(DeadlineBaseGuardSource::Registration);
        let non_timer = local_owner.then_some(deadline_base.non_timer);
        let (cancel_outcome, mut removed) = deadline_base.kernel_timers.cancel(handle);
        let outcome = if removed.is_some() {
            if let Some(non_timer) = non_timer {
                match CpuLocal::update_scheduler_deadline_registration_publication_if_changed(
                    &mut deadline_base,
                    non_timer,
                ) {
                    Ok(Some(update)) => {
                        drop(deadline_base);
                        task_runtime::publish_scheduler_deadline(update);
                        Ok(KernelTimerCancelOutcome::Cancelled)
                    }
                    Ok(None) => Ok(KernelTimerCancelOutcome::Cancelled),
                    Err(error) => {
                        deadline_base.kernel_timers.restore_cancelled(
                            removed
                                .take()
                                .expect("failed cancellation publication must restore its entry"),
                        );
                        Err(error)
                    }
                }
            } else {
                Ok(KernelTimerCancelOutcome::Cancelled)
            }
        } else {
            Ok(cancel_outcome)
        };
        KernelTimerCancellationResult { outcome, removed }
    };
    drop(result.removed);
    result.outcome
}

fn finish_kernel_timer_registration(
    result: KernelTimerRegistrationResult,
) -> Result<KernelTimerHandle, TaskError> {
    match result {
        KernelTimerRegistrationResult::Registered(handle, update) => {
            if let Some(update) = update {
                task_runtime::publish_scheduler_deadline(update);
            }
            Ok(handle)
        }
        KernelTimerRegistrationResult::Rejected(error, entry) => {
            drop(entry);
            Err(error)
        }
    }
}

fn kernel_timer_error(error: TaskDeadlineError) -> TaskError {
    match error {
        TaskDeadlineError::Capacity => TaskError::TimerCapacity,
        TaskDeadlineError::GenerationExhausted | TaskDeadlineError::KindMismatch => {
            TaskError::InvalidConfiguration
        }
    }
}