reth_cli_commands/download/
planning.rs

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

reth_cli_commands/download/
planning.rs
