reth_cli_commands/download/
planning.rs

1use super::{manifest::*, verify::OutputVerifier};
2use eyre::Result;
3use std::{collections::BTreeMap, path::Path};
4use tracing::info;
5
6/// One archive selected from the manifest, along with its component name.
7#[derive(Debug, Clone)]
8pub(crate) struct PlannedArchive {
9    /// Snapshot component type this archive belongs to.
10    pub(crate) ty: SnapshotComponentType,
11    /// User-facing component name used in logs.
12    pub(crate) component: String,
13    /// Concrete snapshot archive metadata resolved from the manifest.
14    pub(crate) archive: SnapshotArchive,
15}
16
17/// The archive list for a modular snapshot download.
18#[derive(Debug)]
19pub(crate) struct PlannedDownloads {
20    /// Concrete archives that still need reuse checks or processing.
21    pub(crate) archives: Vec<PlannedArchive>,
22    /// Total compressed download size of all planned archives.
23    pub(crate) total_download_size: u64,
24    /// Total extracted plain-output size of all planned archives.
25    pub(crate) total_output_size: u64,
26}
27
28impl PlannedDownloads {
29    /// Returns the number of concrete archives queued for this snapshot selection.
30    pub(crate) const fn total_archives(&self) -> usize {
31        self.archives.len()
32    }
33}
34
35/// Returns the sort priority used to schedule archives.
36pub(crate) const fn archive_priority_rank(ty: SnapshotComponentType) -> u8 {
37    match ty {
38        SnapshotComponentType::State => 0,
39        SnapshotComponentType::RocksdbIndices => 1,
40        _ => 2,
41    }
42}
43
44/// Startup summary showing how much of the selected work can be reused.
45#[derive(Debug, Default, Clone, Copy)]
46pub(crate) struct DownloadStartupSummary {
47    /// Archives whose declared outputs already verify on disk.
48    pub(crate) reusable: usize,
49    /// Archives that still need to be downloaded or retried.
50    pub(crate) needs_download: usize,
51}
52
53/// Checks selected archives against existing output files before work begins.
54pub(crate) fn summarize_download_startup(
55    all_downloads: &[PlannedArchive],
56    target_dir: &Path,
57) -> Result<DownloadStartupSummary> {
58    let mut summary = DownloadStartupSummary::default();
59    let verifier = OutputVerifier::new(target_dir);
60
61    for planned in all_downloads {
62        if verifier.verify(&planned.archive.output_files)? {
63            summary.reusable += 1;
64        } else {
65            summary.needs_download += 1;
66        }
67    }
68
69    Ok(summary)
70}
71
72/// Converts a selection into the manifest distance form used for archive lookup.
73fn selection_archive_distance(
74    selection: &ComponentSelection,
75    snapshot_block: u64,
76) -> Option<Option<u64>> {
77    match selection {
78        ComponentSelection::All => Some(None),
79        ComponentSelection::Distance(distance) => Some(Some(*distance)),
80        ComponentSelection::Since(block) => Some(Some(snapshot_block.saturating_sub(*block) + 1)),
81        ComponentSelection::None => None,
82    }
83}
84
85/// Sorts planned archives into a stable processing order.
86fn sort_planned_archives(all_downloads: &mut [PlannedArchive]) {
87    all_downloads.sort_by(|a, b| {
88        archive_priority_rank(a.ty)
89            .cmp(&archive_priority_rank(b.ty))
90            .then_with(|| a.component.cmp(&b.component))
91            .then_with(|| a.archive.file_name.cmp(&b.archive.file_name))
92    });
93}
94
95/// Expands component selections into the archives that need to be processed.
96pub(crate) fn collect_planned_archives(
97    manifest: &SnapshotManifest,
98    selections: &BTreeMap<SnapshotComponentType, ComponentSelection>,
99) -> Result<PlannedDownloads> {
100    let mut archives = Vec::new();
101    let mut total_download_size = 0;
102    let mut total_output_size = 0;
103
104    for (ty, selection) in selections {
105        let Some(distance) = selection_archive_distance(selection, manifest.block) else {
106            continue;
107        };
108        total_download_size += manifest.size_for_distance(*ty, distance);
109        total_output_size += manifest.output_size_for_distance(*ty, distance);
110
111        let snapshot_archives = manifest.snapshot_archives_for_distance(*ty, distance);
112        let component = ty.display_name().to_string();
113        if !snapshot_archives.is_empty() {
114            info!(target: "reth::cli",
115                component = %component,
116                archives = snapshot_archives.len(),
117                selection = %selection,
118                "Queued component for download"
119            );
120        }
121
122        for archive in snapshot_archives {
123            if archive.output_files.is_empty() {
124                eyre::bail!(
125                    "Invalid modular manifest: {} is missing plain output checksum metadata",
126                    archive.file_name
127                );
128            }
129
130            archives.push(PlannedArchive { ty: *ty, component: component.clone(), archive });
131        }
132    }
133
134    sort_planned_archives(&mut archives);
135    Ok(PlannedDownloads { archives, total_download_size, total_output_size })
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use tempfile::tempdir;
142
143    #[test]
144    fn summarize_download_startup_counts_reusable_and_needs_download() {
145        let dir = tempdir().unwrap();
146        let target_dir = dir.path();
147        let ok_file = target_dir.join("ok.bin");
148        std::fs::write(&ok_file, vec![1_u8; 4]).unwrap();
149        let ok_hash = blake3::hash(&[1_u8; 4]).to_hex().to_string();
150
151        let planned = vec![
152            PlannedArchive {
153                ty: SnapshotComponentType::State,
154                component: "State".to_string(),
155                archive: SnapshotArchive {
156                    url: "https://example.com/ok.tar.zst".to_string(),
157                    file_name: "ok.tar.zst".to_string(),
158                    size: 10,
159                    blake3: None,
160                    output_files: vec![OutputFileChecksum {
161                        path: "ok.bin".to_string(),
162                        size: 4,
163                        blake3: ok_hash,
164                    }],
165                },
166            },
167            PlannedArchive {
168                ty: SnapshotComponentType::Headers,
169                component: "Headers".to_string(),
170                archive: SnapshotArchive {
171                    url: "https://example.com/missing.tar.zst".to_string(),
172                    file_name: "missing.tar.zst".to_string(),
173                    size: 10,
174                    blake3: None,
175                    output_files: vec![OutputFileChecksum {
176                        path: "missing.bin".to_string(),
177                        size: 1,
178                        blake3: "deadbeef".to_string(),
179                    }],
180                },
181            },
182            PlannedArchive {
183                ty: SnapshotComponentType::Transactions,
184                component: "Transactions".to_string(),
185                archive: SnapshotArchive {
186                    url: "https://example.com/bad-size.tar.zst".to_string(),
187                    file_name: "bad-size.tar.zst".to_string(),
188                    size: 10,
189                    blake3: None,
190                    output_files: vec![],
191                },
192            },
193        ];
194
195        let summary = summarize_download_startup(&planned, target_dir).unwrap();
196        assert_eq!(summary.reusable, 1);
197        assert_eq!(summary.needs_download, 2);
198    }
199
200    #[test]
201    fn archive_priority_prefers_state_then_rocksdb() {
202        let mut planned = [
203            PlannedArchive {
204                ty: SnapshotComponentType::Transactions,
205                component: "Transactions".to_string(),
206                archive: SnapshotArchive {
207                    url: "u3".to_string(),
208                    file_name: "t.tar.zst".to_string(),
209                    size: 1,
210                    blake3: None,
211                    output_files: vec![OutputFileChecksum {
212                        path: "a".to_string(),
213                        size: 1,
214                        blake3: "x".to_string(),
215                    }],
216                },
217            },
218            PlannedArchive {
219                ty: SnapshotComponentType::RocksdbIndices,
220                component: "RocksDB Indices".to_string(),
221                archive: SnapshotArchive {
222                    url: "u2".to_string(),
223                    file_name: "rocksdb_indices.tar.zst".to_string(),
224                    size: 1,
225                    blake3: None,
226                    output_files: vec![OutputFileChecksum {
227                        path: "b".to_string(),
228                        size: 1,
229                        blake3: "y".to_string(),
230                    }],
231                },
232            },
233            PlannedArchive {
234                ty: SnapshotComponentType::State,
235                component: "State (mdbx)".to_string(),
236                archive: SnapshotArchive {
237                    url: "u1".to_string(),
238                    file_name: "state.tar.zst".to_string(),
239                    size: 1,
240                    blake3: None,
241                    output_files: vec![OutputFileChecksum {
242                        path: "c".to_string(),
243                        size: 1,
244                        blake3: "z".to_string(),
245                    }],
246                },
247            },
248        ];
249
250        planned.sort_by(|a, b| {
251            archive_priority_rank(a.ty)
252                .cmp(&archive_priority_rank(b.ty))
253                .then_with(|| a.component.cmp(&b.component))
254                .then_with(|| a.archive.file_name.cmp(&b.archive.file_name))
255        });
256
257        assert_eq!(planned[0].ty, SnapshotComponentType::State);
258        assert_eq!(planned[1].ty, SnapshotComponentType::RocksdbIndices);
259        assert_eq!(planned[2].ty, SnapshotComponentType::Transactions);
260    }
261
262    #[test]
263    fn collect_planned_archives_tracks_download_and_output_totals() {
264        let mut components = BTreeMap::new();
265        components.insert(
266            "state".to_string(),
267            ComponentManifest::Single(SingleArchive {
268                file: "state.tar.zst".to_string(),
269                size: 10,
270                decompressed_size: 100,
271                blake3: None,
272                output_files: vec![OutputFileChecksum {
273                    path: "db/mdbx.dat".to_string(),
274                    size: 100,
275                    blake3: "h0".to_string(),
276                }],
277            }),
278        );
279        components.insert(
280            "transactions".to_string(),
281            ComponentManifest::Chunked(ChunkedArchive {
282                blocks_per_file: 500_000,
283                total_blocks: 1_000_000,
284                chunk_sizes: vec![20, 30],
285                chunk_decompressed_sizes: vec![200, 300],
286                chunk_output_files: vec![
287                    vec![OutputFileChecksum {
288                        path: "static_files/tx-0".to_string(),
289                        size: 200,
290                        blake3: "h1".to_string(),
291                    }],
292                    vec![OutputFileChecksum {
293                        path: "static_files/tx-1".to_string(),
294                        size: 300,
295                        blake3: "h2".to_string(),
296                    }],
297                ],
298            }),
299        );
300
301        let manifest = SnapshotManifest {
302            block: 1_000_000,
303            chain_id: 1,
304            storage_version: 2,
305            timestamp: 0,
306            base_url: Some("https://example.com".to_string()),
307            reth_version: None,
308            components,
309        };
310
311        let selections = BTreeMap::from([
312            (SnapshotComponentType::State, ComponentSelection::All),
313            (SnapshotComponentType::Transactions, ComponentSelection::Distance(500_000)),
314        ]);
315
316        let planned = collect_planned_archives(&manifest, &selections).unwrap();
317
318        assert_eq!(planned.total_download_size, 40);
319        assert_eq!(planned.total_output_size, 400);
320        assert_eq!(planned.archives.len(), 2);
321    }
322}
reth_cli_commands/download/planning.rs

reth_cli_commands/download/
planning.rs
