reth_cli_commands/download/
extract.rs

1use super::{
2    fetch::{ArchiveFetcher, DownloadedArchive},
3    progress::{
4        ArchiveExtractionProgress, ArchiveExtractionProgressHandle, DownloadProgress,
5        DownloadRequestLimiter, ProgressReader, SharedProgress, SharedProgressReader,
6    },
7    session::DownloadSession,
8    MAX_DOWNLOAD_RETRIES, RETRY_BACKOFF_SECS,
9};
10use eyre::{Result, WrapErr};
11use lz4::Decoder;
12use reqwest::blocking::Client as BlockingClient;
13use reth_cli_util::cancellation::CancellationToken;
14use reth_fs_util as fs;
15use std::{
16    io::Read,
17    path::{Component, Path, PathBuf},
18    sync::{
19        atomic::{AtomicBool, Ordering},
20        Arc,
21    },
22    thread,
23    time::{Duration, Instant},
24};
25use tar::Archive;
26use tokio::task;
27use tracing::{info, warn};
28use url::Url;
29use zstd::stream::read::Decoder as ZstdDecoder;
30
31const EXTENSION_TAR_LZ4: &str = ".tar.lz4";
32const EXTENSION_TAR_ZSTD: &str = ".tar.zst";
33const STREAMING_EXTRACTION_PROGRESS_MIN_FILE_SIZE: u64 = 64 * 1024 * 1024;
34const EXTRACTION_PROGRESS_POLL_INTERVAL: Duration = Duration::from_millis(100);
35
36/// Supported compression formats for snapshots
37#[derive(Debug, Clone, Copy)]
38pub(crate) enum CompressionFormat {
39    /// LZ4-compressed tar archive.
40    Lz4,
41    /// Zstandard-compressed tar archive.
42    Zstd,
43}
44
45impl CompressionFormat {
46    /// Detect compression format from file extension
47    pub(crate) fn from_url(url: &str) -> Result<Self> {
48        let path =
49            Url::parse(url).map(|u| u.path().to_string()).unwrap_or_else(|_| url.to_string());
50
51        if path.ends_with(EXTENSION_TAR_LZ4) {
52            Ok(Self::Lz4)
53        } else if path.ends_with(EXTENSION_TAR_ZSTD) {
54            Ok(Self::Zstd)
55        } else {
56            Err(eyre::eyre!(
57                "Unsupported file format. Expected .tar.lz4 or .tar.zst, got: {}",
58                path
59            ))
60        }
61    }
62}
63
64/// Extracts a compressed tar archive to the target directory with progress tracking.
65fn extract_archive<R: Read>(
66    reader: R,
67    total_size: u64,
68    format: CompressionFormat,
69    target_dir: &Path,
70    cancel_token: CancellationToken,
71) -> Result<()> {
72    let progress_reader = ProgressReader::new(reader, total_size, cancel_token);
73
74    match format {
75        CompressionFormat::Lz4 => {
76            let decoder = Decoder::new(progress_reader)?;
77            Archive::new(decoder).unpack(target_dir)?;
78        }
79        CompressionFormat::Zstd => {
80            let decoder = ZstdDecoder::new(progress_reader)?;
81            Archive::new(decoder).unpack(target_dir)?;
82        }
83    }
84
85    println!();
86    Ok(())
87}
88
89/// Extracts a compressed tar archive without progress tracking.
90pub(crate) fn extract_archive_raw<R: Read>(
91    reader: R,
92    format: CompressionFormat,
93    target_dir: &Path,
94    progress: Option<&mut ArchiveExtractionProgress>,
95) -> Result<()> {
96    match format {
97        CompressionFormat::Lz4 => {
98            unpack_archive(Archive::new(Decoder::new(reader)?), target_dir, progress)?;
99        }
100        CompressionFormat::Zstd => {
101            unpack_archive(Archive::new(ZstdDecoder::new(reader)?), target_dir, progress)?;
102        }
103    }
104
105    Ok(())
106}
107
108fn unpack_archive<R: Read>(
109    mut archive: Archive<R>,
110    target_dir: &Path,
111    mut progress: Option<&mut ArchiveExtractionProgress>,
112) -> Result<()> {
113    let entries = archive.entries().wrap_err_with(|| {
114        format!("failed to read archive entries for `{}`", target_dir.display())
115    })?;
116
117    for entry in entries {
118        let mut entry = entry.wrap_err_with(|| {
119            format!("failed to read archive entry for `{}`", target_dir.display())
120        })?;
121        extract_entry_with_progress(&mut entry, target_dir, progress.as_deref_mut())?;
122    }
123
124    Ok(())
125}
126
127fn extract_entry_with_progress<R: Read>(
128    entry: &mut tar::Entry<'_, R>,
129    target_dir: &Path,
130    progress: Option<&mut ArchiveExtractionProgress>,
131) -> Result<()> {
132    let size = entry.header().entry_size().unwrap_or(0);
133    let entry_type = entry.header().entry_type();
134
135    if !entry_type.is_file() || size == 0 {
136        entry.unpack_in(target_dir).wrap_err_with(|| {
137            format!("failed to extract archive into `{}`", target_dir.display())
138        })?;
139        return Ok(())
140    }
141
142    if size < STREAMING_EXTRACTION_PROGRESS_MIN_FILE_SIZE {
143        entry.unpack_in(target_dir).wrap_err_with(|| {
144            format!("failed to extract archive into `{}`", target_dir.display())
145        })?;
146        if let Some(progress) = progress {
147            progress.record_extracted(size);
148        }
149        return Ok(())
150    }
151
152    let Some(progress_handle) = progress.as_ref().and_then(|progress| progress.handle()) else {
153        entry.unpack_in(target_dir).wrap_err_with(|| {
154            format!("failed to extract archive into `{}`", target_dir.display())
155        })?;
156        return Ok(())
157    };
158
159    let Some(entry_path) = entry_destination_path(entry, target_dir)? else {
160        entry.unpack_in(target_dir).wrap_err_with(|| {
161            format!("failed to extract archive into `{}`", target_dir.display())
162        })?;
163        return Ok(())
164    };
165
166    let stop = Arc::new(AtomicBool::new(false));
167    let monitor = spawn_extraction_progress_monitor(entry_path, progress_handle, Arc::clone(&stop));
168    let unpack_result = entry
169        .unpack_in(target_dir)
170        .wrap_err_with(|| format!("failed to extract archive into `{}`", target_dir.display()));
171    stop.store(true, Ordering::Relaxed);
172
173    let monitor_result = monitor.join();
174    unpack_result?;
175
176    monitor_result.map_err(|_| eyre::eyre!("extraction progress monitor panicked"))?;
177    Ok(())
178}
179
180fn entry_destination_path<R: Read>(
181    entry: &tar::Entry<'_, R>,
182    target_dir: &Path,
183) -> Result<Option<PathBuf>> {
184    let mut file_dst = target_dir.to_path_buf();
185    let path = entry.path().wrap_err("invalid path in archive entry")?;
186
187    for part in path.components() {
188        match part {
189            Component::Prefix(..) | Component::RootDir | Component::CurDir => continue,
190            Component::ParentDir => return Ok(None),
191            Component::Normal(part) => file_dst.push(part),
192        }
193    }
194
195    if file_dst == target_dir {
196        return Ok(None)
197    }
198
199    Ok(Some(file_dst))
200}
201
202fn spawn_extraction_progress_monitor(
203    entry_path: PathBuf,
204    progress: ArchiveExtractionProgressHandle,
205    stop: Arc<AtomicBool>,
206) -> thread::JoinHandle<()> {
207    thread::spawn(move || {
208        let mut extracted = 0_u64;
209
210        loop {
211            record_extracted_file_bytes(&entry_path, &progress, &mut extracted);
212            if stop.load(Ordering::Relaxed) {
213                break;
214            }
215            thread::sleep(EXTRACTION_PROGRESS_POLL_INTERVAL);
216        }
217    })
218}
219
220fn record_extracted_file_bytes(
221    entry_path: &Path,
222    progress: &ArchiveExtractionProgressHandle,
223    extracted: &mut u64,
224) {
225    let Ok(meta) = fs::metadata(entry_path) else { return };
226    let len = meta.len();
227    if len > *extracted {
228        progress.record_extracted(len - *extracted);
229        *extracted = len;
230    }
231}
232
233/// Extracts a snapshot from a local file.
234fn extract_from_file(path: &Path, format: CompressionFormat, target_dir: &Path) -> Result<()> {
235    let file = std::fs::File::open(path)?;
236    let total_size = file.metadata()?.len();
237    info!(target: "reth::cli",
238        file = %path.display(),
239        size = %DownloadProgress::format_size(total_size),
240        "Extracting local archive"
241    );
242    let start = Instant::now();
243    extract_archive(file, total_size, format, target_dir, CancellationToken::new())?;
244    info!(target: "reth::cli",
245        file = %path.display(),
246        elapsed = %DownloadProgress::format_duration(start.elapsed()),
247        "Local extraction complete"
248    );
249    Ok(())
250}
251
252/// Streams a remote archive directly into the extractor without writing to disk.
253///
254/// On failure, retries from scratch up to [`MAX_DOWNLOAD_RETRIES`] times.
255pub(crate) fn streaming_download_and_extract(
256    url: &str,
257    format: CompressionFormat,
258    target_dir: &Path,
259    session: &DownloadSession,
260) -> Result<()> {
261    if let Some(path) = archive_file_url_path(url)? {
262        let size = path.metadata()?.len();
263        extract_from_file(&path, format, target_dir)?;
264        session.record_archive_output_complete(size);
265        return Ok(())
266    }
267
268    let shared = session.progress();
269    let quiet = session.progress().is_some();
270    let mut last_error: Option<eyre::Error> = None;
271
272    for attempt in 1..=MAX_DOWNLOAD_RETRIES {
273        if attempt > 1 {
274            info!(target: "reth::cli",
275                url = %url,
276                attempt,
277                max = MAX_DOWNLOAD_RETRIES,
278                "Retrying streaming download from scratch"
279            );
280        }
281
282        let client = BlockingClient::builder().connect_timeout(Duration::from_secs(30)).build()?;
283        let _request_permit = session
284            .request_limiter()
285            .map(|limiter| limiter.acquire(session.progress(), session.cancel_token()))
286            .transpose()?;
287
288        let response = match client.get(url).send().and_then(|r| r.error_for_status()) {
289            Ok(r) => r,
290            Err(error) => {
291                let err = eyre::Error::from(error);
292                if attempt < MAX_DOWNLOAD_RETRIES {
293                    warn!(target: "reth::cli",
294                        url = %url,
295                        attempt,
296                        max = MAX_DOWNLOAD_RETRIES,
297                        err = %err,
298                        "Streaming request failed, retrying"
299                    );
300                }
301                last_error = Some(err);
302                if attempt < MAX_DOWNLOAD_RETRIES {
303                    std::thread::sleep(Duration::from_secs(RETRY_BACKOFF_SECS));
304                }
305                continue;
306            }
307        };
308
309        if !quiet && let Some(size) = response.content_length() {
310            info!(target: "reth::cli",
311                url = %url,
312                size = %DownloadProgress::format_size(size),
313                "Streaming archive"
314            );
315        }
316
317        let result = if let Some(progress) = shared {
318            let reader = SharedProgressReader { inner: response, progress: Arc::clone(progress) };
319            extract_archive_raw(reader, format, target_dir, None)
320        } else {
321            let total_size = response.content_length().unwrap_or(0);
322            extract_archive(
323                response,
324                total_size,
325                format,
326                target_dir,
327                session.cancel_token().clone(),
328            )
329        };
330
331        match result {
332            Ok(()) => return Ok(()),
333            Err(error) => {
334                if attempt < MAX_DOWNLOAD_RETRIES {
335                    warn!(target: "reth::cli",
336                        url = %url,
337                        attempt,
338                        max = MAX_DOWNLOAD_RETRIES,
339                        err = %error,
340                        "Streaming extraction failed, retrying"
341                    );
342                }
343                last_error = Some(error);
344                if attempt < MAX_DOWNLOAD_RETRIES {
345                    std::thread::sleep(Duration::from_secs(RETRY_BACKOFF_SECS));
346                }
347            }
348        }
349    }
350
351    Err(last_error.unwrap_or_else(|| {
352        eyre::eyre!("Streaming download failed after {MAX_DOWNLOAD_RETRIES} attempts")
353    }))
354}
355
356/// Resolves a `file://` archive URL to its local path.
357fn archive_file_url_path(url: &str) -> Result<Option<PathBuf>> {
358    let Ok(parsed) = Url::parse(url) else { return Ok(None) };
359    if parsed.scheme() != "file" {
360        return Ok(None)
361    }
362
363    parsed
364        .to_file_path()
365        .map(Some)
366        .map_err(|_| eyre::eyre!("Invalid file:// archive URL path: {url}"))
367}
368
369/// Fetches the snapshot from a remote URL with resume support, then extracts it.
370fn download_and_extract(
371    url: &str,
372    format: CompressionFormat,
373    target_dir: &Path,
374    session: DownloadSession,
375) -> Result<()> {
376    let quiet = session.progress().is_some();
377    let fetcher = ArchiveFetcher::new(url.to_string(), target_dir, session.clone());
378    let DownloadedArchive { path: downloaded_path, size: total_size } = fetcher.download(None)?;
379
380    let file_name =
381        downloaded_path.file_name().map(|f| f.to_string_lossy().to_string()).unwrap_or_default();
382
383    if !quiet {
384        info!(target: "reth::cli",
385            file = %file_name,
386            size = %DownloadProgress::format_size(total_size),
387            "Extracting archive"
388        );
389    }
390    let file = fs::open(&downloaded_path)?;
391
392    if quiet {
393        extract_archive_raw(file, format, target_dir, None)?;
394    } else {
395        extract_archive(file, total_size, format, target_dir, session.cancel_token().clone())?;
396        info!(target: "reth::cli",
397            file = %file_name,
398            "Extraction complete"
399        );
400    }
401
402    fetcher.cleanup_downloaded_files();
403    session.record_archive_output_complete(total_size);
404
405    Ok(())
406}
407
408/// Downloads and extracts a snapshot, blocking until finished.
409///
410/// Supports `file://` URLs for local files and HTTP(S) URLs for remote downloads.
411/// When `resumable` is true, downloads to a `.part` file first with HTTP Range resume
412/// support. Otherwise streams directly into the extractor.
413fn blocking_download_and_extract(
414    url: &str,
415    target_dir: &Path,
416    shared: Option<Arc<SharedProgress>>,
417    resumable: bool,
418    request_limiter: Option<Arc<DownloadRequestLimiter>>,
419    cancel_token: CancellationToken,
420) -> Result<()> {
421    let format = CompressionFormat::from_url(url)?;
422
423    if let Ok(parsed_url) = Url::parse(url) &&
424        parsed_url.scheme() == "file"
425    {
426        let session = DownloadSession::new(shared, request_limiter, cancel_token);
427        let file_path = parsed_url
428            .to_file_path()
429            .map_err(|_| eyre::eyre!("Invalid file:// URL path: {}", url))?;
430        let result = extract_from_file(&file_path, format, target_dir);
431        if result.is_ok() {
432            session.record_archive_output_complete(file_path.metadata()?.len());
433        }
434        result
435    } else if let Some(request_limiter) = request_limiter {
436        download_and_extract(
437            url,
438            format,
439            target_dir,
440            DownloadSession::new(shared, Some(request_limiter), cancel_token),
441        )
442    } else if resumable {
443        let session =
444            DownloadSession::new(shared, Some(DownloadRequestLimiter::new(1)), cancel_token);
445        download_and_extract(url, format, target_dir, session)
446    } else {
447        let session = DownloadSession::new(shared, None, cancel_token);
448        let result = streaming_download_and_extract(url, format, target_dir, &session);
449        if result.is_ok() {
450            session.record_archive_output_complete(0);
451        }
452        result
453    }
454}
455
456/// Downloads and extracts a snapshot archive asynchronously.
457///
458/// When `shared` is provided, download progress is reported to the shared
459/// counter for aggregated display. Otherwise uses a local progress bar.
460/// When `resumable` is true, uses two-phase download with `.part` files.
461pub(crate) async fn stream_and_extract(
462    url: &str,
463    target_dir: &Path,
464    shared: Option<Arc<SharedProgress>>,
465    resumable: bool,
466    request_limiter: Option<Arc<DownloadRequestLimiter>>,
467    cancel_token: CancellationToken,
468) -> Result<()> {
469    let target_dir = target_dir.to_path_buf();
470    let url = url.to_string();
471    task::spawn_blocking(move || {
472        blocking_download_and_extract(
473            &url,
474            &target_dir,
475            shared,
476            resumable,
477            request_limiter,
478            cancel_token,
479        )
480    })
481    .await??;
482
483    Ok(())
484}
485
486#[cfg(test)]
487mod tests {
488    use super::*;
489
490    #[test]
491    fn test_compression_format_detection() {
492        assert!(matches!(
493            CompressionFormat::from_url("https://example.com/snapshot.tar.lz4"),
494            Ok(CompressionFormat::Lz4)
495        ));
496        assert!(matches!(
497            CompressionFormat::from_url("https://example.com/snapshot.tar.zst"),
498            Ok(CompressionFormat::Zstd)
499        ));
500        assert!(matches!(
501            CompressionFormat::from_url("file:///path/to/snapshot.tar.lz4"),
502            Ok(CompressionFormat::Lz4)
503        ));
504        assert!(matches!(
505            CompressionFormat::from_url("file:///path/to/snapshot.tar.zst"),
506            Ok(CompressionFormat::Zstd)
507        ));
508        assert!(CompressionFormat::from_url("https://example.com/snapshot.tar.gz").is_err());
509    }
510}
reth_cli_commands/download/extract.rs

reth_cli_commands/download/
extract.rs
