@@ -5,15 +5,21 @@ use apiserver::{
5
5
RemoveNodeExclusionFromLoadBalancerRequest , UncordonBottlerocketShadowRequest ,
6
6
{ CreateBottlerocketShadowRequest , UpdateBottlerocketShadowRequest } ,
7
7
} ;
8
- use models:: node:: {
9
- BottlerocketShadow , BottlerocketShadowSelector , BottlerocketShadowSpec ,
10
- BottlerocketShadowState , BottlerocketShadowStatus ,
11
- } ;
12
-
13
8
use chrono:: { DateTime , Utc } ;
9
+ use governor:: {
10
+ clock:: DefaultClock ,
11
+ middleware:: NoOpMiddleware ,
12
+ state:: { InMemoryState , NotKeyed } ,
13
+ Quota , RateLimiter ,
14
+ } ;
14
15
use k8s_openapi:: api:: core:: v1:: Node ;
15
16
use kube:: runtime:: reflector:: Store ;
16
17
use kube:: Api ;
18
+ use lazy_static:: lazy_static;
19
+ use models:: node:: {
20
+ BottlerocketShadow , BottlerocketShadowSelector , BottlerocketShadowSpec ,
21
+ BottlerocketShadowState , BottlerocketShadowStatus ,
22
+ } ;
17
23
use snafu:: { OptionExt , ResultExt } ;
18
24
use std:: env;
19
25
use tokio:: time:: { sleep, Duration } ;
@@ -30,6 +36,14 @@ const RETRY_MAX_DELAY: Duration = Duration::from_secs(30);
30
36
const NUM_RETRIES : usize = 5 ;
31
37
32
38
const AGENT_SLEEP_DURATION : Duration = Duration :: from_secs ( 5 ) ;
39
+ const UPDATE_CHECK_INTERVAL : Duration = Duration :: from_secs ( 120 ) ;
40
+
41
+ type SimpleRateLimiter = RateLimiter < NotKeyed , InMemoryState , DefaultClock , NoOpMiddleware > ;
42
+
43
+ lazy_static ! {
44
+ static ref CHECK_UPDATE_RATE_LIMITER : SimpleRateLimiter =
45
+ RateLimiter :: direct( Quota :: with_period( UPDATE_CHECK_INTERVAL ) . unwrap( ) ) ;
46
+ }
33
47
34
48
/// The module-wide result type.
35
49
pub type Result < T > = std:: result:: Result < T , agentclient_error:: Error > ;
@@ -189,14 +203,23 @@ impl<T: APIServerClient> BrupopAgent<T> {
189
203
let os_info = apiclient:: get_os_info ( )
190
204
. await
191
205
. context ( agentclient_error:: BottlerocketShadowStatusVersionSnafu ) ?;
192
- let update_version = match apiclient:: get_chosen_update ( )
193
- . await
194
- . context ( agentclient_error:: BottlerocketShadowStatusChosenUpdateSnafu ) ?
195
- {
196
- Some ( chosen_update) => chosen_update. version ,
197
- // if chosen update is null which means current node already in latest version, assign current version value to it.
198
- _ => os_info. version_id . clone ( ) ,
199
- } ;
206
+
207
+ // If enough time has elapsed, ask the Bottlerocket API if any updates are available.
208
+ let update_version = if CHECK_UPDATE_RATE_LIMITER . check ( ) . is_ok ( ) {
209
+ event ! ( Level :: INFO , "Checking for Bottlerocket updates." ) ;
210
+ apiclient:: get_chosen_update ( )
211
+ . await
212
+ . context ( agentclient_error:: BottlerocketShadowStatusChosenUpdateSnafu ) ?
213
+ . map ( |chosen_update| chosen_update. version )
214
+ } else {
215
+ // Rate limited, don't check for updates, just return whatever version we last returned.
216
+ self . fetch_shadow ( )
217
+ . await ?
218
+ . status
219
+ . map ( |status| status. target_version ( ) )
220
+ }
221
+ // If we can't determine a version to update to, just return our current version
222
+ . unwrap_or ( os_info. version_id . clone ( ) ) ;
200
223
201
224
Ok ( BottlerocketShadowStatus :: new (
202
225
os_info. version_id . clone ( ) ,
@@ -572,7 +595,8 @@ impl<T: APIServerClient> BrupopAgent<T> {
572
595
_ => {
573
596
event ! (
574
597
Level :: WARN ,
575
- "An error occurred when invoking Bottlerocket Update API"
598
+ "An error occurred when invoking Bottlerocket Update API: '{}'" ,
599
+ e
576
600
) ;
577
601
match self
578
602
. update_status_in_shadow (
0 commit comments