diff --git a/crates/leader-election/src/callbacks.rs b/crates/leader-election/src/callbacks.rs index 63e5f93..f873ca2 100644 --- a/crates/leader-election/src/callbacks.rs +++ b/crates/leader-election/src/callbacks.rs @@ -24,9 +24,11 @@ pub trait LeaderCallbacks: Send + Sync { /// Receives a cancellation token that is cancelled when leadership is lost. async fn on_started_leading(&self, cancel: CancellationToken); - /// Called when this instance stops being the leader (always called, even if never led). + /// Called after each leadership epoch ends. Also called once if the elector stops before it + /// ever becomes leader. async fn on_stopped_leading(&self); - /// Called when a new leader is observed (fire-and-forget, runs in a separate task). + /// Called when a new leader is observed. Implementations should return quickly or spawn their + /// own work. async fn on_new_leader(&self, identity: String); } diff --git a/crates/leader-election/src/elector.rs b/crates/leader-election/src/elector.rs index db69868..62654cd 100644 --- a/crates/leader-election/src/elector.rs +++ b/crates/leader-election/src/elector.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use std::time::Duration; use chrono::{DateTime, Utc}; -use tokio::sync::RwLock; +use tokio::sync::{RwLock, watch}; use tokio_util::sync::CancellationToken; use tracing::{debug, info, warn}; @@ -122,23 +122,43 @@ impl LeaderElector { &self, callbacks: impl LeaderCallbacks + 'static, cancel: CancellationToken, + ) -> Result<(), Error> { + self.run_with_state(callbacks, cancel, None).await + } + + /// Run the leader election loop with an optional external state channel. + async fn run_with_state( + &self, + callbacks: impl LeaderCallbacks + 'static, + cancel: CancellationToken, + state_tx: Option>, ) -> Result<(), Error> { let callbacks = Arc::new(callbacks); + // Preserve the existing contract that a run cancelled before its first acquisition still + // receives one stopped notification. After a completed leadership epoch, do not emit a + // duplicate notification if cancellation arrives while re-acquiring. + let mut has_notified_stopped = false; info!(identity = %self.config.identity, lock = %self.lock.describe(), "starting leader election"); loop { // Phase 1: acquire if cancel.is_cancelled() { - callbacks.on_stopped_leading().await; + if !has_notified_stopped { + callbacks.on_stopped_leading().await; + } return Ok(()); } - if !self.acquire(&cancel, &callbacks).await { + if !self.acquire(&cancel, &callbacks, state_tx.as_ref()).await { // Cancelled during acquire - callbacks.on_stopped_leading().await; + if !has_notified_stopped { + callbacks.on_stopped_leading().await; + } return Ok(()); } + Self::publish_state(state_tx.as_ref(), LeaderState::Leading); + // We are now the leader. Create a child token for the leading task. let leading_cancel = CancellationToken::new(); @@ -156,8 +176,10 @@ impl LeaderElector { let should_retry = self.renew(&cancel).await; // We lost leadership (or lost renew loop due cancel). - // Stop the leading task. + // Revoke the leading task and externally visible state before any asynchronous cleanup + // or re-acquisition work, so consumers cannot observe stale leadership. leading_cancel.cancel(); + Self::publish_state(state_tx.as_ref(), LeaderState::Pending); // Wait for the leading task to finish. let _ = leading_handle.await; @@ -166,10 +188,12 @@ impl LeaderElector { self.release().await; } + // Phase 5: notify stopped once for every completed leadership epoch. + callbacks.on_stopped_leading().await; + has_notified_stopped = true; + if !should_retry { info!(identity = %self.config.identity, "stopped leading"); - // Phase 5: notify stopped - callbacks.on_stopped_leading().await; return Ok(()); } @@ -192,14 +216,10 @@ impl LeaderElector { { let (state_tx, state_rx) = tokio::sync::watch::channel(LeaderState::Pending); let handle = LeaderElectorHandle { state_rx }; - let join = tokio::spawn(async move { - // Wrap callbacks to also update the watch channel - let wrapped = StateTrackingCallbacks { - inner: callbacks, - state_tx, - }; - self.run(wrapped, cancel).await - }); + let join = + tokio::spawn( + async move { self.run_with_state(callbacks, cancel, Some(state_tx)).await }, + ); (handle, join) } @@ -210,6 +230,7 @@ impl LeaderElector { &self, cancel: &CancellationToken, callbacks: &Arc, + state_tx: Option<&watch::Sender>, ) -> bool { info!(identity = %self.config.identity, "attempting to acquire leader lock"); loop { @@ -223,9 +244,10 @@ impl LeaderElector { } } - if self.try_acquire_or_renew().await { + let acquired = self.try_acquire_or_renew().await; + self.maybe_report_transition(callbacks, state_tx).await; + if acquired { info!(identity = %self.config.identity, "successfully acquired lease"); - self.maybe_report_transition(callbacks).await; return true; } } @@ -531,26 +553,52 @@ impl LeaderElector { /// Check if a new leader has been observed and fire the on_new_leader callback. /// Deduplicates: only fires when the leader identity changes. - async fn maybe_report_transition(&self, callbacks: &Arc) { - let mut observed = self.observed.write().await; - let current_leader = observed - .record - .as_ref() - .map(|r| r.holder_identity.clone()) - .unwrap_or_default(); + async fn maybe_report_transition( + &self, + callbacks: &Arc, + state_tx: Option<&watch::Sender>, + ) { + let current_leader = { + let mut observed = self.observed.write().await; + let current_leader = observed + .record + .as_ref() + .map(|r| r.holder_identity.clone()) + .unwrap_or_default(); - // Dedup: skip if we already reported this leader - if observed.reported_leader.as_deref() == Some(¤t_leader) { + // Dedup: skip if we already reported this leader. + if observed.reported_leader.as_deref() == Some(¤t_leader) { + return; + } + + observed.reported_leader = Some(current_leader.clone()); + current_leader + }; + + if current_leader.is_empty() { + Self::publish_state(state_tx, LeaderState::Pending); return; } - observed.reported_leader = Some(current_leader.clone()); + if current_leader != self.config.identity { + Self::publish_state(state_tx, LeaderState::Following(current_leader.clone())); + } + debug!(new_leader = %current_leader, "observed new leader"); + // Call without holding the observed-state lock. Callbacks are expected to be fast or + // spawn their own work. + callbacks.on_new_leader(current_leader).await; + } - if !current_leader.is_empty() { - debug!(new_leader = %current_leader, "observed new leader"); - // Fire the on_new_leader callback (runs in caller's task; callbacks are expected - // to be fast or spawn their own work). - callbacks.on_new_leader(current_leader).await; + /// Publish an externally observable state only when it actually changes. + fn publish_state(state_tx: Option<&watch::Sender>, state: LeaderState) { + if let Some(state_tx) = state_tx { + state_tx.send_if_modified(|current| { + if *current == state { + return false; + } + *current = state; + true + }); } } @@ -561,34 +609,6 @@ impl LeaderElector { } } -/// Internal callbacks wrapper that updates the watch channel on state transitions. -struct StateTrackingCallbacks { - inner: C, - state_tx: tokio::sync::watch::Sender, -} - -#[async_trait::async_trait] -impl LeaderCallbacks for StateTrackingCallbacks { - async fn on_started_leading(&self, cancel: CancellationToken) { - let _ = self.state_tx.send(LeaderState::Leading); - self.inner.on_started_leading(cancel).await; - } - - async fn on_stopped_leading(&self) { - let _ = self.state_tx.send(LeaderState::Pending); - self.inner.on_stopped_leading().await; - } - - async fn on_new_leader(&self, identity: String) { - // Update state channel: if we're not currently Leading, report Following - let is_leading = matches!(&*self.state_tx.borrow(), LeaderState::Leading); - if !is_leading && !identity.is_empty() { - let _ = self.state_tx.send(LeaderState::Following(identity.clone())); - } - self.inner.on_new_leader(identity).await; - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/leader-election/src/state.rs b/crates/leader-election/src/state.rs index 3d95e1d..72084f8 100644 --- a/crates/leader-election/src/state.rs +++ b/crates/leader-election/src/state.rs @@ -58,7 +58,10 @@ impl LeaderElectorHandle { if rx.changed().await.is_err() { break; } - yield rx.borrow().clone(); + // Drop the watch borrow before yielding. Holding it across a yield would block + // producers from publishing the next state until the stream is polled again. + let state = rx.borrow().clone(); + yield state; } } } diff --git a/crates/leader-election/tests/integration_tests.rs b/crates/leader-election/tests/integration_tests.rs index ef6fdbe..de28e0b 100644 --- a/crates/leader-election/tests/integration_tests.rs +++ b/crates/leader-election/tests/integration_tests.rs @@ -14,17 +14,19 @@ //! Integration tests for leader election. -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, RwLock as StdRwLock}; use std::time::Duration; use async_trait::async_trait; use chrono::{DateTime, Utc}; -use tokio::sync::{Mutex, RwLock}; +use futures::StreamExt; +use tokio::sync::{Mutex, Notify, RwLock}; use tokio_util::sync::CancellationToken; use kube_leader_election::{ - Clock, Error, LeaderCallbacks, LeaderElectionRecord, LeaderElector, LeaderElectorConfig, Lock, + Clock, Error, LeaderCallbacks, LeaderElectionRecord, LeaderElector, LeaderElectorConfig, + LeaderState, Lock, }; // ─── Test Helpers ─────────────────────────────────────────────────────────── @@ -37,6 +39,9 @@ struct FakeLock { resource_version: Arc>, /// If set, update() will fail with Conflict this many times before succeeding. conflict_count: Arc, + /// If true, all update attempts fail until the test explicitly allows them again. + fail_updates: Arc, + update_failure: Arc, } impl FakeLock { @@ -46,6 +51,8 @@ impl FakeLock { record: Arc::new(RwLock::new(None)), resource_version: Arc::new(Mutex::new(0)), conflict_count: Arc::new(AtomicUsize::new(0)), + fail_updates: Arc::new(AtomicBool::new(false)), + update_failure: Arc::new(Notify::new()), } } @@ -61,6 +68,16 @@ impl FakeLock { fn set_conflicts(&self, count: usize) { self.conflict_count.store(count, Ordering::SeqCst); } + + /// Configure every update to fail or succeed. + fn set_fail_updates(&self, fail: bool) { + self.fail_updates.store(fail, Ordering::SeqCst); + } + + /// Wait until an update has failed because `fail_updates` is enabled. + async fn wait_for_update_failure(&self) { + self.update_failure.notified().await; + } } #[async_trait] @@ -82,6 +99,11 @@ impl Lock for FakeLock { } async fn update(&self, record: LeaderElectionRecord) -> Result<(), Error> { + if self.fail_updates.load(Ordering::SeqCst) { + self.update_failure.notify_one(); + return Err(Error::Conflict); + } + // Simulate conflict if configured let remaining = self.conflict_count.load(Ordering::SeqCst); if remaining > 0 { @@ -108,25 +130,30 @@ impl Lock for FakeLock { /// Mock clock with controllable time. #[derive(Clone)] struct MockClock { - now: Arc>>, + now: Arc>>, } impl MockClock { fn new(time: DateTime) -> Self { Self { - now: Arc::new(RwLock::new(time)), + now: Arc::new(StdRwLock::new(time)), } } async fn set(&self, time: DateTime) { - *self.now.write().await = time; + match self.now.write() { + Ok(mut now) => *now = time, + Err(poisoned) => *poisoned.into_inner() = time, + } } } impl Clock for MockClock { fn now(&self) -> DateTime { - // Block on async read (safe in tests) - futures::executor::block_on(async { *self.now.read().await }) + match self.now.read() { + Ok(now) => *now, + Err(poisoned) => *poisoned.into_inner(), + } } } @@ -416,6 +443,42 @@ async fn test_acquire_active_lease() { let _ = tokio::time::timeout(Duration::from_secs(2), handle).await; } +#[tokio::test] +async fn test_spawn_reports_observed_remote_leader() { + let lock = FakeLock::new("node-2"); + let now = Utc::now(); + lock.set_record(LeaderElectionRecord { + holder_identity: "node-1".to_string(), + lease_duration_seconds: 15, + acquire_time: now, + renew_time: now, + leader_transitions: 0, + }) + .await; + + let callbacks = Arc::new(TestCallbacks::new()); + let elector = LeaderElector::new(test_config("node-2"), lock, MockClock::new(now)).unwrap(); + let cancel = CancellationToken::new(); + let (handle, join) = elector.spawn(SharedCallbacks(callbacks.clone()), cancel.clone()); + let mut states = Box::pin(handle.state_stream()); + + let state = tokio::time::timeout(Duration::from_secs(2), states.next()) + .await + .expect("elector should report the observed leader") + .expect("state stream should remain open"); + assert_eq!(state, LeaderState::Following("node-1".to_string())); + assert_eq!(handle.current_leader().as_deref(), Some("node-1")); + assert!(!handle.is_leader()); + + cancel.cancel(); + let result = tokio::time::timeout(Duration::from_secs(2), join) + .await + .expect("elector should stop after cancellation") + .expect("elector task should not panic"); + assert!(result.is_ok()); + assert_eq!(callbacks.stopped_count().await, 1); +} + #[tokio::test] async fn test_acquire_non_positive_lease_without_panicking() { for lease_duration_seconds in [i32::MIN, -1, 0] { @@ -624,6 +687,8 @@ async fn test_concurrent_acquire() { record: shared_record.clone(), resource_version: shared_rv.clone(), conflict_count: Arc::new(AtomicUsize::new(0)), + fail_updates: Arc::new(AtomicBool::new(false)), + update_failure: Arc::new(Notify::new()), }; let lock2 = FakeLock { @@ -631,6 +696,8 @@ async fn test_concurrent_acquire() { record: shared_record.clone(), resource_version: shared_rv.clone(), conflict_count: Arc::new(AtomicUsize::new(0)), + fail_updates: Arc::new(AtomicBool::new(false)), + update_failure: Arc::new(Notify::new()), }; let clock1 = MockClock::new(Utc::now()); @@ -673,3 +740,66 @@ async fn test_concurrent_acquire() { "exactly one candidate should become leader" ); } + +/// A renewal deadline revokes externally visible leadership before re-acquisition. +#[tokio::test] +async fn test_spawn_reports_non_leader_while_reacquiring() { + let now = Utc::now(); + let lock = FakeLock::new("node-1"); + let clock = MockClock::new(now); + let test_clock = clock.clone(); + let callbacks = Arc::new(TestCallbacks::new()); + let elector = LeaderElector::new(test_config("node-1"), lock.clone(), clock).unwrap(); + let cancel = CancellationToken::new(); + + let (handle, join) = elector.spawn(SharedCallbacks(callbacks.clone()), cancel.clone()); + let mut states = Box::pin(handle.state_stream()); + + let first_state = tokio::time::timeout(Duration::from_secs(2), states.next()) + .await + .expect("elector should acquire the lease") + .expect("state stream should remain open"); + assert_eq!(first_state, LeaderState::Leading); + assert!(handle.is_leader()); + + // Force the next renewal cycle to fail, then move the injected clock beyond the deadline. + // Updates remain blocked so the elector cannot immediately re-acquire after stepping down. + lock.set_fail_updates(true); + tokio::time::timeout(Duration::from_secs(2), lock.wait_for_update_failure()) + .await + .expect("renewal should attempt an update"); + test_clock.set(now + chrono::Duration::seconds(11)).await; + + let lost_state = tokio::time::timeout(Duration::from_secs(2), states.next()) + .await + .expect("renewal failure should publish a state change") + .expect("state stream should remain open"); + assert_eq!(lost_state, LeaderState::Pending); + assert!(!handle.is_leader()); + + tokio::time::timeout(Duration::from_secs(2), async { + while callbacks.stopped_count().await != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("leadership loss should invoke the stopped callback once"); + + // Permit updates again and verify that leadership is only re-published after acquisition. + lock.set_fail_updates(false); + let reacquired_state = tokio::time::timeout(Duration::from_secs(2), states.next()) + .await + .expect("elector should re-acquire the lease") + .expect("state stream should remain open"); + assert_eq!(reacquired_state, LeaderState::Leading); + assert!(handle.is_leader()); + + cancel.cancel(); + let result = tokio::time::timeout(Duration::from_secs(2), join) + .await + .expect("elector should stop after cancellation") + .expect("elector task should not panic"); + assert!(result.is_ok()); + assert_eq!(callbacks.started_count().await, 2); + assert_eq!(callbacks.stopped_count().await, 2); +}