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    /// Directory used for cached archive downloads, when enabled.
85    cache_dir: Option<PathBuf>,
86    /// Shared command-scoped download state.
87    session: DownloadSession,
88}
89
90impl ArchiveProcessContext {
91    /// Creates the context used while processing modular archives.
92    pub(crate) fn new(
93        target_dir: PathBuf,
94        cache_dir: Option<PathBuf>,
95        session: DownloadSession,
96    ) -> Self {
97        Self { target_dir, cache_dir, session }
98    }
99
100    /// Returns the directory where extracted outputs should be written.
101    pub(crate) fn target_dir(&self) -> &Path {
102        &self.target_dir
103    }
104
105    /// Returns the cache directory for two-phase downloads, if enabled.
106    pub(crate) fn cache_dir(&self) -> Option<&Path> {
107        self.cache_dir.as_deref()
108    }
109
110    /// Returns the shared download session.
111    pub(crate) fn session(&self) -> &DownloadSession {
112        &self.session
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn retry_delay_preserves_defaults_or_uses_override() {
122        let session = DownloadSession::new(None, None, CancellationToken::new());
123        for default in [Duration::from_secs(2), Duration::from_secs(5), Duration::from_secs(40)] {
124            assert_eq!(session.retry_delay(default), default);
125            for delay in [Duration::ZERO, Duration::from_millis(250)] {
126                assert_eq!(
127                    session.clone().with_retry_backoff(Some(delay)).retry_delay(default),
128                    delay
129                );
130            }
131        }
132    }
133}
reth_cli_commands/download/session.rs

reth_cli_commands/download/
session.rs
