rws_gateway_api/gateway_api/
lease.rs1use 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
17pub const DURATION: Duration = Duration::from_secs(15);
19pub const RENEW: Duration = Duration::from_secs(5);
21
22pub struct LeaseSettings {
23 pub namespace: String,
24 pub name: String,
25 pub identity: String,
27}
28
29#[derive(Debug, PartialEq)]
31pub enum Decision {
32 Create,
34 Take,
36 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
56pub 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
88fn 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(¤t), 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}