OXYGENHUB
◆ v0.6.1 oxygenhub / crates / hub / src / leader.rs
12.4 KiB · Rust · last changed Run the hub as web and worker processes, singletons under a leader lock 16 days ago source blame history raw
1 
//! Leader locks (docs/roles.md, #272): a background task that must run
2 
//! once across every process sharing the hub's database runs only in the
3 
//! process holding a lock named for it.
4 
//!
5 
//! On PostgreSQL the lock is a session advisory lock
6 
//! (`pg_try_advisory_lock`) on one connection of this process's own, kept
7 
//! apart from the pool for as long as the process runs. A process that does
8 
//! not hold a task's lock asks again every [`RETRY`]. One that holds it
9 
//! checks its connection every [`CHECK`]; if the connection has gone (the
10 
//! database restarted, the network dropped, the session was ended), the
11 
//! server has freed the lock, so the task is stopped (its future dropped)
12 
//! and the process contends again on a new connection. A process that dies
13 
//! closes its connection, freeing its locks, and another takes over within
14 
//! seconds. Each task reads its work from the database, so a successor
15 
//! picking up where it stopped is safe.
16 
//!
17 
//! On SQLite (one machine, one process by definition) every lock is held.
18
19 
use crate::Shared;
20 
use std::collections::BTreeSet;
21 
use std::future::Future;
22 
use std::sync::Mutex;
23 
use std::time::Duration;
24
25 
/// How often a process that does not hold a lock asks for it again.
26 
pub const RETRY: Duration = Duration::from_secs(2);
27
28 
/// How often a holder makes sure its connection, and so its lock, is still
29 
/// there; and how long it waits for the database's answer.
30 
pub const CHECK: Duration = Duration::from_secs(2);
31
32 
/// This process's background tasks and the locks it holds.
33 
#[derive(Default)]
34 
pub struct Leader {
35 
    /// Every background task this process started, led or not.
36 
    tasks: Mutex<BTreeSet<&'static str>>,
37 
    /// The singleton tasks running here now: the locks this process holds.
38 
    leading: Mutex<BTreeSet<&'static str>>,
39 
    /// The loops contending for locks, stopped by [`Leader::resign`].
40 
    contenders: Mutex<Vec<tokio::task::AbortHandle>>,
41 
    #[cfg(feature = "postgres")]
42 
    session: pg::Session,
43 
}
44
45 
impl Leader {
46 
    /// The background tasks this process started, by name.
47 
    pub fn tasks(&self) -> Vec<&'static str> {
48 
        self.tasks.lock().unwrap_or_else(|p| p.into_inner()).iter().copied().collect()
49 
    }
50
51 
    /// The leader locks this process holds: the singleton tasks running here.
52 
    pub fn leading(&self) -> Vec<&'static str> {
53 
        self.leading.lock().unwrap_or_else(|p| p.into_inner()).iter().copied().collect()
54 
    }
55
56 
    /// Record a background task that needs no lock (it is safe anywhere).
57 
    pub fn note(&self, name: &'static str) {
58 
        self.tasks.lock().unwrap_or_else(|p| p.into_inner()).insert(name);
59 
    }
60
61 
    /// Stop every singleton task here and give up their locks, as a process
62 
    /// that stops does: another process takes them over.
63 
    pub async fn resign(&self) {
64 
        for c in self.contenders.lock().unwrap_or_else(|p| p.into_inner()).drain(..) {
65 
            c.abort();
66 
        }
67 
        #[cfg(feature = "postgres")]
68 
        self.session.close().await;
69 
    }
70
71 
    /// The process id of the database session holding this process's
72 
    /// locks, if it has one (for tests that end it).
73 
    #[cfg(feature = "postgres")]
74 
    pub async fn session_pid(&self) -> Option<i32> {
75 
        self.session.pid().await
76 
    }
77 
}
78
79 
// Taking, checking and giving back a lock: on PostgreSQL through the
80 
// session; on SQLite every lock is this process's.
81 
#[cfg(feature = "postgres")]
82 
impl Leader {
83 
    /// Take `name`'s lock now if it is free: the session's generation it
84 
    /// was taken in, which [`Leader::held`] checks against.
85 
    async fn acquire(&self, app: &Shared, name: &'static str) -> Option<u64> {
86 
        self.session.acquire(&app.db, name).await
87 
    }
88
89 
    /// Whether a lock taken in generation `term` is still held.
90 
    async fn held(&self, term: u64) -> bool {
91 
        self.session.held(term).await
92 
    }
93
94 
    /// Give `name`'s lock back, its task done.
95 
    async fn release(&self, name: &'static str, term: u64) {
96 
        self.session.release(name, term).await
97 
    }
98 
}
99
100 
#[cfg(not(feature = "postgres"))]
101 
impl Leader {
102 
    async fn acquire(&self, _app: &Shared, _name: &'static str) -> Option<u64> {
103 
        Some(0)
104 
    }
105
106 
    async fn held(&self, _term: u64) -> bool {
107 
        true
108 
    }
109
110 
    async fn release(&self, _name: &'static str, _term: u64) {}
111 
}
112
113 
/// Marks a task as led here while it runs; dropping it (the task ended,
114 
/// was stopped, or the loop was aborted) unmarks it.
115 
struct Leading<'a>(&'a Leader, &'static str);
116
117 
impl<'a> Leading<'a> {
118 
    fn new(leader: &'a Leader, name: &'static str) -> Self {
119 
        leader.leading.lock().unwrap_or_else(|p| p.into_inner()).insert(name);
120 
        Leading(leader, name)
121 
    }
122 
}
123
124 
impl Drop for Leading<'_> {
125 
    fn drop(&mut self) {
126 
        self.0.leading.lock().unwrap_or_else(|p| p.into_inner()).remove(self.1);
127 
    }
128 
}
129
130 
/// Run the singleton task `name` in the background, only while this
131 
/// process holds its lock. `task` makes the task's future; it is made
132 
/// again each time the lock is taken after being lost. A task that ends by
133 
/// itself gives its lock back and is not run again.
134 
pub fn run<F, Fut>(app: &Shared, name: &'static str, task: F)
135 
where
136 
    F: Fn(Shared) -> Fut + Send + 'static,
137 
    Fut: Future<Output = ()> + Send + 'static,
138 
{
139 
    app.leader.note(name);
140 
    let handle = tokio::spawn(contend(app.clone(), name, task));
141 
    app.leader.contenders.lock().unwrap_or_else(|p| p.into_inner()).push(handle.abort_handle());
142 
}
143
144 
async fn contend<F, Fut>(app: Shared, name: &'static str, task: F)
145 
where
146 
    F: Fn(Shared) -> Fut,
147 
    Fut: Future<Output = ()>,
148 
{
149 
    let leader = &app.leader;
150 
    let mut waiting = false;
151 
    loop {
152 
        let Some(term) = leader.acquire(&app, name).await else {
153 
            if !waiting {
154 
                tracing::info!("leader: {name} runs elsewhere; standing by");
155 
                waiting = true;
156 
            }
157 
            tokio::time::sleep(RETRY).await;
158 
            continue;
159 
        };
160 
        if waiting {
161 
            tracing::info!("leader: took over {name}");
162 
            waiting = false;
163 
        }
164 
        let _leading = Leading::new(leader, name);
165 
        let work = task(app.clone());
166 
        tokio::pin!(work);
167 
        let finished = loop {
168 
            tokio::select! {
169 
                _ = &mut work => break true,
170 
                _ = tokio::time::sleep(CHECK) => {
171 
                    if !leader.held(term).await {
172 
                        break false;
173 
                    }
174 
                }
175 
            }
176 
        };
177 
        if finished {
178 
            leader.release(name, term).await;
179 
            return;
180 
        }
181 
        // The future is dropped here: the task stops where it was.
182 
        tracing::warn!("leader: lost the lock for {name}; stopped it, contending again");
183 
    }
184 
}
185
186 
#[cfg(feature = "postgres")]
187 
mod pg {
188 
    use super::CHECK;
189 
    use crate::db::Pool;
190 
    use sqlx::postgres::PgConnection;
191 
    use sqlx::Connection;
192 
    use std::time::Instant;
193
194 
    /// The connection holding this process's locks.
195 
    struct Conn {
196 
        conn: PgConnection,
197 
        /// The schema the hub's tables are in: two hubs in two schemas of
198 
        /// one database are two hubs, and do not share locks.
199 
        schema: String,
200 
        pid: i32,
201 
        /// When the database last answered on it.
202 
        checked: Instant,
203 
    }
204
205 
    #[derive(Default)]
206 
    pub(super) struct Session {
207 
        conn: tokio::sync::Mutex<Option<Conn>>,
208 
        /// Counts connections lost: a lock taken in an earlier generation
209 
        /// went with its connection.
210 
        term: std::sync::atomic::AtomicU64,
211 
        closed: std::sync::atomic::AtomicBool,
212 
    }
213
214 
    /// The lock's key: a stable 64-bit hash of the schema and the name.
215 
    pub(super) fn key(schema: &str, name: &str) -> i64 {
216 
        let h = blake3::hash(format!("oxygenhub leader\0{schema}\0{name}").as_bytes());
217 
        i64::from_le_bytes(h.as_bytes()[..8].try_into().expect("eight bytes"))
218 
    }
219
220 
    impl Session {
221 
        fn term(&self) -> u64 {
222 
            self.term.load(std::sync::atomic::Ordering::SeqCst)
223 
        }
224
225 
        /// The connection is gone, and every lock with it.
226 
        fn lose(&self, slot: &mut Option<Conn>) {
227 
            if slot.take().is_some() {
228 
                self.term.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
229 
            }
230 
        }
231
232 
        async fn connect(db: &Pool) -> anyhow::Result<Conn> {
233 
            let options = db.connect_options();
234 
            let mut conn = PgConnection::connect_with(&options).await?;
235 
            let (schema, pid): (String, i32) =
236 
                sqlx::query_as("SELECT CAST(current_schema() AS TEXT), pg_backend_pid()").fetch_one(&mut conn).await?;
237 
            Ok(Conn { conn, schema, pid, checked: Instant::now() })
238 
        }
239
240 
        pub(super) async fn acquire(&self, db: &Pool, name: &'static str) -> Option<u64> {
241 
            let mut slot = self.conn.lock().await;
242 
            if self.closed.load(std::sync::atomic::Ordering::SeqCst) {
243 
                return None;
244 
            }
245 
            if slot.is_none() {
246 
                match tokio::time::timeout(CHECK * 5, Self::connect(db)).await {
247 
                    Ok(Ok(c)) => *slot = Some(c),
248 
                    Ok(Err(e)) => {
249 
                        tracing::warn!("leader: connecting for {name}'s lock: {e}");
250 
                        return None;
251 
                    }
252 
                    Err(_) => {
253 
                        tracing::warn!("leader: connecting for {name}'s lock: timed out");
254 
                        return None;
255 
                    }
256 
                }
257 
            }
258 
            let term = self.term();
259 
            let c = slot.as_mut()?;
260 
            let ask = sqlx::query_scalar::<_, bool>("SELECT pg_try_advisory_lock($1)").bind(key(&c.schema, name));
261 
            match tokio::time::timeout(CHECK, ask.fetch_one(&mut c.conn)).await {
262 
                Ok(Ok(got)) => {
263 
                    c.checked = Instant::now();
264 
                    got.then_some(term)
265 
                }
266 
                Ok(Err(e)) => {
267 
                    tracing::warn!("leader: asking for {name}'s lock: {e}");
268 
                    self.lose(&mut slot);
269 
                    None
270 
                }
271 
                Err(_) => {
272 
                    tracing::warn!("leader: asking for {name}'s lock: timed out");
273 
                    self.lose(&mut slot);
274 
                    None
275 
                }
276 
            }
277 
        }
278
279 
        pub(super) async fn held(&self, term: u64) -> bool {
280 
            let mut slot = self.conn.lock().await;
281 
            if self.term() != term {
282 
                return false;
283 
            }
284 
            let Some(c) = slot.as_mut() else { return false };
285 
            // One answer serves every task's check made soon after it.
286 
            if c.checked.elapsed() < CHECK / 2 {
287 
                return true;
288 
            }
289 
            let ping = sqlx::query_scalar::<_, i32>("SELECT 1").fetch_one(&mut c.conn);
290 
            match tokio::time::timeout(CHECK, ping).await {
291 
                Ok(Ok(_)) => {
292 
                    c.checked = Instant::now();
293 
                    true
294 
                }
295 
                _ => {
296 
                    tracing::warn!("leader: the connection holding the locks is gone");
297 
                    self.lose(&mut slot);
298 
                    false
299 
                }
300 
            }
301 
        }
302
303 
        pub(super) async fn release(&self, name: &'static str, term: u64) {
304 
            let mut slot = self.conn.lock().await;
305 
            if self.term() != term {
306 
                return;
307 
            }
308 
            let Some(c) = slot.as_mut() else { return };
309 
            let unlock = sqlx::query_scalar::<_, bool>("SELECT pg_advisory_unlock($1)").bind(key(&c.schema, name));
310 
            if !matches!(tokio::time::timeout(CHECK, unlock.fetch_one(&mut c.conn)).await, Ok(Ok(_))) {
311 
                self.lose(&mut slot);
312 
            }
313 
        }
314
315 
        pub(super) async fn close(&self) {
316 
            self.closed.store(true, std::sync::atomic::Ordering::SeqCst);
317 
            let mut slot = self.conn.lock().await;
318 
            if let Some(c) = slot.take() {
319 
                self.term.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
320 
                let _ = c.conn.close().await;
321 
            }
322 
        }
323
324 
        pub(super) async fn pid(&self) -> Option<i32> {
325 
            self.conn.lock().await.as_ref().map(|c| c.pid)
326 
        }
327 
    }
328
329 
    #[cfg(test)]
330 
    mod tests {
331 
        #[test]
332 
        fn keys_are_stable_and_per_schema() {
333 
            assert_eq!(super::key("hub_a", "flow.scheduler"), super::key("hub_a", "flow.scheduler"));
334 
            assert_ne!(super::key("hub_a", "flow.scheduler"), super::key("hub_b", "flow.scheduler"));
335 
            assert_ne!(super::key("hub_a", "flow.scheduler"), super::key("hub_a", "flow.reaper"));
336 
        }
337 
    }
338 
}