krataoci/
progress.rs

1use indexmap::IndexMap;
2use std::sync::Arc;
3use tokio::{
4    sync::{watch, Mutex},
5    task::JoinHandle,
6};
7
8#[derive(Clone, Debug)]
9pub struct OciProgress {
10    pub phase: OciProgressPhase,
11    pub digest: Option<String>,
12    pub layers: IndexMap<String, OciProgressLayer>,
13    pub indication: OciProgressIndication,
14}
15
16impl Default for OciProgress {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl OciProgress {
23    pub fn new() -> Self {
24        OciProgress {
25            phase: OciProgressPhase::Started,
26            digest: None,
27            layers: IndexMap::new(),
28            indication: OciProgressIndication::Hidden,
29        }
30    }
31
32    pub fn start_resolving(&mut self) {
33        self.phase = OciProgressPhase::Resolving;
34        self.indication = OciProgressIndication::Spinner { message: None };
35    }
36
37    pub fn resolved(&mut self, digest: &str) {
38        self.digest = Some(digest.to_string());
39        self.indication = OciProgressIndication::Hidden;
40    }
41
42    pub fn add_layer(&mut self, id: &str) {
43        self.layers.insert(
44            id.to_string(),
45            OciProgressLayer {
46                id: id.to_string(),
47                phase: OciProgressLayerPhase::Waiting,
48                indication: OciProgressIndication::Spinner { message: None },
49            },
50        );
51    }
52
53    pub fn downloading_layer(&mut self, id: &str, downloaded: u64, total: u64) {
54        if let Some(entry) = self.layers.get_mut(id) {
55            entry.phase = OciProgressLayerPhase::Downloading;
56            entry.indication = OciProgressIndication::ProgressBar {
57                message: None,
58                current: downloaded,
59                total,
60                bytes: true,
61            };
62        }
63    }
64
65    pub fn downloaded_layer(&mut self, id: &str, total: u64) {
66        if let Some(entry) = self.layers.get_mut(id) {
67            entry.phase = OciProgressLayerPhase::Downloaded;
68            entry.indication = OciProgressIndication::Completed {
69                message: None,
70                total: Some(total),
71                bytes: true,
72            };
73        }
74    }
75
76    pub fn start_assemble(&mut self) {
77        self.phase = OciProgressPhase::Assemble;
78        self.indication = OciProgressIndication::Hidden;
79    }
80
81    pub fn start_extracting_layer(&mut self, id: &str) {
82        if let Some(entry) = self.layers.get_mut(id) {
83            entry.phase = OciProgressLayerPhase::Extracting;
84            entry.indication = OciProgressIndication::Spinner { message: None };
85        }
86    }
87
88    pub fn extracting_layer(&mut self, id: &str, file: &str) {
89        if let Some(entry) = self.layers.get_mut(id) {
90            entry.phase = OciProgressLayerPhase::Extracting;
91            entry.indication = OciProgressIndication::Spinner {
92                message: Some(file.to_string()),
93            };
94        }
95    }
96
97    pub fn extracted_layer(&mut self, id: &str, count: u64, total_size: u64) {
98        if let Some(entry) = self.layers.get_mut(id) {
99            entry.phase = OciProgressLayerPhase::Extracted;
100            entry.indication = OciProgressIndication::Completed {
101                message: Some(format!("{} files", count)),
102                total: Some(total_size),
103                bytes: true,
104            };
105        }
106    }
107
108    pub fn start_packing(&mut self) {
109        self.phase = OciProgressPhase::Pack;
110        for layer in self.layers.values_mut() {
111            layer.indication = OciProgressIndication::Hidden;
112        }
113        self.indication = OciProgressIndication::Spinner { message: None };
114    }
115
116    pub fn complete(&mut self, size: u64) {
117        self.phase = OciProgressPhase::Complete;
118        self.indication = OciProgressIndication::Completed {
119            message: None,
120            total: Some(size),
121            bytes: true,
122        }
123    }
124}
125
126#[derive(Clone, Debug)]
127pub enum OciProgressPhase {
128    Started,
129    Resolving,
130    Resolved,
131    ConfigDownload,
132    LayerDownload,
133    Assemble,
134    Pack,
135    Complete,
136}
137
138#[derive(Clone, Debug)]
139pub enum OciProgressIndication {
140    Hidden,
141
142    ProgressBar {
143        message: Option<String>,
144        current: u64,
145        total: u64,
146        bytes: bool,
147    },
148
149    Spinner {
150        message: Option<String>,
151    },
152
153    Completed {
154        message: Option<String>,
155        total: Option<u64>,
156        bytes: bool,
157    },
158}
159
160#[derive(Clone, Debug)]
161pub struct OciProgressLayer {
162    pub id: String,
163    pub phase: OciProgressLayerPhase,
164    pub indication: OciProgressIndication,
165}
166
167#[derive(Clone, Debug)]
168pub enum OciProgressLayerPhase {
169    Waiting,
170    Downloading,
171    Downloaded,
172    Extracting,
173    Extracted,
174}
175
176#[derive(Clone)]
177pub struct OciProgressContext {
178    sender: watch::Sender<OciProgress>,
179}
180
181impl OciProgressContext {
182    pub fn create() -> (OciProgressContext, watch::Receiver<OciProgress>) {
183        let (sender, receiver) = watch::channel(OciProgress::new());
184        (OciProgressContext::new(sender), receiver)
185    }
186
187    pub fn new(sender: watch::Sender<OciProgress>) -> OciProgressContext {
188        OciProgressContext { sender }
189    }
190
191    pub fn update(&self, progress: &OciProgress) {
192        let _ = self.sender.send(progress.clone());
193    }
194
195    pub fn subscribe(&self) -> watch::Receiver<OciProgress> {
196        self.sender.subscribe()
197    }
198}
199
200#[derive(Clone)]
201pub struct OciBoundProgress {
202    context: OciProgressContext,
203    instance: Arc<Mutex<OciProgress>>,
204}
205
206impl OciBoundProgress {
207    pub fn new(context: OciProgressContext, progress: OciProgress) -> OciBoundProgress {
208        OciBoundProgress {
209            context,
210            instance: Arc::new(Mutex::new(progress)),
211        }
212    }
213
214    pub async fn update(&self, function: impl FnOnce(&mut OciProgress)) {
215        let mut progress = self.instance.lock().await;
216        function(&mut progress);
217        self.context.update(&progress);
218    }
219
220    pub fn update_blocking(&self, function: impl FnOnce(&mut OciProgress)) {
221        let mut progress = self.instance.blocking_lock();
222        function(&mut progress);
223        self.context.update(&progress);
224    }
225
226    pub async fn also_update(&self, context: OciProgressContext) -> JoinHandle<()> {
227        let progress = self.instance.lock().await.clone();
228        context.update(&progress);
229        let mut receiver = self.context.subscribe();
230        tokio::task::spawn(async move {
231            while receiver.changed().await.is_ok() {
232                context
233                    .sender
234                    .send_replace(receiver.borrow_and_update().clone());
235            }
236        })
237    }
238}