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(
304                        session.retry_delay(Duration::from_secs(RETRY_BACKOFF_SECS)),
305                    );
306                }
307                continue;
308            }
309        };
310
311        if !quiet && let Some(size) = response.content_length() {
312            info!(target: "reth::cli",
313                url = %url,
314                size = %DownloadProgress::format_size(size),
315                "Streaming archive"
316            );
317        }
318
319        let result = if let Some(progress) = shared {
320            let reader = SharedProgressReader { inner: response, progress: Arc::clone(progress) };
321            extract_archive_raw(reader, format, target_dir, None)
322        } else {
323            let total_size = response.content_length().unwrap_or(0);
324            extract_archive(
325                response,
326                total_size,
327                format,
328                target_dir,
329                session.cancel_token().clone(),
330            )
331        };
332
333        match result {
334            Ok(()) => return Ok(()),
335            Err(error) => {
336                if attempt < MAX_DOWNLOAD_RETRIES {
337                    warn!(target: "reth::cli",
338                        url = %url,
339                        attempt,
340                        max = MAX_DOWNLOAD_RETRIES,
341                        err = %error,
342                        "Streaming extraction failed, retrying"
343                    );
344                }
345                last_error = Some(error);
346                if attempt < MAX_DOWNLOAD_RETRIES {
347                    std::thread::sleep(
348                        session.retry_delay(Duration::from_secs(RETRY_BACKOFF_SECS)),
349                    );
350                }
351            }
352        }
353    }
354
355    Err(last_error.unwrap_or_else(|| {
356        eyre::eyre!("Streaming download failed after {MAX_DOWNLOAD_RETRIES} attempts")
357    }))
358}
359
360/// Resolves a `file://` archive URL to its local path.
361fn archive_file_url_path(url: &str) -> Result<Option<PathBuf>> {
362    let Ok(parsed) = Url::parse(url) else { return Ok(None) };
363    if parsed.scheme() != "file" {
364        return Ok(None)
365    }
366
367    parsed
368        .to_file_path()
369        .map(Some)
370        .map_err(|_| eyre::eyre!("Invalid file:// archive URL path: {url}"))
371}
372
373/// Fetches the snapshot from a remote URL with resume support, then extracts it.
374fn download_and_extract(
375    url: &str,
376    format: CompressionFormat,
377    target_dir: &Path,
378    session: DownloadSession,
379) -> Result<()> {
380    let quiet = session.progress().is_some();
381    let fetcher = ArchiveFetcher::new(url.to_string(), target_dir, session.clone());
382    let DownloadedArchive { path: downloaded_path, size: total_size } = fetcher.download(None)?;
383
384    let file_name =
385        downloaded_path.file_name().map(|f| f.to_string_lossy().to_string()).unwrap_or_default();
386
387    if !quiet {
388        info!(target: "reth::cli",
389            file = %file_name,
390            size = %DownloadProgress::format_size(total_size),
391            "Extracting archive"
392        );
393    }
394    let file = fs::open(&downloaded_path)?;
395
396    if quiet {
397        extract_archive_raw(file, format, target_dir, None)?;
398    } else {
399        extract_archive(file, total_size, format, target_dir, session.cancel_token().clone())?;
400        info!(target: "reth::cli",
401            file = %file_name,
402            "Extraction complete"
403        );
404    }
405
406    fetcher.cleanup_downloaded_files();
407    session.record_archive_output_complete(total_size);
408
409    Ok(())
410}
411
412/// Downloads and extracts a snapshot, blocking until finished.
413///
414/// Supports `file://` URLs for local files and HTTP(S) URLs for remote downloads.
415/// When `resumable` is true, downloads to a `.part` file first with HTTP Range resume
416/// support. Otherwise streams directly into the extractor.
417fn blocking_download_and_extract(
418    url: &str,
419    target_dir: &Path,
420    shared: Option<Arc<SharedProgress>>,
421    resumable: bool,
422    request_limiter: Option<Arc<DownloadRequestLimiter>>,
423    cancel_token: CancellationToken,
424    retry_backoff: Option<Duration>,
425) -> Result<()> {
426    let format = CompressionFormat::from_url(url)?;
427
428    if let Ok(parsed_url) = Url::parse(url) &&
429        parsed_url.scheme() == "file"
430    {
431        let session = DownloadSession::new(shared, request_limiter, cancel_token)
432            .with_retry_backoff(retry_backoff);
433        let file_path = parsed_url
434            .to_file_path()
435            .map_err(|_| eyre::eyre!("Invalid file:// URL path: {}", url))?;
436        let result = extract_from_file(&file_path, format, target_dir);
437        if result.is_ok() {
438            session.record_archive_output_complete(file_path.metadata()?.len());
439        }
440        result
441    } else if let Some(request_limiter) = request_limiter {
442        download_and_extract(
443            url,
444            format,
445            target_dir,
446            DownloadSession::new(shared, Some(request_limiter), cancel_token)
447                .with_retry_backoff(retry_backoff),
448        )
449    } else if resumable {
450        let session =
451            DownloadSession::new(shared, Some(DownloadRequestLimiter::new(1)), cancel_token)
452                .with_retry_backoff(retry_backoff);
453        download_and_extract(url, format, target_dir, session)
454    } else {
455        let session =
456            DownloadSession::new(shared, None, cancel_token).with_retry_backoff(retry_backoff);
457        let result = streaming_download_and_extract(url, format, target_dir, &session);
458        if result.is_ok() {
459            session.record_archive_output_complete(0);
460        }
461        result
462    }
463}
464
465/// Downloads and extracts a snapshot archive asynchronously.
466///
467/// When `shared` is provided, download progress is reported to the shared
468/// counter for aggregated display. Otherwise uses a local progress bar.
469/// When `resumable` is true, uses two-phase download with `.part` files.
470pub(crate) async fn stream_and_extract(
471    url: &str,
472    target_dir: &Path,
473    shared: Option<Arc<SharedProgress>>,
474    resumable: bool,
475    request_limiter: Option<Arc<DownloadRequestLimiter>>,
476    cancel_token: CancellationToken,
477    retry_backoff: Option<Duration>,
478) -> Result<()> {
479    let target_dir = target_dir.to_path_buf();
480    let url = url.to_string();
481    task::spawn_blocking(move || {
482        blocking_download_and_extract(
483            &url,
484            &target_dir,
485            shared,
486            resumable,
487            request_limiter,
488            cancel_token,
489            retry_backoff,
490        )
491    })
492    .await??;
493
494    Ok(())
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500
501    #[test]
502    fn test_compression_format_detection() {
503        assert!(matches!(
504            CompressionFormat::from_url("https://example.com/snapshot.tar.lz4"),
505            Ok(CompressionFormat::Lz4)
506        ));
507        assert!(matches!(
508            CompressionFormat::from_url("https://example.com/snapshot.tar.zst"),
509            Ok(CompressionFormat::Zstd)
510        ));
511        assert!(matches!(
512            CompressionFormat::from_url("file:///path/to/snapshot.tar.lz4"),
513            Ok(CompressionFormat::Lz4)
514        ));
515        assert!(matches!(
516            CompressionFormat::from_url("file:///path/to/snapshot.tar.zst"),
517            Ok(CompressionFormat::Zstd)
518        ));
519        assert!(CompressionFormat::from_url("https://example.com/snapshot.tar.gz").is_err());
520    }
521}
reth_cli_commands/download/extract.rs

reth_cli_commands/download/
extract.rs
