krabka_bench_driver/
failover.rs

1//! Kill-broker orchestration for the `failover` scenario.
2//!
3//! This module uses the in-cluster `kube` API through the Job's
4//! `ServiceAccount`. That `ServiceAccount` needs `pods: get,list,delete` in the
5//! target namespace, and it is mounted only when a failover scenario asks for
6//! it.
7//!
8//! `partition0_leader` resolves the topic metadata first and then maps the
9//! partition leader's broker id back to the matching `StatefulSet` pod. If that
10//! metadata probe fails, the caller can fall back to the first broker pod by
11//! name-sort for backwards-compatible smoke runs.
12
13use {
14    crate::{
15        numeric::nonnegative_i64_to_u64,
16        scenario::{FailoverSpec, Stack},
17    },
18    anyhow::{Context, Result, anyhow},
19    chrono::Utc,
20    k8s_openapi::api::core::v1::Pod,
21    krabka_client_core::{
22        Client as KafkaClient, ClientFrameMax, ConnectionDispatchQueueCapacity,
23        security::ClientSecurity,
24    },
25    krabka_protocol::owned::metadata_response::MetadataResponse,
26    krabka_units::convert::TimeExt,
27    kube::{
28        Client as KubeClient,
29        api::{Api, DeleteParams, ListParams},
30    },
31    std::{
32        future::Future,
33        sync::{
34            Arc,
35            atomic::{AtomicU64, Ordering},
36        },
37    },
38    tokio::{
39        task::JoinSet,
40        time::{Instant, sleep_until},
41    },
42    tracing::{info, warn},
43};
44
45pub(crate) struct FaultConfig {
46    pub stack: Stack,
47    pub namespace: String,
48    pub bootstrap: String,
49    pub topic: String,
50    pub security: Option<ClientSecurity>,
51    pub dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
52    pub frame_max: ClientFrameMax,
53}
54
55#[derive(Debug, PartialEq)]
56pub(crate) struct KillEvidence {
57    pub pod: String,
58    pub at_ms: u64,
59}
60
61#[derive(Debug, PartialEq)]
62pub(crate) enum FaultOutcome {
63    Cancelled,
64    TaskFailed(String),
65    ClientUnavailable(String),
66    Finished {
67        leader_lookup_error: Option<String>,
68        deletion: Result<KillEvidence, String>,
69    },
70}
71
72impl FaultOutcome {
73    pub(crate) fn killed_at_ms(&self) -> Option<u64> {
74        match self {
75            Self::Finished {
76                deletion: Ok(evidence),
77                ..
78            } => Some(evidence.at_ms),
79            _ => None,
80        }
81    }
82
83    pub(crate) fn note(&self) -> Option<String> {
84        match self {
85            Self::Cancelled => Some("failover: cancelled at measurement end".into()),
86            Self::TaskFailed(error) => Some(format!("failover: fault task failed: {error}")),
87            Self::ClientUnavailable(error) => {
88                Some(format!("failover: in-cluster client unavailable: {error}"))
89            }
90            Self::Finished {
91                leader_lookup_error,
92                deletion,
93            } => {
94                let fallback = leader_lookup_error
95                    .as_ref()
96                    .map_or_else(String::new, |error| {
97                        format!("leader lookup failed; using first broker pod: {error}; ")
98                    });
99                match deletion {
100                    Ok(evidence) if leader_lookup_error.is_some() => Some(format!(
101                        "failover: {fallback}killed broker {}",
102                        evidence.pod
103                    )),
104                    Ok(_) => None,
105                    Err(error) => Some(format!("failover: {fallback}broker kill failed: {error}")),
106                }
107            }
108        }
109    }
110}
111
112pub(crate) struct FaultTask {
113    tasks: JoinSet<FaultOutcome>,
114}
115
116impl FaultTask {
117    pub(crate) fn spawn(
118        spec: &FailoverSpec,
119        cfg: FaultConfig,
120        started_at: Instant,
121        kill_at: Arc<AtomicU64>,
122    ) -> Self {
123        let target_leader = spec.target == "partition0_leader";
124        let deadline = started_at + spec.kill_after.to_std();
125        Self::schedule(deadline, async move {
126            Box::pin(execute_fault(
127                target_leader,
128                try_client(),
129                partition0_leader_from_metadata(
130                    &cfg.bootstrap,
131                    &cfg.topic,
132                    cfg.security,
133                    cfg.dispatch_queue_capacity,
134                    cfg.frame_max,
135                ),
136                |client, leader| async move {
137                    kill_broker_pod(&client, cfg.stack, &cfg.namespace, leader).await
138                },
139                kill_at,
140            ))
141            .await
142        })
143    }
144
145    fn schedule(
146        deadline: Instant,
147        operation: impl Future<Output = FaultOutcome> + Send + 'static,
148    ) -> Self {
149        let mut tasks = JoinSet::new();
150        tasks.spawn(async move {
151            sleep_until(deadline).await;
152            operation.await
153        });
154        Self { tasks }
155    }
156
157    pub(crate) async fn stop(mut self) -> FaultOutcome {
158        // Keep completed evidence even if the measurement ends before this task is joined.
159        let result = if let Some(result) = self.tasks.try_join_next() {
160            Some(result)
161        } else {
162            self.tasks.abort_all();
163            self.tasks.join_next().await
164        };
165        match result {
166            Some(Ok(outcome)) => outcome,
167            Some(Err(error)) if error.is_cancelled() => FaultOutcome::Cancelled,
168            Some(Err(error)) => FaultOutcome::TaskFailed(error.to_string()),
169            None => FaultOutcome::TaskFailed("fault task missing".into()),
170        }
171    }
172}
173
174async fn execute_fault<C, D, L, K, F>(
175    target_leader: bool,
176    discover: D,
177    lookup: L,
178    delete: F,
179    kill_at: Arc<AtomicU64>,
180) -> FaultOutcome
181where
182    D: Future<Output = Result<C>>,
183    L: Future<Output = Result<i32>>,
184    K: Future<Output = Result<String>>,
185    F: FnOnce(C, Option<i32>) -> K,
186{
187    let client = match discover.await {
188        Ok(client) => client,
189        Err(error) => {
190            warn!(%error, "failover: in-cluster client unavailable");
191            return FaultOutcome::ClientUnavailable(error.to_string());
192        }
193    };
194    let (leader_id, leader_lookup_error) = if target_leader {
195        match lookup.await {
196            Ok(id) => (Some(id), None),
197            Err(error) => {
198                warn!(%error, "failover: leader lookup failed; using first broker pod");
199                (None, Some(error.to_string()))
200            }
201        }
202    } else {
203        (None, None)
204    };
205    let deletion = match delete(client, leader_id).await {
206        Ok(pod) => {
207            let at_ms = nonnegative_i64_to_u64(Utc::now().timestamp_millis());
208            kill_at.store(at_ms, Ordering::SeqCst);
209            info!(%pod, "failover: killed broker");
210            Ok(KillEvidence { pod, at_ms })
211        }
212        Err(error) => {
213            warn!(%error, "failover: broker kill failed");
214            Err(error.to_string())
215        }
216    };
217    FaultOutcome::Finished {
218        leader_lookup_error,
219        deletion,
220    }
221}
222
223/// Discovers an in-cluster Kubernetes client, which uses the in-pod
224/// `serviceAccount` token.
225///
226/// Outside the cluster, which is useful for `cargo test`, this returns an `Err`
227/// and the caller should report the failover as skipped.
228/// # Errors
229/// Returns an error when input data is invalid, required I/O fails, or the destination rejects the generated report or audit event.
230pub async fn try_client() -> Result<KubeClient> {
231    KubeClient::try_default()
232        .await
233        .context("build in-cluster kube client")
234}
235
236/// Queries the topic metadata and returns partition 0's current leader broker
237/// id.
238/// # Errors
239/// Returns an error when input data is invalid, required I/O fails, or the destination rejects the generated report or audit event.
240pub async fn partition0_leader_from_metadata(
241    bootstrap: &str,
242    topic: &str,
243    security: Option<ClientSecurity>,
244    dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
245    frame_max: ClientFrameMax,
246) -> Result<i32> {
247    let client = KafkaClient::builder()
248        .bootstrap(bootstrap.to_owned())
249        .client_id("bench-failover-targeter")
250        .dispatch_queue_capacity(dispatch_queue_capacity.get())
251        .frame_max(frame_max.size())
252        .maybe_security(security)
253        .build()
254        .await
255        .context("build metadata client")?;
256    let md = client
257        .refresh_metadata()
258        .await
259        .context("refresh metadata")?;
260    let leader = partition0_leader_id(&md, topic)
261        .ok_or_else(|| anyhow!("metadata did not contain {topic} partition 0 leader"))?;
262    client.close();
263    Ok(leader)
264}
265
266/// Deletes the requested broker pod.
267///
268/// When `leader_id` is `Some`, this targets the pod whose ordinal matches that
269/// broker id. Otherwise it deletes the first matching pod, for
270/// backwards-compatible smoke runs. `grace_period_seconds = 0` means SIGKILL.
271/// # Errors
272/// Returns an error when input data is invalid, required I/O fails, or the destination rejects the generated report or audit event.
273pub async fn kill_broker_pod(
274    client: &KubeClient,
275    stack: Stack,
276    namespace: &str,
277    leader_id: Option<i32>,
278) -> Result<String> {
279    let prefix = stack.broker_pod_regex().trim_start_matches('^');
280
281    let pods: Api<Pod> = Api::namespaced(client.clone(), namespace);
282    let list = pods
283        .list(&ListParams::default())
284        .await
285        .context("list pods in namespace")?;
286    let mut names: Vec<String> = list
287        .items
288        .iter()
289        .filter_map(|p| p.metadata.name.clone())
290        .filter(|n| n.starts_with(prefix))
291        .collect();
292    names.sort();
293    let Some(target) = choose_target_pod(&names, stack, leader_id) else {
294        return Err(anyhow!("no broker pod matched prefix {prefix}"));
295    };
296
297    let dp = DeleteParams::default().grace_period(0);
298    pods.delete(&target, &dp)
299        .await
300        .with_context(|| format!("delete pod {target}"))?;
301
302    Ok(target)
303}
304
305/// Deletes the first broker pod in alphabetical order that matches the stack's
306/// pod-name regex root. This stays as a compatibility wrapper for runs that do
307/// not target a partition.
308/// # Errors
309/// Returns an error when input data is invalid, required I/O fails, or the destination rejects the generated report or audit event.
310pub async fn kill_first_broker(
311    client: &KubeClient,
312    stack: Stack,
313    namespace: &str,
314) -> Result<String> {
315    kill_broker_pod(client, stack, namespace, None).await
316}
317
318fn partition0_leader_id(md: &MetadataResponse, topic: &str) -> Option<i32> {
319    md.topics
320        .iter()
321        .find(|t| t.name.as_deref() == Some(topic))
322        .and_then(|t| t.partitions.iter().find(|p| p.partition_index == 0))
323        .map(|p| p.leader_id)
324        .filter(|id| *id >= 0)
325}
326
327fn choose_target_pod(names: &[String], stack: Stack, leader_id: Option<i32>) -> Option<String> {
328    let prefix = stack.broker_pod_regex().trim_start_matches('^');
329    let mut matching: Vec<&String> = names.iter().filter(|n| n.starts_with(prefix)).collect();
330    matching.sort();
331
332    if let Some(id) = leader_id {
333        let broker_suffix = format!("-{id}");
334        let nodepool_suffix = format!("-{id}-0");
335        if let Some(name) = matching
336            .iter()
337            .find(|n| n.ends_with(&broker_suffix) || n.ends_with(&nodepool_suffix))
338        {
339            return Some((*name).clone());
340        }
341    }
342
343    matching.first().map(|n| (*n).clone())
344}
345
346#[cfg(test)]
347mod tests {
348    use {
349        super::*,
350        krabka_protocol::owned::metadata_response::{
351            MetadataResponse, MetadataResponsePartition, MetadataResponseTopic,
352        },
353        std::{sync, time::Duration},
354        tokio::{sync::oneshot, task, time},
355    };
356
357    #[tokio::test(start_paused = true)]
358    async fn stopping_before_fault_deadline_prevents_later_deletion() {
359        let kill_at = Arc::new(AtomicU64::new(0));
360        let deletes = Arc::new(AtomicU64::new(0));
361        let observed = sync::Arc::clone(&deletes);
362        let marker = sync::Arc::clone(&kill_at);
363        let task = FaultTask::schedule(Instant::now() + Duration::from_secs(60), async move {
364            execute_fault(
365                true,
366                async { Ok(()) },
367                async { Ok(2) },
368                |(), _| async move {
369                    observed.fetch_add(1, Ordering::SeqCst);
370                    Ok("broker-2".into())
371                },
372                marker,
373            )
374            .await
375        });
376        task::yield_now().await;
377        time::advance(Duration::from_secs(5)).await;
378        assert2::assert!(task.stop().await == FaultOutcome::Cancelled);
379        time::advance(Duration::from_secs(120)).await;
380        assert2::assert!(deletes.load(Ordering::SeqCst) == 0);
381        assert2::assert!(kill_at.load(Ordering::SeqCst) == 0);
382    }
383
384    #[tokio::test(start_paused = true)]
385    async fn dropping_fault_owner_cancels_scheduled_operation() {
386        let operations = Arc::new(AtomicU64::new(0));
387        let observed = sync::Arc::clone(&operations);
388        let task = FaultTask::schedule(Instant::now() + Duration::from_secs(30), async move {
389            observed.fetch_add(1, Ordering::SeqCst);
390            FaultOutcome::ClientUnavailable("injected".into())
391        });
392        task::yield_now().await;
393        drop(task);
394        task::yield_now().await;
395        time::advance(Duration::from_secs(60)).await;
396        assert2::assert!(operations.load(Ordering::SeqCst) == 0);
397    }
398
399    #[tokio::test(start_paused = true)]
400    async fn stopping_cancels_and_joins_an_in_flight_fault() {
401        let kill_at = Arc::new(AtomicU64::new(0));
402        let marker = sync::Arc::clone(&kill_at);
403        let (entered_tx, entered_rx) = oneshot::channel();
404        let (release_tx, release_rx) = oneshot::channel();
405        let task = FaultTask::schedule(Instant::now(), async move {
406            execute_fault(
407                true,
408                async { Ok(()) },
409                async { Ok(2) },
410                |(), _| async move {
411                    assert2::assert!(entered_tx.send(()).is_ok());
412                    release_rx.await.context("release injected deletion")?;
413                    Ok("broker".into())
414                },
415                marker,
416            )
417            .await
418        });
419        entered_rx.await.expect("fault operation entered");
420        assert2::assert!(task.stop().await == FaultOutcome::Cancelled);
421        // Joining cancellation drops the operation, including its receiver.
422        assert2::assert!(release_tx.send(()).is_err());
423        assert2::assert!(kill_at.load(Ordering::SeqCst) == 0);
424    }
425
426    #[tokio::test]
427    async fn fault_outcomes_preserve_failures_and_only_mark_successful_deletions() {
428        for (discovery_fails, lookup_fails, deletion_fails, expected_note) in [
429            (
430                true,
431                false,
432                false,
433                Some("failover: in-cluster client unavailable: discovery"),
434            ),
435            (
436                false,
437                true,
438                false,
439                Some(
440                    "failover: leader lookup failed; using first broker pod: lookup; killed broker broker",
441                ),
442            ),
443            (
444                false,
445                false,
446                true,
447                Some("failover: broker kill failed: deletion"),
448            ),
449            (
450                false,
451                true,
452                true,
453                Some(
454                    "failover: leader lookup failed; using first broker pod: lookup; broker kill failed: deletion",
455                ),
456            ),
457            (false, false, false, None),
458        ] {
459            let kill_at = Arc::new(AtomicU64::new(0));
460            let deletes = Arc::new(AtomicU64::new(0));
461            let observed = sync::Arc::clone(&deletes);
462            let outcome = execute_fault(
463                true,
464                async {
465                    if discovery_fails {
466                        Err(anyhow!("discovery"))
467                    } else {
468                        Ok(())
469                    }
470                },
471                async {
472                    if lookup_fails {
473                        Err(anyhow!("lookup"))
474                    } else {
475                        Ok(2)
476                    }
477                },
478                |(), leader| async move {
479                    assert2::assert!(leader == if lookup_fails { None } else { Some(2) });
480                    observed.fetch_add(1, Ordering::SeqCst);
481                    if deletion_fails {
482                        Err(anyhow!("deletion"))
483                    } else {
484                        Ok("broker".into())
485                    }
486                },
487                sync::Arc::clone(&kill_at),
488            )
489            .await;
490            assert2::assert!(outcome.note().as_deref() == expected_note);
491            if discovery_fails {
492                assert2::assert!(outcome == FaultOutcome::ClientUnavailable("discovery".into()));
493                assert2::assert!(deletes.load(Ordering::SeqCst) == 0);
494            } else {
495                let deletion = if deletion_fails {
496                    Err("deletion".into())
497                } else {
498                    Ok(KillEvidence {
499                        pod: "broker".into(),
500                        at_ms: kill_at.load(Ordering::SeqCst),
501                    })
502                };
503                assert2::assert!(
504                    outcome
505                        == FaultOutcome::Finished {
506                            leader_lookup_error: lookup_fails.then(|| "lookup".into()),
507                            deletion,
508                        }
509                );
510                assert2::assert!(deletes.load(Ordering::SeqCst) == 1);
511            }
512            assert2::assert!(
513                (kill_at.load(Ordering::SeqCst) != 0) == (!discovery_fails && !deletion_fails)
514            );
515            assert2::assert!(
516                outcome.killed_at_ms().is_some() == (!discovery_fails && !deletion_fails)
517            );
518        }
519    }
520
521    #[tokio::test(start_paused = true)]
522    async fn completed_fault_evidence_survives_stop() {
523        let task = FaultTask::schedule(Instant::now(), async {
524            FaultOutcome::Finished {
525                leader_lookup_error: None,
526                deletion: Ok(KillEvidence {
527                    pod: "broker".into(),
528                    at_ms: 42,
529                }),
530            }
531        });
532        task::yield_now().await;
533        task::yield_now().await;
534        assert2::assert!(task.stop().await.killed_at_ms() == Some(42));
535    }
536
537    #[test]
538    fn partition0_leader_is_read_from_metadata() {
539        let md = MetadataResponse {
540            topics: vec![MetadataResponseTopic {
541                name: Some("bench-topic".into()),
542                partitions: vec![
543                    MetadataResponsePartition {
544                        partition_index: 1,
545                        leader_id: 0,
546                        ..Default::default()
547                    },
548                    MetadataResponsePartition {
549                        partition_index: 0,
550                        leader_id: 2,
551                        ..Default::default()
552                    },
553                ],
554                ..Default::default()
555            }],
556            ..Default::default()
557        };
558
559        assert2::assert!(partition0_leader_id(&md, "bench-topic") == Some(2));
560    }
561
562    #[test]
563    fn target_pod_uses_actual_leader_id_for_each_stack_naming_shape() {
564        let krabka_multi = vec![
565            "demo-broker-0-0".to_owned(),
566            "demo-broker-1-0".to_owned(),
567            "demo-broker-2-0".to_owned(),
568        ];
569        assert2::assert!(
570            choose_target_pod(&krabka_multi, Stack::Krabka, Some(1)).as_deref()
571                == Some("demo-broker-1-0")
572        );
573
574        let krabka_single_pool = vec![
575            "demo-brokers-0".to_owned(),
576            "demo-brokers-1".to_owned(),
577            "demo-brokers-2".to_owned(),
578        ];
579        assert2::assert!(
580            choose_target_pod(&krabka_single_pool, Stack::Krabka, Some(2)).as_deref()
581                == Some("demo-brokers-2")
582        );
583
584        let kafka = vec![
585            "demo-kafka-0".to_owned(),
586            "demo-kafka-1".to_owned(),
587            "demo-kafka-2".to_owned(),
588        ];
589        assert2::assert!(
590            choose_target_pod(&kafka, Stack::Kafka, Some(1)).as_deref() == Some("demo-kafka-1")
591        );
592    }
593
594    #[test]
595    fn target_pod_falls_back_to_first_matching_broker_when_leader_unknown() {
596        let pods = vec![
597            "demo-kafka-2".to_owned(),
598            "demo-kafka-0".to_owned(),
599            "demo-kafka-1".to_owned(),
600        ];
601
602        assert2::assert!(
603            choose_target_pod(&pods, Stack::Kafka, None).as_deref() == Some("demo-kafka-0")
604        );
605    }
606}
krabka_bench_driver/failover.rs

krabka_bench_driver/
failover.rs
