reth_cli_commands/download/
planning.rs

1use super::{manifest::*, verify::OutputVerifier};
2use eyre::Result;
3use serde::Serialize;
4use std::{collections::BTreeMap, io::Write, path::Path};
5use tracing::info;
6
7/// Machine-readable description of the archives selected for a modular snapshot download.
8#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
9#[serde(rename_all = "camelCase")]
10pub struct DownloadPlan {
11    /// Version of this JSON schema.
12    pub schema_version: u8,
13    /// Chain ID declared by the snapshot manifest.
14    pub chain_id: u64,
15    /// Block declared by the snapshot manifest.
16    pub block: u64,
17    /// Total compressed bytes for all selected archives.
18    pub total_download_size: u64,
19    /// Total extracted bytes for all selected archives.
20    pub total_output_size: u64,
21    /// Archives selected by the command's component and range flags.
22    pub archives: Vec<DownloadPlanArchive>,
23}
24
25impl DownloadPlan {
26    const SCHEMA_VERSION: u8 = 1;
27
28    pub(crate) fn from_planned(manifest: &SnapshotManifest, planned: &PlannedDownloads) -> Self {
29        Self {
30            schema_version: Self::SCHEMA_VERSION,
31            chain_id: manifest.chain_id,
32            block: manifest.block,
33            total_download_size: planned.total_download_size,
34            total_output_size: planned.total_output_size,
35            archives: planned.archives.iter().map(DownloadPlanArchive::from_planned).collect(),
36        }
37    }
38
39    /// Adds an archive supplied by a command that extends the base Reth snapshot format.
40    pub fn push_archive(&mut self, archive: DownloadPlanArchive) {
41        self.total_download_size = self.total_download_size.saturating_add(archive.download_size);
42        self.total_output_size = self.total_output_size.saturating_add(archive.output_size);
43        self.archives.push(archive);
44    }
45
46    /// Writes the plan as pretty-printed JSON followed by a newline.
47    pub fn write_json(&self, mut writer: impl Write) -> Result<()> {
48        serde_json::to_writer_pretty(&mut writer, self)?;
49        writeln!(writer)?;
50        Ok(())
51    }
52}
53
54/// One concrete archive in a [`DownloadPlan`].
55#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
56#[serde(rename_all = "camelCase")]
57pub struct DownloadPlanArchive {
58    /// Stable manifest component key, such as `state` or `transactions`.
59    pub component: String,
60    /// Archive file name relative to the manifest's base URL.
61    pub file_name: String,
62    /// Resolved archive source URL.
63    pub url: String,
64    /// Compressed bytes fetched for this archive.
65    pub download_size: u64,
66    /// Extracted bytes produced by this archive.
67    pub output_size: u64,
68    /// Optional checksum of the compressed archive.
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub blake3: Option<String>,
71}
72
73impl DownloadPlanArchive {
74    /// Creates an archive entry for a download-command extension.
75    pub fn new(
76        component: impl Into<String>,
77        file_name: impl Into<String>,
78        url: impl Into<String>,
79        download_size: u64,
80        output_size: u64,
81        blake3: Option<String>,
82    ) -> Self {
83        Self {
84            component: component.into(),
85            file_name: file_name.into(),
86            url: url.into(),
87            download_size,
88            output_size,
89            blake3,
90        }
91    }
92
93    fn from_planned(planned: &PlannedArchive) -> Self {
94        Self::new(
95            planned.ty.key(),
96            planned.archive.file_name.clone(),
97            planned.archive.url.clone(),
98            planned.archive.size,
99            planned.archive.output_size(),
100            planned.archive.blake3.clone(),
101        )
102    }
103}
104
105/// One archive selected from the manifest, along with its component name.
106#[derive(Debug, Clone)]
107pub(crate) struct PlannedArchive {
108    /// Snapshot component type this archive belongs to.
109    pub(crate) ty: SnapshotComponentType,
110    /// User-facing component name used in logs.
111    pub(crate) component: String,
112    /// Concrete snapshot archive metadata resolved from the manifest.
113    pub(crate) archive: SnapshotArchive,
114}
115
116/// The archive list for a modular snapshot download.
117#[derive(Debug)]
118pub(crate) struct PlannedDownloads {
119    /// Concrete archives that still need reuse checks or processing.
120    pub(crate) archives: Vec<PlannedArchive>,
121    /// Total compressed download size of all planned archives.
122    pub(crate) total_download_size: u64,
123    /// Total extracted plain-output size of all planned archives.
124    pub(crate) total_output_size: u64,
125}
126
127impl PlannedDownloads {
128    /// Returns the number of concrete archives queued for this snapshot selection.
129    pub(crate) const fn total_archives(&self) -> usize {
130        self.archives.len()
131    }
132}
133
134/// Returns the sort priority used to schedule archives.
135pub(crate) const fn archive_priority_rank(ty: SnapshotComponentType) -> u8 {
136    match ty {
137        SnapshotComponentType::State => 0,
138        SnapshotComponentType::RocksdbIndices => 1,
139        _ => 2,
140    }
141}
142
143/// Startup summary showing how much of the selected work can be reused.
144#[derive(Debug, Default, Clone, Copy)]
145pub(crate) struct DownloadStartupSummary {
146    /// Archives whose declared outputs already verify on disk.
147    pub(crate) reusable: usize,
148    /// Archives that still need to be downloaded or retried.
149    pub(crate) needs_download: usize,
150}
151
152/// Checks selected archives against existing output files before work begins.
153pub(crate) fn summarize_download_startup(
154    all_downloads: &[PlannedArchive],
155    target_dir: &Path,
156) -> Result<DownloadStartupSummary> {
157    let mut summary = DownloadStartupSummary::default();
158    let verifier = OutputVerifier::new(target_dir);
159
160    for planned in all_downloads {
161        if verifier.verify(&planned.archive.output_files)? {
162            summary.reusable += 1;
163        } else {
164            summary.needs_download += 1;
165        }
166    }
167
168    Ok(summary)
169}
170
171/// Converts a selection into the manifest distance form used for archive lookup.
172fn selection_archive_distance(
173    selection: &ComponentSelection,
174    snapshot_block: u64,
175) -> Option<Option<u64>> {
176    match selection {
177        ComponentSelection::All => Some(None),
178        ComponentSelection::Distance(distance) => Some(Some(*distance)),
179        ComponentSelection::Since(block) => Some(Some(snapshot_block.saturating_sub(*block) + 1)),
180        ComponentSelection::None => None,
181    }
182}
183
184/// Sorts planned archives into a stable processing order.
185fn sort_planned_archives(all_downloads: &mut [PlannedArchive]) {
186    all_downloads.sort_by(|a, b| {
187        archive_priority_rank(a.ty)
188            .cmp(&archive_priority_rank(b.ty))
189            .then_with(|| a.component.cmp(&b.component))
190            .then_with(|| a.archive.file_name.cmp(&b.archive.file_name))
191    });
192}
193
194/// Expands component selections into the archives that need to be processed.
195pub(crate) fn collect_planned_archives(
196    manifest: &SnapshotManifest,
197    selections: &BTreeMap<SnapshotComponentType, ComponentSelection>,
198) -> Result<PlannedDownloads> {
199    let mut archives = Vec::new();
200    let mut total_download_size = 0;
201    let mut total_output_size = 0;
202
203    for (ty, selection) in selections {
204        let Some(distance) = selection_archive_distance(selection, manifest.block) else {
205            continue;
206        };
207        total_download_size += manifest.size_for_distance(*ty, distance);
208        total_output_size += manifest.output_size_for_distance(*ty, distance);
209
210        let snapshot_archives = manifest.snapshot_archives_for_distance(*ty, distance);
211        let component = ty.display_name().to_string();
212        if !snapshot_archives.is_empty() {
213            info!(target: "reth::cli",
214                component = %component,
215                archives = snapshot_archives.len(),
216                selection = %selection,
217                "Queued component for download"
218            );
219        }
220
221        for archive in snapshot_archives {
222            if archive.output_files.is_empty() {
223                eyre::bail!(
224                    "Invalid modular manifest: {} is missing plain output checksum metadata",
225                    archive.file_name
226                );
227            }
228
229            archives.push(PlannedArchive { ty: *ty, component: component.clone(), archive });
230        }
231    }
232
233    sort_planned_archives(&mut archives);
234    Ok(PlannedDownloads { archives, total_download_size, total_output_size })
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use tempfile::tempdir;
241
242    #[test]
243    fn summarize_download_startup_counts_reusable_and_needs_download() {
244        let dir = tempdir().unwrap();
245        let target_dir = dir.path();
246        let ok_file = target_dir.join("ok.bin");
247        std::fs::write(&ok_file, vec![1_u8; 4]).unwrap();
248        let ok_hash = blake3::hash(&[1_u8; 4]).to_hex().to_string();
249
250        let planned = vec![
251            PlannedArchive {
252                ty: SnapshotComponentType::State,
253                component: "State".to_string(),
254                archive: SnapshotArchive {
255                    url: "https://example.com/ok.tar.zst".to_string(),
256                    file_name: "ok.tar.zst".to_string(),
257                    size: 10,
258                    blake3: None,
259                    output_files: vec![OutputFileChecksum {
260                        path: "ok.bin".to_string(),
261                        size: 4,
262                        blake3: ok_hash,
263                    }],
264                },
265            },
266            PlannedArchive {
267                ty: SnapshotComponentType::Headers,
268                component: "Headers".to_string(),
269                archive: SnapshotArchive {
270                    url: "https://example.com/missing.tar.zst".to_string(),
271                    file_name: "missing.tar.zst".to_string(),
272                    size: 10,
273                    blake3: None,
274                    output_files: vec![OutputFileChecksum {
275                        path: "missing.bin".to_string(),
276                        size: 1,
277                        blake3: "deadbeef".to_string(),
278                    }],
279                },
280            },
281            PlannedArchive {
282                ty: SnapshotComponentType::Transactions,
283                component: "Transactions".to_string(),
284                archive: SnapshotArchive {
285                    url: "https://example.com/bad-size.tar.zst".to_string(),
286                    file_name: "bad-size.tar.zst".to_string(),
287                    size: 10,
288                    blake3: None,
289                    output_files: vec![],
290                },
291            },
292        ];
293
294        let summary = summarize_download_startup(&planned, target_dir).unwrap();
295        assert_eq!(summary.reusable, 1);
296        assert_eq!(summary.needs_download, 2);
297    }
298
299    #[test]
300    fn archive_priority_prefers_state_then_rocksdb() {
301        let mut planned = [
302            PlannedArchive {
303                ty: SnapshotComponentType::Transactions,
304                component: "Transactions".to_string(),
305                archive: SnapshotArchive {
306                    url: "u3".to_string(),
307                    file_name: "t.tar.zst".to_string(),
308                    size: 1,
309                    blake3: None,
310                    output_files: vec![OutputFileChecksum {
311                        path: "a".to_string(),
312                        size: 1,
313                        blake3: "x".to_string(),
314                    }],
315                },
316            },
317            PlannedArchive {
318                ty: SnapshotComponentType::RocksdbIndices,
319                component: "RocksDB Indices".to_string(),
320                archive: SnapshotArchive {
321                    url: "u2".to_string(),
322                    file_name: "rocksdb_indices.tar.zst".to_string(),
323                    size: 1,
324                    blake3: None,
325                    output_files: vec![OutputFileChecksum {
326                        path: "b".to_string(),
327                        size: 1,
328                        blake3: "y".to_string(),
329                    }],
330                },
331            },
332            PlannedArchive {
333                ty: SnapshotComponentType::State,
334                component: "State (mdbx)".to_string(),
335                archive: SnapshotArchive {
336                    url: "u1".to_string(),
337                    file_name: "state.tar.zst".to_string(),
338                    size: 1,
339                    blake3: None,
340                    output_files: vec![OutputFileChecksum {
341                        path: "c".to_string(),
342                        size: 1,
343                        blake3: "z".to_string(),
344                    }],
345                },
346            },
347        ];
348
349        planned.sort_by(|a, b| {
350            archive_priority_rank(a.ty)
351                .cmp(&archive_priority_rank(b.ty))
352                .then_with(|| a.component.cmp(&b.component))
353                .then_with(|| a.archive.file_name.cmp(&b.archive.file_name))
354        });
355
356        assert_eq!(planned[0].ty, SnapshotComponentType::State);
357        assert_eq!(planned[1].ty, SnapshotComponentType::RocksdbIndices);
358        assert_eq!(planned[2].ty, SnapshotComponentType::Transactions);
359    }
360
361    #[test]
362    fn collect_planned_archives_tracks_download_and_output_totals() {
363        let mut components = BTreeMap::new();
364        components.insert(
365            "state".to_string(),
366            ComponentManifest::Single(SingleArchive {
367                file: "state.tar.zst".to_string(),
368                size: 10,
369                decompressed_size: 100,
370                blake3: None,
371                output_files: vec![OutputFileChecksum {
372                    path: "db/mdbx.dat".to_string(),
373                    size: 100,
374                    blake3: "h0".to_string(),
375                }],
376            }),
377        );
378        components.insert(
379            "transactions".to_string(),
380            ComponentManifest::Chunked(ChunkedArchive {
381                blocks_per_file: 500_000,
382                total_blocks: 1_000_000,
383                chunk_sizes: vec![20, 30],
384                chunk_decompressed_sizes: vec![200, 300],
385                chunk_output_files: vec![
386                    vec![OutputFileChecksum {
387                        path: "static_files/tx-0".to_string(),
388                        size: 200,
389                        blake3: "h1".to_string(),
390                    }],
391                    vec![OutputFileChecksum {
392                        path: "static_files/tx-1".to_string(),
393                        size: 300,
394                        blake3: "h2".to_string(),
395                    }],
396                ],
397            }),
398        );
399
400        let manifest = SnapshotManifest {
401            block: 1_000_000,
402            chain_id: 1,
403            storage_version: 2,
404            timestamp: 0,
405            base_url: Some("https://example.com".to_string()),
406            reth_version: None,
407            components,
408        };
409
410        let selections = BTreeMap::from([
411            (SnapshotComponentType::State, ComponentSelection::All),
412            (SnapshotComponentType::Transactions, ComponentSelection::Distance(500_000)),
413        ]);
414
415        let planned = collect_planned_archives(&manifest, &selections).unwrap();
416
417        assert_eq!(planned.total_download_size, 40);
418        assert_eq!(planned.total_output_size, 400);
419        assert_eq!(planned.archives.len(), 2);
420
421        let plan = DownloadPlan::from_planned(&manifest, &planned);
422        assert_eq!(
423            serde_json::to_value(plan).unwrap(),
424            serde_json::json!({
425                "schemaVersion": 1,
426                "chainId": 1,
427                "block": 1_000_000,
428                "totalDownloadSize": 40,
429                "totalOutputSize": 400,
430                "archives": [
431                    {
432                        "component": "state",
433                        "fileName": "state.tar.zst",
434                        "url": "https://example.com/state.tar.zst",
435                        "downloadSize": 10,
436                        "outputSize": 100
437                    },
438                    {
439                        "component": "transactions",
440                        "fileName": "transactions-500000-999999.tar.zst",
441                        "url": "https://example.com/transactions-500000-999999.tar.zst",
442                        "downloadSize": 30,
443                        "outputSize": 300
444                    }
445                ]
446            })
447        );
448    }
449}
reth_cli_commands/download/planning.rs

reth_cli_commands/download/
planning.rs
