pingora_load_balancing/
background.rs

1// Copyright 2026 Cloudflare, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Implement [BackgroundService] for [LoadBalancer]
16
17use std::future::Future;
18use std::time::{Duration, Instant};
19
20use super::{BackendIter, BackendSelection, HealthCheckService, LoadBalancer, LoadBalancerGroup};
21use async_trait::async_trait;
22use pingora_core::{
23    server::ShutdownWatch,
24    services::{background::BackgroundService, ServiceReadyNotifier},
25};
26
27/// Why a [`ShutdownWatch`] woke a service loop.
28enum WatchChange {
29    /// Shutdown was signaled.
30    ShuttingDown,
31    /// The watch changed to a non-shutdown value. Nothing needs to stop.
32    Spurious,
33    /// Every sender was dropped, so shutdown can never be signaled from here.
34    ///
35    /// Distinct from [`WatchChange::ShuttingDown`] on purpose. Treating a
36    /// dropped sender as shutdown would make a service abandon work that
37    /// nothing is ever going to ask it to stop doing.
38    Closed,
39}
40
41/// Await the next change to `shutdown` and classify it.
42async fn shutdown_changed(shutdown: &mut ShutdownWatch) -> WatchChange {
43    match shutdown.changed().await {
44        // Read the latest value rather than the one that triggered the change,
45        // so a `false` immediately followed by a `true` still stops the service.
46        Ok(()) if *shutdown.borrow() => WatchChange::ShuttingDown,
47        Ok(()) => WatchChange::Spurious,
48        Err(_) => WatchChange::Closed,
49    }
50}
51
52/// The outcome of racing a unit of background work against shutdown.
53///
54/// Returned by [`race_shutdown`]. Matching this exhaustively keeps each service
55/// loop explicit about both outcomes.
56enum Raced<T> {
57    /// The work ran to completion, yielding its output.
58    Completed(T),
59    /// Shutdown was signaled. The work was dropped mid-flight and the service
60    /// loop should return.
61    ShuttingDown,
62}
63
64/// Run `work`, returning early if shutdown is signaled.
65///
66/// Shutdown is polled first, but [`tokio::sync::watch::Receiver::changed`] only
67/// reports versions this receiver has not already seen, so callers must check
68/// `*shutdown.borrow()` themselves as well.
69///
70/// A shutdown drops `work` mid-flight, cancelling an in-flight discovery or
71/// health check. Any other outcome resumes the *same* `work` future rather than
72/// restarting it, so repeated watch changes can neither starve it nor repeat
73/// effects it already applied.
74async fn race_shutdown<F: Future>(shutdown: &mut ShutdownWatch, work: F) -> Raced<F::Output> {
75    tokio::pin!(work);
76    loop {
77        tokio::select! {
78            biased;
79            change = shutdown_changed(&mut *shutdown) => match change {
80                WatchChange::ShuttingDown => return Raced::ShuttingDown,
81                WatchChange::Spurious => continue,
82                // Nothing can signal shutdown any more, so stop racing
83                // altogether and let the work finish.
84                WatchChange::Closed => break,
85            },
86            output = &mut work => return Raced::Completed(output),
87        }
88    }
89    Raced::Completed(work.await)
90}
91
92/// Timing state shared by update and health-check loops.
93struct TaskSchedule {
94    frequency: Option<Duration>,
95    next: Option<Instant>,
96}
97
98/// The independently optional tasks run by a background service.
99struct BackgroundSchedule {
100    update: Option<TaskSchedule>,
101    health_check: Option<TaskSchedule>,
102}
103
104impl TaskSchedule {
105    fn new(frequency: Option<Duration>, now: Instant) -> Self {
106        Self {
107            frequency,
108            next: Some(now),
109        }
110    }
111
112    fn is_due(&self, now: Instant) -> bool {
113        self.next.is_some_and(|next| next <= now)
114    }
115
116    fn finished_at(&mut self, base: Instant) {
117        self.next = self.frequency.map(|frequency| base + frequency);
118    }
119
120    fn next(&self) -> Option<Instant> {
121        self.next
122    }
123
124    fn is_once(&self) -> bool {
125        self.frequency.is_none()
126    }
127}
128
129impl BackgroundSchedule {
130    // 136 years, used when no scheduled work remains but another event can wake the service.
131    const NEVER: Duration = Duration::from_secs(u32::MAX as u64);
132
133    fn new(
134        update_frequency: Option<Option<Duration>>,
135        health_check_frequency: Option<Option<Duration>>,
136        now: Instant,
137    ) -> Self {
138        Self {
139            update: update_frequency.map(|frequency| TaskSchedule::new(frequency, now)),
140            health_check: health_check_frequency.map(|frequency| TaskSchedule::new(frequency, now)),
141        }
142    }
143
144    fn update_is_due(&self, now: Instant) -> bool {
145        self.update
146            .as_ref()
147            .is_some_and(|schedule| schedule.is_due(now))
148    }
149
150    fn health_check_is_due(&self, now: Instant) -> bool {
151        self.health_check
152            .as_ref()
153            .is_some_and(|schedule| schedule.is_due(now))
154    }
155
156    fn update_finished_at(&mut self, base: Instant) {
157        if let Some(schedule) = self.update.as_mut() {
158            schedule.finished_at(base);
159        }
160    }
161
162    fn health_check_finished_at(&mut self, base: Instant) {
163        if let Some(schedule) = self.health_check.as_mut() {
164            schedule.finished_at(base);
165        }
166    }
167
168    fn next_scheduled(&self) -> Option<Instant> {
169        self.update
170            .iter()
171            .chain(self.health_check.iter())
172            .filter_map(TaskSchedule::next)
173            .min()
174    }
175
176    fn next(&self, now: Instant) -> Instant {
177        self.next_scheduled().unwrap_or(now + Self::NEVER)
178    }
179
180    fn next_health_check(&self) -> Option<Instant> {
181        self.health_check.as_ref().and_then(TaskSchedule::next)
182    }
183
184    fn health_check_is_once(&self) -> bool {
185        self.health_check
186            .as_ref()
187            .is_some_and(TaskSchedule::is_once)
188    }
189
190    fn is_idle(&self) -> bool {
191        self.next_scheduled().is_none()
192    }
193}
194
195impl<S: Send + Sync + BackendSelection + 'static> LoadBalancer<S>
196where
197    S::Iter: BackendIter,
198{
199    pub async fn run(
200        &self,
201        mut shutdown: ShutdownWatch,
202        mut ready_opt: Option<ServiceReadyNotifier>,
203    ) -> () {
204        let mut now = Instant::now();
205        let mut schedule = BackgroundSchedule::new(
206            Some(self.update_frequency),
207            // Private views schedule probes here. Shared views are probed once by
208            // their registry's HealthCheckService.
209            self.backends
210                .owns_health_checks()
211                .then_some(self.health_check_frequency),
212            now,
213        );
214        loop {
215            if *shutdown.borrow() {
216                return;
217            }
218
219            if schedule.update_is_due(now) {
220                // TODO: log err
221                match race_shutdown(&mut shutdown, self.update()).await {
222                    Raced::ShuttingDown => return,
223                    Raced::Completed(_) => schedule.update_finished_at(now),
224                }
225            }
226
227            // After the first update, discovery and selection setup will be
228            // done, so dependent services can start receiving traffic.
229            if let Some(ready) = ready_opt.take() {
230                ServiceReadyNotifier::notify_ready(ready)
231            }
232
233            if schedule.health_check_is_due(now) {
234                let health_check = self.backends.run_health_check(self.parallel_health_check);
235                match race_shutdown(&mut shutdown, health_check).await {
236                    Raced::ShuttingDown => return,
237                    Raced::Completed(()) => schedule.health_check_finished_at(now),
238                }
239            }
240
241            if schedule.is_idle() {
242                return;
243            }
244            let to_wake = schedule.next(now);
245            match race_shutdown(&mut shutdown, tokio::time::sleep_until(to_wake.into())).await {
246                Raced::ShuttingDown => return,
247                Raced::Completed(()) => {}
248            }
249            now = Instant::now();
250        }
251    }
252}
253
254/// Implement [BackgroundService] for [LoadBalancer]. For backward-compatibility
255/// reasons, we implement both the `start` and `start_with_ready_notifier`
256/// methods.
257#[async_trait]
258impl<S: Send + Sync + BackendSelection + 'static> BackgroundService for LoadBalancer<S>
259where
260    S::Iter: BackendIter,
261{
262    async fn start_with_ready_notifier(
263        &self,
264        shutdown: pingora_core::server::ShutdownWatch,
265        ready: ServiceReadyNotifier,
266    ) -> () {
267        self.run(shutdown, Some(ready)).await
268    }
269
270    async fn start(&self, shutdown: pingora_core::server::ShutdownWatch) -> () {
271        self.run(shutdown, None).await
272    }
273}
274
275impl<S: Send + Sync + BackendSelection + 'static> LoadBalancerGroup<S>
276where
277    S::Config: 'static,
278    S::Iter: BackendIter,
279{
280    /// Run discovery, selector rebuilds, and privately managed health checks
281    /// until shutdown.
282    pub async fn run(
283        &self,
284        mut shutdown: pingora_core::server::ShutdownWatch,
285        mut ready_opt: Option<ServiceReadyNotifier>,
286    ) {
287        let mut now = Instant::now();
288        let mut schedule = BackgroundSchedule::new(
289            Some(self.update_frequency),
290            // Private groups schedule probes here. Shared groups only consume the
291            // health state maintained by their registry's HealthCheckService.
292            self.backends()
293                .owns_health_checks()
294                .then_some(self.health_check_frequency),
295            now,
296        );
297        let mut ready_generation = None;
298
299        loop {
300            if *shutdown.borrow() {
301                return;
302            }
303
304            if schedule.update_is_due(now) {
305                match race_shutdown(&mut shutdown, self.update()).await {
306                    Raced::ShuttingDown => return,
307                    Raced::Completed(Ok(())) => {
308                        // Until readiness is signaled, always target the latest
309                        // backend generation. Pinning to the first successful
310                        // update would let a burst of updates satisfy readiness
311                        // with selectors that are already stale relative to the
312                        // current membership.
313                        if ready_opt.is_some() {
314                            ready_generation = Some(self.backend_generation());
315                        }
316                        schedule.update_finished_at(now);
317                    }
318                    Raced::Completed(Err(error)) => {
319                        log::error!("load balancer group update failed: {error}");
320                        schedule.update_finished_at(now);
321                    }
322                }
323            }
324
325            if ready_generation.is_some_and(|generation| self.selectors_ready_for(generation)) {
326                if let Some(ready) = ready_opt.take() {
327                    // Every selector reached the target generation, so dependent
328                    // services can start receiving traffic.
329                    ServiceReadyNotifier::notify_ready(ready)
330                }
331            }
332
333            if schedule.health_check_is_due(now) {
334                let health_check = self.backends().run_health_check(self.parallel_health_check);
335                match race_shutdown(&mut shutdown, health_check).await {
336                    Raced::ShuttingDown => return,
337                    Raced::Completed(()) => schedule.health_check_finished_at(now),
338                }
339            }
340
341            // Discovery and health checks have independent schedules. One-shot
342            // discovery can still feed recurring checks over its last membership.
343            if ready_opt.is_none() && schedule.is_idle() {
344                // No readiness notification or periodic work remains.
345                // Any queued selector rebuilds finish independently.
346                return;
347            }
348
349            let to_wake = schedule.next(now);
350            // Selector completion may satisfy startup readiness before the next
351            // scheduled task. After readiness, rebuilds no longer wake this loop.
352            let wake = async {
353                tokio::select! {
354                    _ = tokio::time::sleep_until(to_wake.into()) => {}
355                    _ = self.rebuild_notified(), if ready_opt.is_some() => {} // re-trigger loop to notify waiters
356                }
357            };
358            match race_shutdown(&mut shutdown, wake).await {
359                Raced::ShuttingDown => return,
360                Raced::Completed(()) => {}
361            }
362            now = Instant::now();
363        }
364    }
365}
366
367#[async_trait]
368impl<S: Send + Sync + BackendSelection + 'static> BackgroundService for LoadBalancerGroup<S>
369where
370    S::Config: 'static,
371    S::Iter: BackendIter,
372{
373    async fn start_with_ready_notifier(
374        &self,
375        shutdown: pingora_core::server::ShutdownWatch,
376        ready: ServiceReadyNotifier,
377    ) {
378        self.run(shutdown, Some(ready)).await
379    }
380
381    async fn start(&self, shutdown: pingora_core::server::ShutdownWatch) {
382        self.run(shutdown, None).await
383    }
384}
385
386impl HealthCheckService {
387    /// Run health checks for the shared registry until shutdown.
388    pub async fn run(
389        &self,
390        mut shutdown: pingora_core::server::ShutdownWatch,
391        mut ready_opt: Option<ServiceReadyNotifier>,
392    ) {
393        if !self.registry.has_health_check() {
394            log::error!("HealthCheckService requires a configured HealthRegistry health check");
395            // Keep the notifier alive so dependents cannot observe a false-ready
396            // signal from its Drop implementation. A non-shutdown change must
397            // not release it early; a closed watch has no server left to
398            // mislead, so parking on it would only leak this task.
399            if ready_opt.is_some() && !*shutdown.borrow() {
400                while let WatchChange::Spurious = shutdown_changed(&mut shutdown).await {}
401            }
402            return;
403        }
404
405        let mut schedule =
406            BackgroundSchedule::new(None, Some(self.health_check_frequency), Instant::now());
407        loop {
408            if *shutdown.borrow() {
409                return;
410            }
411
412            // One-shot mode (no frequency) runs a single pass and returns, so it
413            // must observe at least one target first. Otherwise it would check
414            // an empty registry, signal ready, and never look at targets
415            // published afterwards. Periodic mode instead signals ready eagerly
416            // and relies on later passes to pick up newly published targets.
417            if schedule.health_check_is_once() && self.registry.target_count() == 0 {
418                match race_shutdown(&mut shutdown, self.registry.wait_for_targets()).await {
419                    Raced::ShuttingDown => return,
420                    Raced::Completed(()) => {}
421                }
422                // A shutdown already observed by an earlier `changed()` will not
423                // wake the race above, so re-check the current value before
424                // starting a pass.
425                if *shutdown.borrow() {
426                    return;
427                }
428            }
429
430            let health_check = self.registry.run_health_check(self.parallel_health_check);
431            match race_shutdown(&mut shutdown, health_check).await {
432                Raced::ShuttingDown => return,
433                Raced::Completed(()) => {
434                    if let Some(ready) = ready_opt.take() {
435                        // The initial health pass completed, so services depending on
436                        // this registry can start receiving traffic.
437                        ServiceReadyNotifier::notify_ready(ready);
438                    }
439                    schedule.health_check_finished_at(Instant::now());
440                }
441            }
442
443            let Some(next_health_check) = schedule.next_health_check() else {
444                // no more checks
445                return;
446            };
447            loop {
448                let has_targets = self.registry.target_count() > 0;
449                // `true` means the next pass is due; a view removal only
450                // re-evaluates whether the registry still has targets.
451                let wake = async {
452                    tokio::select! {
453                        _ = self.registry.wait_for_targets(), if !has_targets => true,
454                        _ = tokio::time::sleep_until(next_health_check.into()), if has_targets => true,
455                        _ = self.registry.wait_for_view_removal() => false,
456                    }
457                };
458                match race_shutdown(&mut shutdown, wake).await {
459                    Raced::ShuttingDown => return,
460                    Raced::Completed(true) => break,
461                    Raced::Completed(false) => {}
462                }
463            }
464        }
465    }
466}
467
468#[async_trait]
469impl BackgroundService for HealthCheckService {
470    async fn start_with_ready_notifier(
471        &self,
472        shutdown: pingora_core::server::ShutdownWatch,
473        ready: ServiceReadyNotifier,
474    ) {
475        self.run(shutdown, Some(ready)).await
476    }
477
478    async fn start(&self, shutdown: pingora_core::server::ShutdownWatch) {
479        self.run(shutdown, None).await
480    }
481}
482
483#[cfg(test)]
484mod tests {
485    use std::collections::{BTreeSet, HashMap};
486    use std::future;
487    use std::sync::Arc;
488
489    use async_trait::async_trait;
490    use pingora_error::Result;
491    use tokio::sync::{watch, Notify};
492
493    use super::*;
494    use crate::discovery::{ServiceDiscovery, Static};
495    use crate::health_check::HealthCheck;
496    use crate::selection;
497    use crate::{Backend, Backends, HealthRegistry};
498
499    #[test]
500    fn schedule_without_tasks_is_idle() {
501        let now = Instant::now();
502        let mut schedule = BackgroundSchedule::new(None, None, now);
503
504        assert!(!schedule.update_is_due(now));
505        assert!(!schedule.health_check_is_due(now));
506        schedule.update_finished_at(now);
507        schedule.health_check_finished_at(now);
508        assert!(schedule.is_idle());
509        assert_eq!(schedule.next(now), now + BackgroundSchedule::NEVER);
510        assert_eq!(schedule.next_health_check(), None);
511        assert!(!schedule.health_check_is_once());
512    }
513
514    #[test]
515    fn one_time_update_schedule_finishes_after_first_run() {
516        let now = Instant::now();
517        let mut schedule = BackgroundSchedule::new(Some(None), None, now);
518
519        assert!(schedule.update_is_due(now));
520        assert!(!schedule.is_idle());
521        schedule.update_finished_at(now);
522        assert!(!schedule.update_is_due(now));
523        assert!(schedule.is_idle());
524        assert_eq!(schedule.next(now), now + BackgroundSchedule::NEVER);
525    }
526
527    #[test]
528    fn periodic_health_check_schedule_tracks_next_pass() {
529        let now = Instant::now();
530        let frequency = Duration::from_secs(5);
531        let mut schedule = BackgroundSchedule::new(None, Some(Some(frequency)), now);
532
533        assert!(schedule.health_check_is_due(now));
534        assert!(!schedule.health_check_is_once());
535        schedule.health_check_finished_at(now);
536        assert!(!schedule.health_check_is_due(now));
537        assert!(schedule.health_check_is_due(now + frequency));
538        assert!(!schedule.is_idle());
539        assert_eq!(schedule.next_health_check(), Some(now + frequency));
540    }
541
542    #[test]
543    fn one_time_health_check_schedule_finishes_after_first_run() {
544        let now = Instant::now();
545        let mut schedule = BackgroundSchedule::new(None, Some(None), now);
546
547        assert!(schedule.health_check_is_once());
548        assert!(schedule.health_check_is_due(now));
549        schedule.health_check_finished_at(now);
550        assert!(!schedule.health_check_is_due(now));
551        assert_eq!(schedule.next_health_check(), None);
552        assert!(schedule.is_idle());
553        assert_eq!(schedule.next(now), now + BackgroundSchedule::NEVER);
554    }
555
556    #[test]
557    fn schedule_uses_earliest_task_deadline() {
558        let now = Instant::now();
559        let update_frequency = Duration::from_secs(10);
560        let health_check_frequency = Duration::from_secs(5);
561        let mut schedule = BackgroundSchedule::new(
562            Some(Some(update_frequency)),
563            Some(Some(health_check_frequency)),
564            now,
565        );
566
567        schedule.update_finished_at(now);
568        schedule.health_check_finished_at(now);
569        assert_eq!(schedule.next(now), now + health_check_frequency);
570        assert!(!schedule.update_is_due(now + health_check_frequency));
571        assert!(schedule.health_check_is_due(now + health_check_frequency));
572        assert!(schedule.update_is_due(now + update_frequency));
573    }
574
575    struct NotifyingDiscovery {
576        notify: Arc<Notify>,
577    }
578
579    #[async_trait]
580    impl ServiceDiscovery for NotifyingDiscovery {
581        async fn discover(&self) -> Result<(BTreeSet<Backend>, HashMap<u64, bool>)> {
582            self.notify.notify_one();
583            Ok((BTreeSet::new(), HashMap::new()))
584        }
585    }
586
587    struct PendingDiscovery {
588        notify: Arc<Notify>,
589    }
590
591    #[async_trait]
592    impl ServiceDiscovery for PendingDiscovery {
593        async fn discover(&self) -> Result<(BTreeSet<Backend>, HashMap<u64, bool>)> {
594            self.notify.notify_one();
595            future::pending().await
596        }
597    }
598
599    struct PendingHealthCheck {
600        notify: Arc<Notify>,
601        drop_notify: Arc<Notify>,
602    }
603
604    #[async_trait]
605    impl HealthCheck for PendingHealthCheck {
606        async fn check(&self, _target: &Backend) -> Result<()> {
607            struct NotifyOnDrop(Arc<Notify>);
608
609            impl Drop for NotifyOnDrop {
610                fn drop(&mut self) {
611                    self.0.notify_one();
612                }
613            }
614
615            let _notify_on_drop = NotifyOnDrop(self.drop_notify.clone());
616            self.notify.notify_one();
617            future::pending().await
618        }
619
620        fn health_threshold(&self, _success: bool) -> usize {
621            1
622        }
623    }
624
625    async fn assert_run_exits_on_shutdown(
626        lb: LoadBalancer<selection::RoundRobin>,
627        notify: Arc<Notify>,
628    ) {
629        let (shutdown_tx, shutdown_rx) = watch::channel(false);
630        let handle = tokio::spawn(async move {
631            lb.run(shutdown_rx, None).await;
632        });
633
634        notify.notified().await;
635        shutdown_tx.send(true).unwrap();
636
637        tokio::time::timeout(Duration::from_secs(1), handle)
638            .await
639            .expect("background service should observe shutdown promptly")
640            .expect("background service task should not panic");
641    }
642
643    #[tokio::test]
644    async fn run_returns_when_shutdown_while_sleeping() {
645        let notify = Arc::new(Notify::new());
646        let discovery = NotifyingDiscovery {
647            notify: notify.clone(),
648        };
649        let mut lb = LoadBalancer::<selection::RoundRobin>::from_backends(Backends::new(Box::new(
650            discovery,
651        )));
652        lb.update_frequency = Some(Duration::from_secs(60));
653
654        assert_run_exits_on_shutdown(lb, notify).await;
655    }
656
657    #[tokio::test]
658    async fn run_updates_when_no_shutdown_sender_remains() {
659        let notify = Arc::new(Notify::new());
660        let discovery = NotifyingDiscovery {
661            notify: notify.clone(),
662        };
663        let mut lb = LoadBalancer::<selection::RoundRobin>::from_backends(Backends::new(Box::new(
664            discovery,
665        )));
666        lb.update_frequency = Some(Duration::from_secs(60));
667
668        let (shutdown_tx, shutdown_rx) = watch::channel(false);
669        // No sender remains, so shutdown can never be signaled. That must not
670        // be read as "already shutting down": the service would return before
671        // ever running discovery, and its dropped notifier would signal ready.
672        drop(shutdown_tx);
673
674        let handle = tokio::spawn(async move { lb.run(shutdown_rx, None).await });
675
676        tokio::time::timeout(Duration::from_secs(1), notify.notified())
677            .await
678            .expect("discovery should still run with no shutdown sender left");
679        handle.abort();
680    }
681
682    #[tokio::test]
683    async fn run_continues_after_a_non_shutdown_watch_change() {
684        let notify = Arc::new(Notify::new());
685        let discovery = NotifyingDiscovery {
686            notify: notify.clone(),
687        };
688        let mut lb = LoadBalancer::<selection::RoundRobin>::from_backends(Backends::new(Box::new(
689            discovery,
690        )));
691        lb.update_frequency = Some(Duration::from_secs(60));
692
693        let (shutdown_tx, shutdown_rx) = watch::channel(false);
694        let handle = tokio::spawn(async move { lb.run(shutdown_rx, None).await });
695
696        // Let the first update land so the loop is parked in the sleep race.
697        notify.notified().await;
698        shutdown_tx.send(false).unwrap();
699
700        // Yield generously so a service that was going to return has run.
701        tokio::time::sleep(Duration::from_millis(50)).await;
702        assert!(
703            !handle.is_finished(),
704            "a non-shutdown watch change stopped the service"
705        );
706
707        shutdown_tx.send(true).unwrap();
708        tokio::time::timeout(Duration::from_secs(1), handle)
709            .await
710            .expect("a real shutdown should still stop the service")
711            .expect("background service task should not panic");
712    }
713
714    #[tokio::test]
715    async fn run_returns_when_shutdown_while_updating() {
716        let notify = Arc::new(Notify::new());
717        let discovery = PendingDiscovery {
718            notify: notify.clone(),
719        };
720        let lb = LoadBalancer::<selection::RoundRobin>::from_backends(Backends::new(Box::new(
721            discovery,
722        )));
723
724        assert_run_exits_on_shutdown(lb, notify).await;
725    }
726
727    #[tokio::test]
728    async fn run_returns_when_shutdown_while_health_checking() {
729        let notify = Arc::new(Notify::new());
730        let drop_notify = Arc::new(Notify::new());
731        let mut lb =
732            LoadBalancer::<selection::RoundRobin>::try_from_iter(["127.0.0.1:80"]).unwrap();
733        lb.set_health_check(Box::new(PendingHealthCheck {
734            notify: notify.clone(),
735            drop_notify: drop_notify.clone(),
736        }));
737
738        assert_run_exits_on_shutdown(lb, notify).await;
739        tokio::time::timeout(Duration::from_secs(1), drop_notify.notified())
740            .await
741            .expect("pending health check should be cancelled");
742    }
743
744    #[tokio::test]
745    async fn run_aborts_parallel_health_check_on_shutdown() {
746        let notify = Arc::new(Notify::new());
747        let drop_notify = Arc::new(Notify::new());
748        let mut lb =
749            LoadBalancer::<selection::RoundRobin>::try_from_iter(["127.0.0.1:80"]).unwrap();
750        lb.parallel_health_check = true;
751        lb.set_health_check(Box::new(PendingHealthCheck {
752            notify: notify.clone(),
753            drop_notify: drop_notify.clone(),
754        }));
755
756        assert_run_exits_on_shutdown(lb, notify).await;
757        tokio::time::timeout(Duration::from_secs(1), drop_notify.notified())
758            .await
759            .expect("parallel health check task should be aborted");
760    }
761
762    async fn assert_group_run_exits_on_shutdown(
763        group: LoadBalancerGroup<selection::RoundRobin>,
764        notify: Arc<Notify>,
765    ) {
766        let (shutdown_tx, shutdown_rx) = watch::channel(false);
767        let handle = tokio::spawn(async move {
768            group.run(shutdown_rx, None).await;
769        });
770
771        notify.notified().await;
772        shutdown_tx.send(true).unwrap();
773
774        tokio::time::timeout(Duration::from_secs(1), handle)
775            .await
776            .expect("background service should observe shutdown promptly")
777            .expect("background service task should not panic");
778    }
779
780    #[tokio::test]
781    async fn group_run_returns_when_shutdown_while_updating() {
782        let notify = Arc::new(Notify::new());
783        let discovery = PendingDiscovery {
784            notify: notify.clone(),
785        };
786        let group = LoadBalancerGroup::<selection::RoundRobin>::from_backends_with_configs(
787            Backends::new(Box::new(discovery)),
788            [None],
789        );
790
791        assert_group_run_exits_on_shutdown(group, notify).await;
792    }
793
794    #[tokio::test]
795    async fn group_run_returns_when_shutdown_while_health_checking() {
796        let notify = Arc::new(Notify::new());
797        let drop_notify = Arc::new(Notify::new());
798        let mut backends = Backends::new(Static::try_from_iter(["127.0.0.1:80"]).unwrap());
799        backends.set_health_check(Box::new(PendingHealthCheck {
800            notify: notify.clone(),
801            drop_notify: drop_notify.clone(),
802        }));
803        let group = LoadBalancerGroup::<selection::RoundRobin>::from_backends_with_configs(
804            backends,
805            [None],
806        );
807
808        assert_group_run_exits_on_shutdown(group, notify).await;
809        tokio::time::timeout(Duration::from_secs(1), drop_notify.notified())
810            .await
811            .expect("pending health check should be cancelled");
812    }
813
814    #[tokio::test]
815    async fn health_check_service_run_returns_when_shutdown_while_health_checking() {
816        let notify = Arc::new(Notify::new());
817        let drop_notify = Arc::new(Notify::new());
818        let registry = Arc::new(HealthRegistry::new());
819        registry.set_health_check(Box::new(PendingHealthCheck {
820            notify: notify.clone(),
821            drop_notify: drop_notify.clone(),
822        }));
823        // The view must outlive the service so its targets stay registered.
824        let view = Backends::new_with_health_registry(
825            Static::try_from_iter(["127.0.0.1:80"]).unwrap(),
826            Arc::clone(&registry),
827        );
828        view.update(|_| {}).await.unwrap();
829
830        let service = HealthCheckService::new(registry);
831        let (shutdown_tx, shutdown_rx) = watch::channel(false);
832        let handle = tokio::spawn(async move { service.run(shutdown_rx, None).await });
833
834        notify.notified().await;
835        shutdown_tx.send(true).unwrap();
836
837        tokio::time::timeout(Duration::from_secs(1), handle)
838            .await
839            .expect("health check service should observe shutdown promptly")
840            .expect("health check service task should not panic");
841        tokio::time::timeout(Duration::from_secs(1), drop_notify.notified())
842            .await
843            .expect("pending health check should be cancelled");
844    }
845}
pingora_load_balancing/background.rs

pingora_load_balancing/
background.rs
