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    time::Duration,
8};
9
10/// Shared state for one run of `reth download`.
11#[derive(Clone)]
12pub(crate) struct DownloadSession {
13    /// Shared progress counters for this command, when enabled.
14    progress: Option<Arc<SharedProgress>>,
15    /// Shared limit for concurrent HTTP requests, when enabled.
16    request_limiter: Option<Arc<DownloadRequestLimiter>>,
17    /// Cancellation token shared by the whole command.
18    cancel_token: CancellationToken,
19    /// Optional fixed delay overriding the default retry backoff.
20    retry_backoff: Option<Duration>,
21}
22
23impl DownloadSession {
24    /// Stores the shared progress, request limiter, and cancellation token.
25    pub(crate) fn new(
26        progress: Option<Arc<SharedProgress>>,
27        request_limiter: Option<Arc<DownloadRequestLimiter>>,
28        cancel_token: CancellationToken,
29    ) -> Self {
30        Self { progress, request_limiter, cancel_token, retry_backoff: None }
31    }
32
33    /// Overrides the delay between retry attempts for this download.
34    pub(crate) const fn with_retry_backoff(mut self, retry_backoff: Option<Duration>) -> Self {
35        self.retry_backoff = retry_backoff;
36        self
37    }
38
39    /// Returns the configured delay, preserving the caller's default when no override is set.
40    pub(crate) fn retry_delay(&self, default: Duration) -> Duration {
41        self.retry_backoff.unwrap_or(default)
42    }
43
44    /// Returns the shared progress tracker, if this flow uses one.
45    pub(crate) fn progress(&self) -> Option<&Arc<SharedProgress>> {
46        self.progress.as_ref()
47    }
48
49    /// Returns the shared HTTP request limiter, if this flow uses one.
50    pub(crate) fn request_limiter(&self) -> Option<&Arc<DownloadRequestLimiter>> {
51        self.request_limiter.as_ref()
52    }
53
54    /// Returns the request limiter or errors if the caller needs one.
55    pub(crate) fn require_request_limiter(&self) -> Result<&Arc<DownloadRequestLimiter>> {
56        self.request_limiter().ok_or_else(|| eyre::eyre!("Missing download request limiter"))
57    }
58
59    /// Returns the cancellation token for this command.
60    pub(crate) fn cancel_token(&self) -> &CancellationToken {
61        &self.cancel_token
62    }
63
64    /// Records one archive whose outputs were already reusable on disk.
65    pub(crate) fn record_reused_archive(&self, download_bytes: u64, output_bytes: u64) {
66        if let Some(progress) = self.progress() {
67            progress.record_reused_archive(download_bytes, output_bytes);
68        }
69    }
70
71    /// Records one archive whose extracted outputs fully verified.
72    pub(crate) fn record_archive_output_complete(&self, bytes: u64) {
73        if let Some(progress) = self.progress() {
74            progress.record_archive_output_complete(bytes);
75        }
76    }
77}
78
79/// Paths used while processing one archive, plus the shared download session.
80#[derive(Clone)]
81pub(crate) struct ArchiveProcessContext {
82    /// Directory where extracted output files are written.
83    target_dir: PathBuf,
84    /// Custom location of static file outputs.
85    static_files_dir: Option<PathBuf>,
86    /// Directory used for cached archive downloads, when enabled.
87    cache_dir: Option<PathBuf>,
88    /// Shared command-scoped download state.
89    session: DownloadSession,
90}
91
92impl ArchiveProcessContext {
93    /// Creates the context used while processing modular archives.
94    pub(crate) fn new(
95        target_dir: PathBuf,
96        static_files_dir: Option<PathBuf>,
97        cache_dir: Option<PathBuf>,
98        session: DownloadSession,
99    ) -> Self {
100        Self { target_dir, static_files_dir, cache_dir, session }
101    }
102
103    /// Returns the directory where extracted outputs should be written.
104    pub(crate) fn target_dir(&self) -> &Path {
105        &self.target_dir
106    }
107
108    /// Returns the custom static files directory, if configured.
109    pub(crate) fn static_files_dir(&self) -> Option<&Path> {
110        self.static_files_dir.as_deref()
111    }
112
113    /// Returns the cache directory for two-phase downloads, if enabled.
114    pub(crate) fn cache_dir(&self) -> Option<&Path> {
115        self.cache_dir.as_deref()
116    }
117
118    /// Returns the shared download session.
119    pub(crate) fn session(&self) -> &DownloadSession {
120        &self.session
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn retry_delay_preserves_defaults_or_uses_override() {
130        let session = DownloadSession::new(None, None, CancellationToken::new());
131        for default in [Duration::from_secs(2), Duration::from_secs(5), Duration::from_secs(40)] {
132            assert_eq!(session.retry_delay(default), default);
133            for delay in [Duration::ZERO, Duration::from_millis(250)] {
134                assert_eq!(
135                    session.clone().with_retry_backoff(Some(delay)).retry_delay(default),
136                    delay
137                );
138            }
139        }
140    }
141}
reth_cli_commands/download/session.rs

reth_cli_commands/download/
session.rs
