Skip to main content

reth_cli_commands/download/
session.rs

1use super::progress::{DownloadRequestLimiter, SharedProgress};
2use eyre::Result;
3use reth_cli_util::cancellation::CancellationToken;
4use std::{
5    path::{Path, PathBuf},
6    sync::Arc,
7};
8
9/// Shared state for one run of `reth download`.
10#[derive(Clone)]
11pub(crate) struct DownloadSession {
12    /// Shared progress counters for this command, when enabled.
13    progress: Option<Arc<SharedProgress>>,
14    /// Shared limit for concurrent HTTP requests, when enabled.
15    request_limiter: Option<Arc<DownloadRequestLimiter>>,
16    /// Cancellation token shared by the whole command.
17    cancel_token: CancellationToken,
18}
19
20impl DownloadSession {
21    /// Stores the shared progress, request limiter, and cancellation token.
22    pub(crate) fn new(
23        progress: Option<Arc<SharedProgress>>,
24        request_limiter: Option<Arc<DownloadRequestLimiter>>,
25        cancel_token: CancellationToken,
26    ) -> Self {
27        Self { progress, request_limiter, cancel_token }
28    }
29
30    /// Returns the shared progress tracker, if this flow uses one.
31    pub(crate) fn progress(&self) -> Option<&Arc<SharedProgress>> {
32        self.progress.as_ref()
33    }
34
35    /// Returns the shared HTTP request limiter, if this flow uses one.
36    pub(crate) fn request_limiter(&self) -> Option<&Arc<DownloadRequestLimiter>> {
37        self.request_limiter.as_ref()
38    }
39
40    /// Returns the request limiter or errors if the caller needs one.
41    pub(crate) fn require_request_limiter(&self) -> Result<&Arc<DownloadRequestLimiter>> {
42        self.request_limiter().ok_or_else(|| eyre::eyre!("Missing download request limiter"))
43    }
44
45    /// Returns the cancellation token for this command.
46    pub(crate) fn cancel_token(&self) -> &CancellationToken {
47        &self.cancel_token
48    }
49
50    /// Records one archive whose outputs were already reusable on disk.
51    pub(crate) fn record_reused_archive(&self, download_bytes: u64, output_bytes: u64) {
52        if let Some(progress) = self.progress() {
53            progress.record_reused_archive(download_bytes, output_bytes);
54        }
55    }
56
57    /// Records one archive whose extracted outputs fully verified.
58    pub(crate) fn record_archive_output_complete(&self, bytes: u64) {
59        if let Some(progress) = self.progress() {
60            progress.record_archive_output_complete(bytes);
61        }
62    }
63}
64
65/// Paths used while processing one archive, plus the shared download session.
66#[derive(Clone)]
67pub(crate) struct ArchiveProcessContext {
68    /// Directory where extracted output files are written.
69    target_dir: PathBuf,
70    /// Directory used for cached archive downloads, when enabled.
71    cache_dir: Option<PathBuf>,
72    /// Shared command-scoped download state.
73    session: DownloadSession,
74}
75
76impl ArchiveProcessContext {
77    /// Creates the context used while processing modular archives.
78    pub(crate) fn new(
79        target_dir: PathBuf,
80        cache_dir: Option<PathBuf>,
81        session: DownloadSession,
82    ) -> Self {
83        Self { target_dir, cache_dir, session }
84    }
85
86    /// Returns the directory where extracted outputs should be written.
87    pub(crate) fn target_dir(&self) -> &Path {
88        &self.target_dir
89    }
90
91    /// Returns the cache directory for two-phase downloads, if enabled.
92    pub(crate) fn cache_dir(&self) -> Option<&Path> {
93        self.cache_dir.as_deref()
94    }
95
96    /// Returns the shared download session.
97    pub(crate) fn session(&self) -> &DownloadSession {
98        &self.session
99    }
100}