Skip to main content

a3s_box_runtime/oci/registry/
progress.rs

1use std::sync::Arc;
2use std::time::{Duration, Instant};
3
4const REPORT_INTERVAL_BYTES: u64 = 1024 * 1024;
5const REPORT_INTERVAL_TIME: Duration = Duration::from_secs(5);
6
7/// Current phase of a registry layer transfer.
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum PullProgressState {
10    Downloading,
11    Retrying,
12    Reused,
13    Complete,
14}
15
16/// Structured registry layer progress with actual transferred byte counts.
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct PullProgress {
19    pub current_layer: usize,
20    pub total_layers: usize,
21    pub digest: String,
22    pub downloaded_bytes: u64,
23    pub total_bytes: u64,
24    pub attempt: usize,
25    pub max_attempts: usize,
26    pub retry_delay_ms: Option<u64>,
27    pub state: PullProgressState,
28}
29
30/// Callback for structured registry layer progress events.
31pub type PullProgressEventFn = Arc<dyn Fn(PullProgress) + Send + Sync>;
32
33pub(super) struct ProgressReporter {
34    callback: Option<PullProgressEventFn>,
35    current_layer: usize,
36    total_layers: usize,
37    digest: String,
38    total_bytes: u64,
39    max_attempts: usize,
40    last_bytes: u64,
41    last_report: Instant,
42}
43
44impl ProgressReporter {
45    pub(super) fn new(
46        callback: Option<PullProgressEventFn>,
47        current_layer: usize,
48        total_layers: usize,
49        digest: String,
50        total_bytes: u64,
51        max_attempts: usize,
52    ) -> Self {
53        Self {
54            callback,
55            current_layer,
56            total_layers,
57            digest,
58            total_bytes,
59            max_attempts,
60            last_bytes: 0,
61            last_report: Instant::now(),
62        }
63    }
64
65    pub(super) fn downloading(&mut self, downloaded_bytes: u64, attempt: usize, force: bool) {
66        let now = Instant::now();
67        if !force
68            && downloaded_bytes.saturating_sub(self.last_bytes) < REPORT_INTERVAL_BYTES
69            && now.duration_since(self.last_report) < REPORT_INTERVAL_TIME
70        {
71            return;
72        }
73        self.last_bytes = downloaded_bytes;
74        self.last_report = now;
75        self.emit(
76            PullProgressState::Downloading,
77            downloaded_bytes,
78            attempt,
79            None,
80        );
81    }
82
83    pub(super) fn retrying(&mut self, downloaded_bytes: u64, next_attempt: usize, delay: Duration) {
84        self.emit(
85            PullProgressState::Retrying,
86            downloaded_bytes,
87            next_attempt,
88            Some(duration_millis(delay)),
89        );
90    }
91
92    pub(super) fn reused(&mut self) {
93        self.emit(PullProgressState::Reused, self.total_bytes, 0, None);
94    }
95
96    pub(super) fn complete(&mut self, downloaded_bytes: u64, attempt: usize) {
97        self.emit(PullProgressState::Complete, downloaded_bytes, attempt, None);
98    }
99
100    fn emit(
101        &self,
102        state: PullProgressState,
103        downloaded_bytes: u64,
104        attempt: usize,
105        retry_delay_ms: Option<u64>,
106    ) {
107        tracing::info!(
108            layer = self.current_layer,
109            layers = self.total_layers,
110            digest = %self.digest,
111            downloaded_bytes,
112            total_bytes = self.total_bytes,
113            attempt,
114            max_attempts = self.max_attempts,
115            ?state,
116            "Registry layer transfer progress"
117        );
118        if let Some(callback) = &self.callback {
119            callback(PullProgress {
120                current_layer: self.current_layer,
121                total_layers: self.total_layers,
122                digest: self.digest.clone(),
123                downloaded_bytes,
124                total_bytes: self.total_bytes,
125                attempt,
126                max_attempts: self.max_attempts,
127                retry_delay_ms,
128                state,
129            });
130        }
131    }
132}
133
134fn duration_millis(duration: Duration) -> u64 {
135    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
136}