11use std:: collections:: { BTreeSet , HashMap } ;
22use std:: io;
33use std:: net:: SocketAddr ;
4+ use std:: process;
45use std:: sync:: Arc ;
6+ use std:: sync:: Mutex ;
57use std:: time:: Duration ;
68use std:: time:: Instant ;
79
@@ -34,64 +36,52 @@ pub(super) trait FluxHealthCheck: Send + Sync + 'static {
3436 }
3537}
3638
37- #[ derive( Clone ) ]
3839struct FluxBackendHealthInner {
3940 healthy : bool ,
4041 enabled : bool ,
4142 consecutive_counter : usize ,
4243}
4344
44- struct FluxBackendHealth ( ArcSwap < FluxBackendHealthInner > ) ;
45+ #[ derive( Clone ) ]
46+ struct FluxBackendHealth ( Arc < Mutex < FluxBackendHealthInner > > ) ;
4547
4648impl Default for FluxBackendHealth {
4749 fn default ( ) -> Self {
48- Self ( ArcSwap :: new ( Arc :: new ( FluxBackendHealthInner {
50+ Self ( Arc :: new ( Mutex :: new ( FluxBackendHealthInner {
4951 healthy : true ,
5052 enabled : true ,
5153 consecutive_counter : 0 ,
5254 } ) ) )
5355 }
5456}
5557
56- impl Clone for FluxBackendHealth {
57- fn clone ( & self ) -> Self {
58- Self ( ArcSwap :: new ( self . 0 . load_full ( ) ) )
58+ impl FluxBackendHealth {
59+ fn lock ( & self ) -> std :: sync :: MutexGuard < ' _ , FluxBackendHealthInner > {
60+ self . 0 . lock ( ) . unwrap_or_else ( |_| process :: abort ( ) )
5961 }
60- }
6162
62- impl FluxBackendHealth {
6363 fn ready ( & self ) -> bool {
64- let health = self . 0 . load ( ) ;
64+ let health = self . lock ( ) ;
6565 health. healthy && health. enabled
6666 }
6767
6868 fn enable ( & self , enabled : bool ) {
69- let health = self . 0 . load ( ) ;
70- if health. enabled != enabled {
71- let mut next = ( * * health) . clone ( ) ;
72- next. enabled = enabled;
73- self . 0 . store ( Arc :: new ( next) ) ;
74- }
69+ self . lock ( ) . enabled = enabled;
7570 }
7671
7772 fn observe ( & self , healthy : bool , flip_threshold : usize ) -> bool {
78- let health = self . 0 . load ( ) ;
79- let mut flipped = false ;
73+ let mut health = self . lock ( ) ;
8074 if health. healthy != healthy {
81- let mut next = ( * * health) . clone ( ) ;
82- next. consecutive_counter = next. consecutive_counter . saturating_add ( 1 ) ;
83- if next. consecutive_counter >= flip_threshold {
84- next. healthy = healthy;
85- next. consecutive_counter = 0 ;
86- flipped = true ;
75+ health. consecutive_counter = health. consecutive_counter . saturating_add ( 1 ) ;
76+ if health. consecutive_counter >= flip_threshold {
77+ health. healthy = healthy;
78+ health. consecutive_counter = 0 ;
79+ return true ;
8780 }
88- self . 0 . store ( Arc :: new ( next) ) ;
89- } else if health. consecutive_counter > 0 {
90- let mut next = ( * * health) . clone ( ) ;
91- next. consecutive_counter = 0 ;
92- self . 0 . store ( Arc :: new ( next) ) ;
81+ } else {
82+ health. consecutive_counter = 0 ;
9383 }
94- flipped
84+ false
9585 }
9686}
9787
@@ -149,8 +139,8 @@ impl FluxLoadBalancerRuntime {
149139 }
150140 next_health. insert ( key, backend_health) ;
151141 }
152- self . health . store ( Arc :: new ( next_health) ) ;
153142 self . backends . store ( Arc :: new ( new_backends) ) ;
143+ self . health . store ( Arc :: new ( next_health) ) ;
154144 } else {
155145 let health = self . health . load ( ) ;
156146 for ( key, enabled) in enablement {
@@ -187,7 +177,7 @@ impl FluxLoadBalancerRuntime {
187177 if let Err ( error) = self . update ( ) . await {
188178 log:: warn!( "load-balancer discovery update failed: {error}" ) ;
189179 }
190- next_update = now + self . update_frequency . unwrap_or ( NEVER ) ;
180+ next_update = checked_next_wake ( now, self . update_frequency . unwrap_or ( NEVER ) ) ;
191181 }
192182
193183 if let Some ( ready) = ready. take ( ) {
@@ -196,7 +186,8 @@ impl FluxLoadBalancerRuntime {
196186
197187 if next_health_check <= now {
198188 self . run_health_check ( self . parallel_health_check ) . await ;
199- next_health_check = now + self . health_check_frequency . unwrap_or ( NEVER ) ;
189+ next_health_check =
190+ checked_next_wake ( now, self . health_check_frequency . unwrap_or ( NEVER ) ) ;
200191 }
201192
202193 if self . update_frequency . is_none ( ) && self . health_check_frequency . is_none ( ) {
@@ -259,6 +250,11 @@ impl FluxLoadBalancerRuntime {
259250 }
260251}
261252
253+ fn checked_next_wake ( now : Instant , delay : Duration ) -> Instant {
254+ now. checked_add ( delay)
255+ . unwrap_or_else ( || now + Duration :: from_secs ( 3600 ) )
256+ }
257+
262258pub ( crate ) trait BackendIdentity {
263259 fn authority ( & self ) -> String ;
264260
0 commit comments