launcher/crates/odian-packs/src/progress.rs

//! Install progress, shaped like `InstallProgress` in the launcher's docs/CONTRACT.md,
//! throttled to about ten reports per second.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use serde::{Deserialize, Serialize};

/// Install stages. This crate reports `manifest`, `app`, `db`, `src` and `tiles`;
/// `restore` and `finalize` exist so the launcher can reuse the type for its own steps.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Stage {
    Manifest,
    App,
    Db,
    Src,
    Tiles,
    Restore,
    Finalize,
}

/// One progress report (`InstallProgress` on the wire).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Progress {
    pub stage: Stage,
    /// 1-based.
    pub stage_index: u32,
    pub stage_count: u32,
    pub done_bytes: u64,
    /// 0 when unknown.
    pub total_bytes: u64,
    pub done_items: u64,
    pub total_items: u64,
    pub bytes_per_sec: f64,
    pub message: Option<String>,
}

const INTERVAL: Duration = Duration::from_millis(100);
/// Time constant of the smoothed transfer rate.
const RATE_WINDOW_SECS: f64 = 2.0;

pub(crate) type ProgressFn = Arc<dyn Fn(Progress) + Send + Sync>;

struct State {
    stage: Stage,
    stage_index: u32,
    stage_count: u32,
    total_bytes: u64,
    total_items: u64,
    label: String,
    unit: &'static str,
    sub_total: u64,
    last_emit: Option<Instant>,
    rate_bytes: u64,
    rate_at: Instant,
    rate: f64,
}

/// Shared by every worker of one install. Counters are atomics; a report is built and
/// sent at most every [`INTERVAL`], except stage starts and ends, which always go out.
pub(crate) struct Reporter {
    callback: ProgressFn,
    done_bytes: AtomicU64,
    done_items: AtomicU64,
    sub_done: AtomicU64,
    transferred: AtomicU64,
    state: Mutex<State>,
}

impl Reporter {
    pub fn new(callback: ProgressFn, stage_count: u32) -> Self {
        Reporter {
            callback,
            done_bytes: AtomicU64::new(0),
            done_items: AtomicU64::new(0),
            sub_done: AtomicU64::new(0),
            transferred: AtomicU64::new(0),
            state: Mutex::new(State {
                stage: Stage::Manifest,
                stage_index: 1,
                stage_count,
                total_bytes: 0,
                total_items: 0,
                label: String::new(),
                unit: "",
                sub_total: 0,
                last_emit: None,
                rate_bytes: 0,
                rate_at: Instant::now(),
                rate: 0.0,
            }),
        }
    }

    pub fn begin_stage(&self, stage: Stage, stage_index: u32, total_bytes: u64, total_items: u64) {
        {
            let mut st = self.state.lock().unwrap_or_else(|p| p.into_inner());
            st.stage = stage;
            st.stage_index = stage_index;
            st.total_bytes = total_bytes;
            st.total_items = total_items;
            st.label.clear();
            st.unit = "";
            st.sub_total = 0;
            self.done_bytes.store(0, Ordering::Relaxed);
            self.done_items.store(0, Ordering::Relaxed);
            self.sub_done.store(0, Ordering::Relaxed);
        }
        self.emit(true);
    }

    /// The part of the stage the message describes ("tiles dark/ds", 1402 "blocks").
    pub fn begin_sub(&self, label: String, unit: &'static str, sub_total: u64) {
        {
            let mut st = self.state.lock().unwrap_or_else(|p| p.into_inner());
            st.label = label;
            st.unit = unit;
            st.sub_total = sub_total;
            self.sub_done.store(0, Ordering::Relaxed);
        }
        self.emit(true);
    }

    /// Bytes that just crossed the network (or were copied from a drive).
    pub fn add_transferred(&self, n: u64) {
        self.transferred.fetch_add(n, Ordering::Relaxed);
        self.done_bytes.fetch_add(n, Ordering::Relaxed);
        self.emit(false);
    }

    /// Bytes found already in place (skipped files, resumed prefixes, verified chunks).
    pub fn add_done_bytes(&self, n: u64) {
        self.done_bytes.fetch_add(n, Ordering::Relaxed);
        self.emit(false);
    }

    /// Take back bytes counted for a download that failed verification and restarts.
    pub fn retract_bytes(&self, n: u64) {
        let _ = self.done_bytes.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| Some(v.saturating_sub(n)));
    }

    pub fn items_done(&self, n: u64) {
        self.done_items.fetch_add(n, Ordering::Relaxed);
        self.sub_done.fetch_add(n, Ordering::Relaxed);
        self.emit(false);
    }

    pub fn finish_stage(&self) {
        self.emit(true);
    }

    pub fn transferred(&self) -> u64 {
        self.transferred.load(Ordering::Relaxed)
    }

    fn emit(&self, force: bool) {
        let now = Instant::now();
        let progress = {
            let mut st = if force {
                self.state.lock().unwrap_or_else(|p| p.into_inner())
            } else {
                match self.state.try_lock() {
                    Ok(guard) => guard,
                    Err(_) => return, // another worker is reporting right now
                }
            };
            if !force && st.last_emit.is_some_and(|last| now.duration_since(last) < INTERVAL) {
                return;
            }
            st.last_emit = Some(now);
            let transferred = self.transferred.load(Ordering::Relaxed);
            let dt = now.duration_since(st.rate_at).as_secs_f64();
            if dt >= 0.05 {
                let instant = transferred.saturating_sub(st.rate_bytes) as f64 / dt;
                let alpha = (dt / RATE_WINDOW_SECS).min(1.0);
                st.rate += alpha * (instant - st.rate);
                st.rate_bytes = transferred;
                st.rate_at = now;
            }
            let mut done_bytes = self.done_bytes.load(Ordering::Relaxed);
            if st.total_bytes > 0 {
                done_bytes = done_bytes.min(st.total_bytes);
            }
            let message = (!st.label.is_empty()).then(|| {
                if st.sub_total == 0 {
                    st.label.clone()
                } else {
                    format!("{}: {}/{} {}", st.label, self.sub_done.load(Ordering::Relaxed), st.sub_total, st.unit)
                }
            });
            Progress {
                stage: st.stage,
                stage_index: st.stage_index,
                stage_count: st.stage_count,
                done_bytes,
                total_bytes: st.total_bytes,
                done_items: self.done_items.load(Ordering::Relaxed),
                total_items: st.total_items,
                bytes_per_sec: st.rate.max(0.0).round(),
                message,
            }
        };
        (self.callback)(progress);
    }
}