pipedash_core/application/
refresh_manager.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use serde::{
5    Deserialize,
6    Serialize,
7};
8use tokio::sync::{
9    mpsc,
10    watch,
11    Mutex,
12};
13use tokio::task::{
14    JoinHandle,
15    JoinSet,
16};
17use tokio::time::{
18    interval,
19    Instant,
20    MissedTickBehavior,
21};
22use tokio_util::sync::CancellationToken;
23
24use super::fetch_scheduler::{
25    ClaimKind,
26    FetchClaim,
27};
28use super::services::metrics_service::MetricsService;
29use super::services::pipeline_service::PipelineService;
30
31const TICK: Duration = Duration::from_secs(5);
32const METRICS_CLEANUP_INTERVAL: Duration = Duration::from_secs(6 * 3600);
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(rename_all = "lowercase")]
36pub enum RefreshMode {
37    Active,
38    Idle,
39}
40
41impl RefreshMode {
42    pub fn as_str(&self) -> &'static str {
43        match self {
44            RefreshMode::Active => "active",
45            RefreshMode::Idle => "idle",
46        }
47    }
48}
49
50impl std::str::FromStr for RefreshMode {
51    type Err = std::convert::Infallible;
52
53    fn from_str(s: &str) -> Result<Self, Self::Err> {
54        Ok(match s.to_lowercase().as_str() {
55            "active" => RefreshMode::Active,
56            _ => RefreshMode::Idle,
57        })
58    }
59}
60
61pub enum RefreshCommand {
62    FetchNow { claims: Vec<FetchClaim> },
63}
64
65#[derive(Clone)]
66pub struct RefreshDispatcher(mpsc::UnboundedSender<RefreshCommand>);
67
68impl RefreshDispatcher {
69    pub fn channel() -> (Self, mpsc::UnboundedReceiver<RefreshCommand>) {
70        let (tx, rx) = mpsc::unbounded_channel();
71
72        (Self(tx), rx)
73    }
74
75    pub fn fetch_now(&self, claims: Vec<FetchClaim>) {
76        let _ = self.0.send(RefreshCommand::FetchNow { claims });
77    }
78}
79
80pub struct RefreshManager {
81    pipeline_service: Arc<PipelineService>,
82    metrics_service: Option<Arc<MetricsService>>,
83    mode: watch::Sender<RefreshMode>,
84    commands: Mutex<Option<mpsc::UnboundedReceiver<RefreshCommand>>>,
85    cancel: CancellationToken,
86    task: Mutex<Option<JoinHandle<()>>>,
87}
88
89impl RefreshManager {
90    pub fn new(
91        pipeline_service: Arc<PipelineService>, metrics_service: Option<Arc<MetricsService>>,
92        mode: watch::Sender<RefreshMode>, commands: mpsc::UnboundedReceiver<RefreshCommand>,
93    ) -> Self {
94        Self {
95            pipeline_service,
96            metrics_service,
97            mode,
98            commands: Mutex::new(Some(commands)),
99            cancel: CancellationToken::new(),
100            task: Mutex::new(None),
101        }
102    }
103
104    pub async fn start(&self) {
105        let mut task = self.task.lock().await;
106        if task.is_some() || self.cancel.is_cancelled() {
107            return;
108        }
109        let Some(commands) = self.commands.lock().await.take() else {
110            return;
111        };
112
113        *task = Some(tokio::spawn(run(
114            Arc::clone(&self.pipeline_service),
115            self.metrics_service.clone(),
116            self.mode.subscribe(),
117            self.cancel.clone(),
118            commands,
119        )));
120    }
121
122    pub async fn stop(&self) {
123        self.cancel.cancel();
124        let task = self.task.lock().await.take();
125        if let Some(task) = task {
126            if let Err(e) = task.await {
127                tracing::error!(error = %e, "RefreshManager task failed");
128            }
129        }
130    }
131
132    pub fn get_mode(&self) -> RefreshMode {
133        *self.mode.borrow()
134    }
135
136    pub fn set_mode(&self, mode: RefreshMode) {
137        if self.mode.send_if_modified(|current| {
138            let changed = *current != mode;
139            *current = mode;
140            changed
141        }) {
142            tracing::debug!(mode = mode.as_str(), "Refresh mode changed");
143        }
144    }
145}
146
147impl Drop for RefreshManager {
148    fn drop(&mut self) {
149        self.cancel.cancel();
150    }
151}
152
153async fn run(
154    pipeline_service: Arc<PipelineService>, metrics_service: Option<Arc<MetricsService>>,
155    mode: watch::Receiver<RefreshMode>, cancel: CancellationToken,
156    mut commands: mpsc::UnboundedReceiver<RefreshCommand>,
157) {
158    let mut tick = interval(TICK);
159    tick.set_missed_tick_behavior(MissedTickBehavior::Skip);
160    let mut attempts: JoinSet<()> = JoinSet::new();
161    let mut last_cleanup: Option<Instant> = None;
162
163    tracing::info!("RefreshManager started");
164
165    loop {
166        tokio::select! {
167            _ = cancel.cancelled() => break,
168            _ = tick.tick() => {
169                let active = *mode.borrow() == RefreshMode::Active;
170                if active {
171                    for claim in pipeline_service.claim_due(Instant::now()).await {
172                        spawn_fetch(&mut attempts, &pipeline_service, claim, ClaimKind::Scheduled);
173                    }
174                }
175                if let Some(metrics_service) = &metrics_service {
176                    if last_cleanup.is_none_or(|last| last.elapsed() >= METRICS_CLEANUP_INTERVAL) {
177                        last_cleanup = Some(Instant::now());
178                        let metrics_service = Arc::clone(metrics_service);
179                        attempts.spawn(async move {
180                            match metrics_service.cleanup_old_metrics().await {
181                                Ok(deleted) => {
182                                    tracing::info!(deleted = deleted, "Metrics cleanup complete");
183                                }
184                                Err(e) => {
185                                    tracing::error!(error = %e, "Metrics cleanup failed");
186                                }
187                            }
188                        });
189                    }
190                }
191            }
192            Some(command) = commands.recv() => match command {
193                RefreshCommand::FetchNow { claims } => {
194                    for claim in claims {
195                        spawn_fetch(&mut attempts, &pipeline_service, claim, ClaimKind::Forced);
196                    }
197                }
198            },
199            Some(finished) = attempts.join_next(), if !attempts.is_empty() => {
200                if let Err(e) = finished {
201                    tracing::error!(error = %e, "Provider fetch task failed");
202                }
203            }
204        }
205    }
206
207    attempts.shutdown().await;
208    tracing::info!("RefreshManager stopped");
209}
210
211fn spawn_fetch(
212    attempts: &mut JoinSet<()>, pipeline_service: &Arc<PipelineService>, claim: FetchClaim,
213    kind: ClaimKind,
214) {
215    let pipeline_service = Arc::clone(pipeline_service);
216    attempts.spawn(async move {
217        pipeline_service.fetch_provider(claim, kind).await;
218    });
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use crate::application::test_support::{
225        Gate,
226        Harness,
227        StubProvider,
228    };
229    use crate::domain::RefreshState;
230    use crate::event::CoreEvent;
231
232    #[tokio::test]
233    async fn the_first_tick_fetches_each_due_provider_once() {
234        let h = Harness::new().await;
235        let (_, first) = h.add_stub(StubProvider::new).await;
236        let (_, second) = h.add_stub(StubProvider::new).await;
237        h.provider_service.load_all_providers().await.unwrap();
238
239        h.start().await;
240        h.start().await;
241        h.bus
242            .wait_for("both providers", |events| {
243                events
244                    .iter()
245                    .filter(|event| matches!(event, CoreEvent::PipelinesUpdated { .. }))
246                    .count()
247                    == 2
248            })
249            .await;
250        tokio::time::sleep(Duration::from_millis(200)).await;
251
252        assert_eq!((first.fetch_count(), second.fetch_count()), (1, 1));
253        assert_eq!(h.bus.count("pipelines-updated"), 2);
254    }
255
256    #[tokio::test]
257    async fn providers_missing_from_the_startup_load_are_fetched_on_schedule() {
258        let h = Harness::new().await;
259        let (_, stub) = h.add_stub(StubProvider::new).await;
260
261        h.start().await;
262        h.bus
263            .wait_for("the unloaded provider", |events| {
264                events
265                    .iter()
266                    .any(|event| matches!(event, CoreEvent::PipelinesUpdated { .. }))
267            })
268            .await;
269
270        assert_eq!(stub.fetch_count(), 1);
271    }
272
273    #[tokio::test]
274    async fn idle_mode_keeps_scheduled_fetches_from_starting() {
275        let h = Harness::new().await;
276        let (_, stub) = h.add_stub(StubProvider::new).await;
277        h.provider_service.load_all_providers().await.unwrap();
278        h.refresh_manager.set_mode(RefreshMode::Idle);
279        assert_eq!(h.refresh_manager.get_mode(), RefreshMode::Idle);
280
281        h.start().await;
282        tokio::time::sleep(Duration::from_millis(300)).await;
283
284        assert_eq!(stub.fetch_count(), 0);
285        h.refresh_manager.set_mode(RefreshMode::Active);
286        assert_eq!(h.refresh_manager.get_mode(), RefreshMode::Active);
287    }
288
289    #[tokio::test]
290    async fn stopping_cancels_running_attempts_and_their_waiters() {
291        let h = Harness::new().await;
292        h.start().await;
293        let gate = Gate::new();
294        let (id, stub) = h.add_stub(|id| StubProvider::new(id).gated(&gate)).await;
295        let waiting = tokio::spawn({
296            let service = Arc::clone(&h.pipeline_service);
297            async move { service.fetch_now(&[id], Duration::from_secs(10)).await }
298        });
299        gate.arrived().await;
300
301        h.refresh_manager.stop().await;
302        let results = waiting.await.unwrap();
303        h.refresh_manager.start().await;
304
305        assert_eq!(results[0].state, RefreshState::Pending);
306        let view = h.scheduler.view(id, Instant::now()).unwrap();
307        assert_eq!((view.failures, view.in_flight), (0, false));
308        assert_eq!(stub.fetch_count(), 1);
309        assert_eq!(h.bus.count("pipelines-updated"), 0);
310        assert!(h.refresh_manager.task.lock().await.is_none());
311    }
312
313    #[tokio::test]
314    async fn a_stopped_manager_hands_new_claims_back_as_cancelled() {
315        let h = Harness::new().await;
316        h.start().await;
317        h.refresh_manager.stop().await;
318        let (id, _) = h.add_stub(StubProvider::new).await;
319
320        let results = h
321            .pipeline_service
322            .fetch_now(&[id], Duration::from_millis(300))
323            .await;
324
325        assert_eq!(results[0].state, RefreshState::Pending);
326    }
327}
pipedash_core/application/refresh_manager.rs

pipedash_core/application/
refresh_manager.rs
