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    fn test_manifest_with_component_tip(block: u64, total_blocks: u64) -> SnapshotManifest {
273        let chunk_count = total_blocks.div_ceil(10);
274        SnapshotManifest {
275            block,
276            chain_id: 1,
277            storage_version: 2,
278            timestamp: 0,
279            base_url: Some("https://example.com".to_string()),
280            reth_version: None,
281            components: BTreeMap::from([(
282                "transactions".to_string(),
283                ComponentManifest::Chunked(ChunkedArchive {
284                    blocks_per_file: 10,
285                    total_blocks,
286                    chunk_sizes: (1..=chunk_count).map(|i| i * 10).collect(),
287                    chunk_decompressed_sizes: (1..=chunk_count).map(|i| i * 100).collect(),
288                    chunk_files: vec![],
289                    chunk_output_files: (0..chunk_count)
290                        .map(|i| {
291                            vec![OutputFileChecksum {
292                                path: format!("static_files/tx-{i}"),
293                                size: (i + 1) * 100,
294                                blake3: format!("h{i}"),
295                            }]
296                        })
297                        .collect(),
298                }),
299            )]),
300            extensions: Default::default(),
301        }
302    }
303
304    #[test]
305    fn partition_reusable_moves_only_verified_archives() {
306        let dir = tempdir().unwrap();
307        let target_dir = dir.path();
308        let ok_file = target_dir.join("ok.bin");
309        std::fs::write(&ok_file, vec![1_u8; 4]).unwrap();
310        let ok_hash = blake3::hash(&[1_u8; 4]).to_hex().to_string();
311
312        let planned = vec![
313            PlannedArchive {
314                ty: SnapshotComponentType::State,
315                component: "State".to_string(),
316                archive: SnapshotArchive {
317                    url: "https://example.com/ok.tar.zst".to_string(),
318                    file_name: "ok.tar.zst".to_string(),
319                    size: 10,
320                    blake3: None,
321                    output_files: vec![OutputFileChecksum {
322                        path: "ok.bin".to_string(),
323                        size: 4,
324                        blake3: ok_hash,
325                    }],
326                },
327            },
328            PlannedArchive {
329                ty: SnapshotComponentType::Headers,
330                component: "Headers".to_string(),
331                archive: SnapshotArchive {
332                    url: "https://example.com/missing.tar.zst".to_string(),
333                    file_name: "missing.tar.zst".to_string(),
334                    size: 10,
335                    blake3: None,
336                    output_files: vec![OutputFileChecksum {
337                        path: "missing.bin".to_string(),
338                        size: 1,
339                        blake3: "deadbeef".to_string(),
340                    }],
341                },
342            },
343            PlannedArchive {
344                ty: SnapshotComponentType::Transactions,
345                component: "Transactions".to_string(),
346                archive: SnapshotArchive {
347                    url: "https://example.com/bad-size.tar.zst".to_string(),
348                    file_name: "bad-size.tar.zst".to_string(),
349                    size: 10,
350                    blake3: None,
351                    output_files: vec![],
352                },
353            },
354        ];
355
356        let downloads =
357            PlannedDownloads { archives: planned, total_download_size: 30, total_output_size: 5 }
358                .partition_reusable(target_dir, None)
359                .unwrap();
360        assert_eq!(downloads.total_archives(), 3);
361
362        let names = |archives: &[PlannedArchive]| {
363            archives.iter().map(|planned| planned.archive.file_name.clone()).collect::<Vec<_>>()
364        };
365        assert_eq!(names(&downloads.reused), ["ok.tar.zst"]);
366        assert_eq!(names(&downloads.pending), ["missing.tar.zst", "bad-size.tar.zst"]);
367    }
368
369    #[test]
370    fn archive_priority_prefers_state_then_rocksdb() {
371        let mut planned = [
372            PlannedArchive {
373                ty: SnapshotComponentType::Transactions,
374                component: "Transactions".to_string(),
375                archive: SnapshotArchive {
376                    url: "u3".to_string(),
377                    file_name: "t.tar.zst".to_string(),
378                    size: 1,
379                    blake3: None,
380                    output_files: vec![OutputFileChecksum {
381                        path: "a".to_string(),
382                        size: 1,
383                        blake3: "x".to_string(),
384                    }],
385                },
386            },
387            PlannedArchive {
388                ty: SnapshotComponentType::RocksdbIndices,
389                component: "RocksDB Indices".to_string(),
390                archive: SnapshotArchive {
391                    url: "u2".to_string(),
392                    file_name: "rocksdb_indices.tar.zst".to_string(),
393                    size: 1,
394                    blake3: None,
395                    output_files: vec![OutputFileChecksum {
396                        path: "b".to_string(),
397                        size: 1,
398                        blake3: "y".to_string(),
399                    }],
400                },
401            },
402            PlannedArchive {
403                ty: SnapshotComponentType::State,
404                component: "State (mdbx)".to_string(),
405                archive: SnapshotArchive {
406                    url: "u1".to_string(),
407                    file_name: "state.tar.zst".to_string(),
408                    size: 1,
409                    blake3: None,
410                    output_files: vec![OutputFileChecksum {
411                        path: "c".to_string(),
412                        size: 1,
413                        blake3: "z".to_string(),
414                    }],
415                },
416            },
417        ];
418
419        planned.sort_by(|a, b| {
420            archive_priority_rank(a.ty)
421                .cmp(&archive_priority_rank(b.ty))
422                .then_with(|| a.component.cmp(&b.component))
423                .then_with(|| a.archive.file_name.cmp(&b.archive.file_name))
424        });
425
426        assert_eq!(planned[0].ty, SnapshotComponentType::State);
427        assert_eq!(planned[1].ty, SnapshotComponentType::RocksdbIndices);
428        assert_eq!(planned[2].ty, SnapshotComponentType::Transactions);
429    }
430
431    #[test]
432    fn collect_planned_archives_tracks_download_and_output_totals() {
433        let mut components = BTreeMap::new();
434        components.insert(
435            "state".to_string(),
436            ComponentManifest::Single(SingleArchive {
437                file: "state.tar.zst".to_string(),
438                size: 10,
439                decompressed_size: 100,
440                blake3: None,
441                output_files: vec![OutputFileChecksum {
442                    path: "db/mdbx.dat".to_string(),
443                    size: 100,
444                    blake3: "h0".to_string(),
445                }],
446            }),
447        );
448        components.insert(
449            "transactions".to_string(),
450            ComponentManifest::Chunked(ChunkedArchive {
451                blocks_per_file: 500_000,
452                total_blocks: 1_000_000,
453                chunk_sizes: vec![20, 30],
454                chunk_decompressed_sizes: vec![200, 300],
455                chunk_files: vec![],
456                chunk_output_files: vec![
457                    vec![OutputFileChecksum {
458                        path: "static_files/tx-0".to_string(),
459                        size: 200,
460                        blake3: "h1".to_string(),
461                    }],
462                    vec![OutputFileChecksum {
463                        path: "static_files/tx-1".to_string(),
464                        size: 300,
465                        blake3: "h2".to_string(),
466                    }],
467                ],
468            }),
469        );
470
471        let manifest = SnapshotManifest {
472            block: 1_000_000,
473            chain_id: 1,
474            storage_version: 2,
475            timestamp: 0,
476            base_url: Some("https://example.com".to_string()),
477            reth_version: None,
478            components,
479            extensions: Default::default(),
480        };
481
482        let selections = BTreeMap::from([
483            (SnapshotComponentType::State, ComponentSelection::All),
484            (SnapshotComponentType::Transactions, ComponentSelection::Distance(500_000)),
485        ]);
486
487        let planned = collect_planned_archives(&manifest, &selections).unwrap();
488
489        assert_eq!(planned.total_download_size, 40);
490        assert_eq!(planned.total_output_size, 400);
491        assert_eq!(planned.archives.len(), 2);
492
493        let plan = DownloadPlan::from_planned(&manifest, &planned);
494        assert_eq!(
495            serde_json::to_value(plan).unwrap(),
496            serde_json::json!({
497                "schemaVersion": 1,
498                "chainId": 1,
499                "block": 1_000_000,
500                "totalDownloadSize": 40,
501                "totalOutputSize": 400,
502                "archives": [
503                    {
504                        "component": "state",
505                        "fileName": "state.tar.zst",
506                        "url": "https://example.com/state.tar.zst",
507                        "downloadSize": 10,
508                        "outputSize": 100
509                    },
510                    {
511                        "component": "transactions",
512                        "fileName": "transactions-500000-999999.tar.zst",
513                        "url": "https://example.com/transactions-500000-999999.tar.zst",
514                        "downloadSize": 30,
515                        "outputSize": 300
516                    }
517                ]
518            })
519        );
520    }
521
522    #[test]
523    fn collect_planned_archives_rejects_mismatched_chunk_files_length() {
524        let mut components = BTreeMap::new();
525        components.insert(
526            "transactions".to_string(),
527            ComponentManifest::Chunked(ChunkedArchive {
528                blocks_per_file: 500_000,
529                total_blocks: 1_000_000,
530                chunk_sizes: vec![20, 30],
531                chunk_decompressed_sizes: vec![200, 300],
532                chunk_files: vec!["static_files/transactions-0-499999.tar.zst".to_string()],
533                chunk_output_files: vec![
534                    vec![OutputFileChecksum {
535                        path: "static_files/tx-0".to_string(),
536                        size: 200,
537                        blake3: "h1".to_string(),
538                    }],
539                    vec![OutputFileChecksum {
540                        path: "static_files/tx-1".to_string(),
541                        size: 300,
542                        blake3: "h2".to_string(),
543                    }],
544                ],
545            }),
546        );
547
548        let manifest = SnapshotManifest {
549            block: 1_000_000,
550            chain_id: 1,
551            storage_version: 2,
552            timestamp: 0,
553            base_url: Some("https://example.com/mainnet".to_string()),
554            reth_version: None,
555            components,
556            extensions: Default::default(),
557        };
558        let selections =
559            BTreeMap::from([(SnapshotComponentType::Transactions, ComponentSelection::All)]);
560
561        let err = collect_planned_archives(&manifest, &selections).unwrap_err();
562        assert!(
563            err.to_string().contains("chunk_files length"),
564            "expected chunk_files length error, got: {err}"
565        );
566    }
567
568    #[test]
569    fn range_selections_use_snapshot_block_when_component_extends_past_it() {
570        // The snapshot ends at 20, but the final archive contains blocks 20 through 24.
571        let manifest = test_manifest_with_component_tip(20, 25);
572        let component = SnapshotComponentType::Transactions;
573        for (selection, expected_chunks) in [
574            (ComponentSelection::Since(19), vec![1, 2]),
575            (ComponentSelection::Distance(2), vec![1, 2]),
576            (ComponentSelection::Since(20), vec![2]),
577            (ComponentSelection::Distance(1), vec![2]),
578            (ComponentSelection::Since(10), vec![1, 2]),
579            (ComponentSelection::Since(9), vec![0, 1, 2]),
580            (ComponentSelection::Distance(0), vec![]),
581            (ComponentSelection::Distance(u64::MAX), vec![0, 1, 2]),
582            (ComponentSelection::All, vec![0, 1, 2]),
583        ] {
584            let planned =
585                collect_planned_archives(&manifest, &BTreeMap::from([(component, selection)]))
586                    .unwrap();
587            let expected_files = expected_chunks
588                .iter()
589                .map(|i| format!("transactions-{}-{}.tar.zst", i * 10, (i + 1) * 10 - 1))
590                .collect::<Vec<_>>();
591            let expected_size = expected_chunks.iter().map(|i| (i + 1) * 10).sum::<u64>();
592            assert_eq!(
593                planned.archives.iter().map(|a| a.archive.file_name.clone()).collect::<Vec<_>>(),
594                expected_files,
595                "{selection}"
596            );
597            assert_eq!(planned.total_download_size, expected_size, "{selection}");
598            assert_eq!(planned.total_output_size, expected_size * 10, "{selection}");
599
600            // The URL, count, and size APIs used by the UI must agree with the download plan.
601            let distance = selection_archive_distance(&selection, manifest.block).unwrap();
602            assert_eq!(
603                manifest.archive_urls_for_distance(component, distance),
604                expected_files
605                    .iter()
606                    .map(|file| format!("https://example.com/{file}"))
607                    .collect::<Vec<_>>(),
608                "{selection}"
609            );
610            assert_eq!(
611                manifest.chunks_for_distance(component, distance),
612                expected_chunks.len() as u64,
613                "{selection}"
614            );
615            assert_eq!(
616                manifest.size_for_distance(component, distance),
617                expected_size,
618                "{selection}"
619            );
620            assert_eq!(
621                manifest.output_size_for_distance(component, distance),
622                expected_size * 10,
623                "{selection}"
624            );
625        }
626    }
627
628    #[test]
629    fn range_selections_do_not_shift_back_to_a_shorter_component_tip() {
630        let manifest = test_manifest_with_component_tip(20, 18);
631        let component = SnapshotComponentType::Transactions;
632        for (selection, expected_files) in [
633            (ComponentSelection::Since(19), vec![]),
634            (ComponentSelection::Distance(2), vec![]),
635            (ComponentSelection::Since(17), vec!["transactions-10-19.tar.zst"]),
636            (ComponentSelection::Distance(4), vec!["transactions-10-19.tar.zst"]),
637        ] {
638            let planned =
639                collect_planned_archives(&manifest, &BTreeMap::from([(component, selection)]))
640                    .unwrap();
641            assert_eq!(
642                planned.archives.iter().map(|a| a.archive.file_name.as_str()).collect::<Vec<_>>(),
643                expected_files,
644                "{selection}"
645            );
646            let expected_size = if expected_files.is_empty() { 0 } else { 20 };
647            assert_eq!(planned.total_download_size, expected_size, "{selection}");
648            assert_eq!(planned.total_output_size, expected_size * 10, "{selection}");
649        }
650    }
651
652    #[test]
653    fn range_selections_preserve_legacy_tip_archives() {
654        // Legacy manifests stored the snapshot height in total_blocks, not the block count.
655        // At an exact chunk boundary, the legacy generator did not emit the next archive.
656        for (block, expected_files, expected_size) in
657            [(24, vec!["transactions-20-29.tar.zst"], 30), (20, vec![], 0)]
658        {
659            let manifest = test_manifest_with_component_tip(block, block);
660            let component = SnapshotComponentType::Transactions;
661            for selection in [ComponentSelection::Since(block), ComponentSelection::Distance(1)] {
662                let planned =
663                    collect_planned_archives(&manifest, &BTreeMap::from([(component, selection)]))
664                        .unwrap();
665                assert_eq!(
666                    planned
667                        .archives
668                        .iter()
669                        .map(|a| a.archive.file_name.as_str())
670                        .collect::<Vec<_>>(),
671                    expected_files,
672                    "{selection}"
673                );
674                assert_eq!(planned.total_download_size, expected_size, "{selection}");
675                assert_eq!(planned.total_output_size, expected_size * 10, "{selection}");
676                assert_eq!(
677                    manifest.chunks_for_distance(component, Some(1)),
678                    expected_files.len() as u64,
679                    "{selection}"
680                );
681            }
682        }
683    }
684}
reth_cli_commands/download/planning.rs

reth_cli_commands/download/
planning.rs
