Skip to main content

rws_gateway_api/gateway_api/
lease.rs

1//! Leader election over a `coordination.k8s.io/v1` Lease: every replica
2//! serves traffic, only the holder writes status and Services. The holder
3//! renews every [`RENEW`]; another replica takes over once a renewal is older
4//! than the lease's duration. Updates carry the Lease's `resourceVersion`,
5//! so two replicas can't both win.
6
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::Arc;
9use std::time::Duration;
10
11use k8s_openapi::api::coordination::v1::{Lease, LeaseSpec};
12use k8s_openapi::apimachinery::pkg::apis::meta::v1::{MicroTime, ObjectMeta};
13use kube::api::{Api, PostParams};
14use kube::Client;
15use tokio::sync::Notify;
16
17/// How long a holder's renewal counts.
18pub const DURATION: Duration = Duration::from_secs(15);
19/// How often the holder renews (and others check).
20pub const RENEW: Duration = Duration::from_secs(5);
21
22pub struct LeaseSettings {
23    pub namespace: String,
24    pub name: String,
25    /// This replica (the pod's name).
26    pub identity: String,
27}
28
29/// What a replica should do with the Lease as it is now.
30#[derive(Debug, PartialEq)]
31pub enum Decision {
32    /// Create it, holding it.
33    Create,
34    /// Write it back holding it (renewing, or taking over an expired one).
35    Take,
36    /// Someone else holds it.
37    Follow,
38}
39
40pub fn decide(lease: Option<&Lease>, identity: &str, now: k8s_openapi::jiff::Timestamp) -> Decision {
41    let Some(spec) = lease.and_then(|l| l.spec.as_ref()) else {
42        return if lease.is_some() { Decision::Take } else { Decision::Create };
43    };
44    if spec.holder_identity.as_deref() == Some(identity) {
45        return Decision::Take;
46    }
47    let duration = spec.lease_duration_seconds.unwrap_or(DURATION.as_secs() as i32) as i64;
48    let renewed = spec.renew_time.as_ref().map(|t| t.0.as_second()).unwrap_or(0);
49    if spec.holder_identity.as_deref().unwrap_or("").is_empty() || renewed + duration < now.as_second() {
50        Decision::Take
51    } else {
52        Decision::Follow
53    }
54}
55
56/// Hold or follow the Lease forever; `leader` says which, `gained` fires
57/// when this replica becomes the leader (so it can reconcile at once).
58pub async fn run(client: Client, settings: LeaseSettings, leader: Arc<AtomicBool>, gained: Arc<Notify>) {
59    let api: Api<Lease> = Api::namespaced(client, &settings.namespace);
60    loop {
61        let now = k8s_openapi::jiff::Timestamp::now();
62        let current = api.get_opt(&settings.name).await;
63        let holding = match current {
64            Err(e) => {
65                eprintln!("[gateway-api] lease: {e}");
66                false
67            }
68            Ok(current) => match decide(current.as_ref(), &settings.identity, now) {
69                Decision::Follow => false,
70                Decision::Create => {
71                    let lease = held(&settings, None, now);
72                    api.create(&PostParams::default(), &lease).await.is_ok()
73                }
74                Decision::Take => {
75                    let lease = held(&settings, current.as_ref(), now);
76                    api.replace(&settings.name, &PostParams::default(), &lease).await.is_ok()
77                }
78            },
79        };
80        let was = leader.swap(holding, Ordering::SeqCst);
81        if holding && !was {
82            gained.notify_one();
83        }
84        tokio::time::sleep(RENEW).await;
85    }
86}
87
88/// The Lease held by this replica, based on `current` (keeping its
89/// `resourceVersion` and counting a change of holder).
90fn held(settings: &LeaseSettings, current: Option<&Lease>, now: k8s_openapi::jiff::Timestamp) -> Lease {
91    let previous = current.and_then(|l| l.spec.clone()).unwrap_or_default();
92    let changed = previous.holder_identity.as_deref() != Some(settings.identity.as_str());
93    Lease {
94        metadata: ObjectMeta {
95            name: Some(settings.name.clone()),
96            namespace: Some(settings.namespace.clone()),
97            resource_version: current.and_then(|l| l.metadata.resource_version.clone()),
98            ..Default::default()
99        },
100        spec: Some(LeaseSpec {
101            holder_identity: Some(settings.identity.clone()),
102            lease_duration_seconds: Some(DURATION.as_secs() as i32),
103            acquire_time: if changed { Some(MicroTime(now)) } else { previous.acquire_time.clone() },
104            renew_time: Some(MicroTime(now)),
105            lease_transitions: Some(previous.lease_transitions.unwrap_or(0) + changed as i32),
106            ..Default::default()
107        }),
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    fn lease(holder: &str, renewed: i64) -> Lease {
116        Lease {
117            metadata: Default::default(),
118            spec: Some(LeaseSpec {
119                holder_identity: Some(holder.to_string()),
120                lease_duration_seconds: Some(15),
121                renew_time: Some(MicroTime(k8s_openapi::jiff::Timestamp::from_second(renewed).unwrap())),
122                ..Default::default()
123            }),
124        }
125    }
126
127    #[test]
128    fn replicas_create_renew_follow_and_take_over() {
129        let now = k8s_openapi::jiff::Timestamp::from_second(1_000).unwrap();
130        assert_eq!(Decision::Create, decide(None, "a", now));
131        assert_eq!(Decision::Take, decide(Some(&lease("a", 999)), "a", now), "the holder renews");
132        assert_eq!(Decision::Follow, decide(Some(&lease("b", 990)), "a", now), "a live holder keeps it");
133        assert_eq!(Decision::Take, decide(Some(&lease("b", 980)), "a", now), "an expired one is taken over");
134        assert_eq!(Decision::Take, decide(Some(&lease("", 999)), "a", now), "nobody holds it");
135    }
136
137    #[test]
138    fn taking_over_counts_a_transition_and_keeps_the_resource_version() {
139        let settings = LeaseSettings { namespace: "ns".to_string(), name: "l".to_string(), identity: "a".to_string() };
140        let mut current = lease("b", 0);
141        current.metadata.resource_version = Some("42".to_string());
142        let now = k8s_openapi::jiff::Timestamp::from_second(100).unwrap();
143        let taken = held(&settings, Some(&current), now);
144        assert_eq!(Some("42".to_string()), taken.metadata.resource_version);
145        let spec = taken.spec.unwrap();
146        assert_eq!((Some("a".to_string()), Some(1)), (spec.holder_identity.clone(), spec.lease_transitions));
147        let renewed = held(&settings, Some(&Lease { spec: Some(spec.clone()), ..current }), now);
148        assert_eq!(Some(1), renewed.spec.unwrap().lease_transitions, "renewing isn't a transition");
149    }
150}