s2_lite/backend/bgtasks/
mod.rs

1use std::{error::Error, future::Future, pin::Pin, time::Duration};
2
3use tokio::{sync::broadcast, time::Instant};
4use tracing::{info, warn};
5
6use crate::backend::Backend;
7
8mod basin_deletion;
9mod stream_doe;
10mod stream_trim;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub(super) enum BgtaskTrigger {
14    BasinDeletion,
15    StreamDeleteOnEmpty,
16    StreamTrim,
17}
18
19pub fn spawn(backend: &Backend) {
20    spawn_bgtask(
21        "stream-trim",
22        Duration::from_secs(60),
23        &[BgtaskTrigger::StreamTrim],
24        backend.bgtask_trigger_subscribe(),
25        move |backend| backend.clone().tick_stream_trim(),
26        backend.clone(),
27    );
28    spawn_bgtask(
29        "stream-delete-on-empty",
30        Duration::from_secs(60),
31        &[BgtaskTrigger::StreamDeleteOnEmpty],
32        backend.bgtask_trigger_subscribe(),
33        move |backend| backend.clone().tick_stream_doe(),
34        backend.clone(),
35    );
36    spawn_bgtask(
37        "basin-deletion",
38        Duration::from_secs(60),
39        &[BgtaskTrigger::BasinDeletion],
40        backend.bgtask_trigger_subscribe(),
41        move |backend| backend.clone().tick_basin_deletion(),
42        backend.clone(),
43    );
44}
45
46fn spawn_bgtask<Tick, Fut, E>(
47    name: &'static str,
48    interval: Duration,
49    triggers: &'static [BgtaskTrigger],
50    mut trigger_rx: broadcast::Receiver<BgtaskTrigger>,
51    tick: Tick,
52    backend: Backend,
53) where
54    Tick: Fn(&Backend) -> Fut + Send + Sync + 'static,
55    Fut: Future<Output = Result<bool, E>> + Send,
56    E: Error + Send + Sync + 'static,
57{
58    tokio::spawn(async move {
59        let sleep = tokio::time::sleep(jittered_delay(interval));
60        tokio::pin!(sleep);
61        let reset_sleep = |sleep: &mut Pin<&mut tokio::time::Sleep>| {
62            sleep
63                .as_mut()
64                .reset(Instant::now() + jittered_delay(interval));
65        };
66        loop {
67            tokio::select! {
68                _ = &mut sleep => {
69                    run_tick(name, &tick, &backend).await;
70                    reset_sleep(&mut sleep);
71                }
72                res = trigger_rx.recv() => {
73                    match res {
74                        Ok(trigger)  => {
75                            if triggers.contains(&trigger) {
76                                run_tick(name, &tick, &backend).await;
77                                reset_sleep(&mut sleep);
78                            }
79                        }
80                        Err(broadcast::error::RecvError::Lagged(skipped)) => {
81                            warn!(
82                                task = name,
83                                skipped,
84                                "bgtask trigger channel lagged, running tick immediately"
85                            );
86                            run_tick(name, &tick, &backend).await;
87                            reset_sleep(&mut sleep);
88                        }
89                        Err(broadcast::error::RecvError::Closed) => {
90                            info!(task = name, "bgtask trigger channel closed, exiting");
91                            break;
92                        }
93                    }
94                }
95            }
96        }
97    });
98}
99
100fn jittered_delay(interval: Duration) -> Duration {
101    if interval.is_zero() {
102        return interval;
103    }
104    let max_jitter = interval / 10;
105    let max_ms = max_jitter.as_millis() as i64;
106    if max_ms == 0 {
107        return interval;
108    }
109    let jitter_ms = rand::random_range(-max_ms..=max_ms);
110    if jitter_ms >= 0 {
111        interval + Duration::from_millis(jitter_ms as u64)
112    } else {
113        interval - Duration::from_millis((-jitter_ms) as u64)
114    }
115}
116
117async fn run_tick<Tick, Fut, E>(task: &'static str, tick: &Tick, backend: &Backend)
118where
119    Tick: Fn(&Backend) -> Fut + Send + Sync,
120    Fut: Future<Output = Result<bool, E>> + Send,
121    E: Error + Send + Sync,
122{
123    loop {
124        match tick(backend).await {
125            Ok(true) => continue,
126            Ok(false) => break,
127            Err(error) => {
128                warn!(task, %error, error_source = error.source().map(|s| s.to_string()), "bgtask tick failed");
129                break;
130            }
131        }
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use std::sync::{
138        Arc,
139        atomic::{AtomicUsize, Ordering},
140    };
141
142    use bytesize::ByteSize;
143    use slatedb::object_store::memory::InMemory;
144
145    use super::*;
146
147    pub(super) async fn test_backend() -> Backend {
148        let object_store = Arc::new(InMemory::new());
149        let db = slatedb::Db::builder("/test", object_store)
150            .build()
151            .await
152            .unwrap();
153        Backend::new(db, ByteSize::mib(10))
154    }
155
156    #[tokio::test]
157    async fn run_tick_repeats_until_done() {
158        let backend = test_backend().await;
159        let calls = Arc::new(AtomicUsize::new(0));
160        let tick = {
161            let calls = Arc::clone(&calls);
162            move |_backend: &Backend| {
163                let calls = Arc::clone(&calls);
164                async move {
165                    let count = calls.fetch_add(1, Ordering::SeqCst);
166                    Ok::<bool, std::io::Error>(count < 2)
167                }
168            }
169        };
170
171        run_tick("test", &tick, &backend).await;
172
173        assert_eq!(calls.load(Ordering::SeqCst), 3);
174    }
175}
s2_lite/backend/bgtasks/mod.rs

s2_lite/backend/bgtasks/
mod.rs
