aria2-rust-pro/crates/aria2-rust-pro-core/src/engine/scheduling.rs

use std::time::{SystemTime, UNIX_EPOCH};

use super::{
    CoreError, DownloadEngine, DownloadId, DownloadStatus, RequestGroup, Result, RetryAttempt,
    RetryHistoryEntry, RuntimeConfig, RuntimeEvent, RuntimeEventKind, ScheduleDecision,
    SegmentAssignment, SegmentState, effective_piece_length, infer_total_length, usize_to_u64,
};

impl DownloadEngine {
    /// Advances the scheduler clock and emits a scheduler tick event.
    pub fn scheduler_tick(&mut self) -> Result<()> {
        let _ = self.scheduler.tick();
        self.events
            .emit(RuntimeEvent::new(RuntimeEventKind::SchedulerTick));
        Ok(())
    }

    /// Prepares a single HTTP download by applying the scheduler decision for its current state.
    pub fn prepare_http_download(&mut self, gid: DownloadId) -> Result<RequestGroup> {
        let decision = {
            let group = self
                .registry
                .get(gid)
                .ok_or(CoreError::UnknownDownloadId(gid))?;
            match group.status() {
                DownloadStatus::Waiting => ScheduleDecision::Queue(gid),
                DownloadStatus::Active => ScheduleDecision::RunNow(gid),
                DownloadStatus::Error => ScheduleDecision::RetryLater(gid),
                DownloadStatus::Paused | DownloadStatus::Complete | DownloadStatus::Removed => {
                    ScheduleDecision::Noop
                }
            }
        };

        self.apply_schedule_decision(gid, &decision);
        self.registry
            .get(gid)
            .cloned()
            .ok_or(CoreError::UnknownDownloadId(gid))
    }

    /// Runs one scheduler pass and returns the first actionable decision.
    #[must_use]
    pub fn schedule_once(&mut self) -> ScheduleDecision {
        self.scheduler.record_schedule_run();
        let _ = self.scheduler.tick();
        let mut gids = self
            .registry
            .groups
            .iter()
            .filter_map(|(gid, group)| (group.status() == &DownloadStatus::Active).then_some(*gid))
            .collect::<Vec<_>>();
        gids.sort_by_key(|gid| gid.as_u64());
        gids.extend(
            self.reserved_queue
                .iter()
                .copied()
                .filter(|gid| self.registry.get(*gid).is_some()),
        );
        let mut retry_gids = self
            .registry
            .groups
            .iter()
            .filter_map(|(gid, group)| (group.status() == &DownloadStatus::Error).then_some(*gid))
            .collect::<Vec<_>>();
        retry_gids.sort_by_key(|gid| gid.as_u64());
        gids.extend(retry_gids);

        for gid in gids {
            let Some(group) = self.registry.get(gid) else {
                continue;
            };
            let decision = self.scheduler.decide(group);
            match decision {
                ScheduleDecision::RunNow(_)
                | ScheduleDecision::Queue(_)
                | ScheduleDecision::RetryLater(_) => {
                    self.apply_schedule_decision(gid, &decision);
                    self.events
                        .emit(RuntimeEvent::new(RuntimeEventKind::SchedulerTick).with_gid(gid));
                    return decision;
                }
                ScheduleDecision::Pause(_) | ScheduleDecision::Remove(_) => {
                    self.scheduler.record_decision(&decision);
                    return decision;
                }
                ScheduleDecision::Noop => {}
            }
        }

        self.scheduler.record_decision(&ScheduleDecision::Noop);
        ScheduleDecision::Noop
    }
}

/// Converts retry attempts into scheduler-facing retry-history entries.
pub(super) fn retry_history_from_group(attempts: &[RetryAttempt]) -> Vec<RetryHistoryEntry> {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_or(0, |duration| duration.as_secs());
    attempts
        .iter()
        .map(|attempt| RetryHistoryEntry {
            at_unix_secs: now,
            reason: attempt.error.clone().unwrap_or_else(|| "retry".to_string()),
        })
        .collect()
}

/// Builds runtime segment assignments for the scheduler-selected active segment count.
pub(super) fn build_segment_assignments(
    group: &RequestGroup,
    runtime: &RuntimeConfig,
    desired_segments: usize,
) -> Vec<SegmentAssignment> {
    if desired_segments == 0 {
        return Vec::new();
    }

    let total_length = infer_total_length(group, effective_piece_length(group));
    let start_offset = group
        .resume_state()
        .map_or(0, |state| state.resume_offset)
        .max(group.completed_length())
        .min(total_length);
    let remaining = total_length.saturating_sub(start_offset);
    if remaining == 0 {
        return Vec::new();
    }

    let piece_length = effective_piece_length(group);
    let alignment = piece_length.max(runtime.min_split_size.max(1));
    let segment_count = desired_segments
        .min(
            usize::try_from(remaining.div_ceil(runtime.min_split_size.max(1)))
                .unwrap_or(usize::MAX),
        )
        .max(1);
    let target_span = remaining.div_ceil(usize_to_u64(segment_count));

    let mut cursor = start_offset;
    let mut assignments = Vec::with_capacity(segment_count);
    for slot in 0..segment_count {
        if cursor >= total_length {
            break;
        }

        let end = if slot + 1 == segment_count {
            total_length
        } else {
            let raw_end = cursor.saturating_add(target_span).min(total_length);
            let aligned_end = raw_end
                .div_ceil(alignment)
                .saturating_mul(alignment)
                .min(total_length);
            aligned_end.max(cursor.saturating_add(1))
        };

        let mut assignment =
            SegmentAssignment::new(slot, crate::piece::PieceRange::new(cursor, end));
        assignment.state = SegmentState::Active;
        assignments.push(assignment);
        cursor = end;
    }

    assignments
}