From b7e3b4c84e6634e64f543ab6f043b4fac2f76e58 Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Tue, 15 Jul 2025 15:22:22 -0700 Subject: [PATCH 01/10] let tracelimit define custom period and limit --- support/tracelimit/src/lib.rs | 129 +++++++++++++++++++++++++++++++++- 1 file changed, 127 insertions(+), 2 deletions(-) diff --git a/support/tracelimit/src/lib.rs b/support/tracelimit/src/lib.rs index 4a2aeaac77..5383af6104 100644 --- a/support/tracelimit/src/lib.rs +++ b/support/tracelimit/src/lib.rs @@ -23,8 +23,10 @@ use std::time::Instant; #[doc(hidden)] pub use tracing; -const PERIOD_MS: u32 = 5000; -const EVENTS_PER_PERIOD: u32 = 10; +#[doc(hidden)] +pub const PERIOD_MS: u32 = 5000; +#[doc(hidden)] +pub const EVENTS_PER_PERIOD: u32 = 10; static DISABLE_RATE_LIMITING: AtomicBool = AtomicBool::new(false); @@ -97,8 +99,49 @@ impl RateLimiter { } /// As [`tracing::error!`], but rate limited. +/// +/// Can be called with optional parameters to customize rate limiting: +/// - `period: ` - rate limiting period in milliseconds +/// - `limit: ` - maximum events per period +/// +/// Examples: +/// ``` +/// use tracelimit::error_ratelimited; +/// error_ratelimited!("simple message"); +/// error_ratelimited!(period: 1000, limit: 5, "custom rate limit"); +/// error_ratelimited!(period: 10000, "custom period only"); +/// error_ratelimited!(limit: 50, "custom limit only"); +/// ``` #[macro_export] macro_rules! error_ratelimited { + // With both period and limit + (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $limit); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // With period only + (period: $period:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $crate::EVENTS_PER_PERIOD); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // With limit only + (limit: $limit:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($crate::PERIOD_MS, $limit); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // Default case (no custom parameters) ($($rest:tt)*) => { { static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); @@ -110,8 +153,49 @@ macro_rules! error_ratelimited { } /// As [`tracing::warn!`], but rate limited. +/// +/// Can be called with optional parameters to customize rate limiting: +/// - `period: ` - rate limiting period in milliseconds +/// - `limit: ` - maximum events per period +/// +/// Examples: +/// ``` +/// use tracelimit::warn_ratelimited; +/// warn_ratelimited!("simple message"); +/// warn_ratelimited!(period: 1000, limit: 5, "custom rate limit"); +/// warn_ratelimited!(period: 10000, "custom period only"); +/// warn_ratelimited!(limit: 50, "custom limit only"); +/// ``` #[macro_export] macro_rules! warn_ratelimited { + // With both period and limit + (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $limit); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // With period only + (period: $period:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $crate::EVENTS_PER_PERIOD); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // With limit only + (limit: $limit:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($crate::PERIOD_MS, $limit); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // Default case (no custom parameters) ($($rest:tt)*) => { { static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); @@ -123,8 +207,49 @@ macro_rules! warn_ratelimited { } /// As [`tracing::info!`], but rate limited. +/// +/// Can be called with optional parameters to customize rate limiting: +/// - `period: ` - rate limiting period in milliseconds +/// - `limit: ` - maximum events per period +/// +/// Examples: +/// ``` +/// use tracelimit::info_ratelimited; +/// info_ratelimited!("simple message"); +/// info_ratelimited!(period: 1000, limit: 5, "custom rate limit"); +/// info_ratelimited!(period: 10000, "custom period only"); +/// info_ratelimited!(limit: 50, "custom limit only"); +/// ``` #[macro_export] macro_rules! info_ratelimited { + // With both period and limit + (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $limit); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // With period only + (period: $period:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $crate::EVENTS_PER_PERIOD); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // With limit only + (limit: $limit:expr, $($rest:tt)*) => { + { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($crate::PERIOD_MS, $limit); + if let Ok(missed_events) = RATE_LIMITER.event() { + $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); + } + } + }; + // Default case (no custom parameters) ($($rest:tt)*) => { { static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); From 6358f3ffe11ea2424e677badf1b3c5308e8c5b26 Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Tue, 15 Jul 2025 15:39:15 -0700 Subject: [PATCH 02/10] Use new tracelimit with limit --- .../firmware_uefi/src/service/diagnostics.rs | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs index 14b49f8a12..f29b51878b 100644 --- a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs +++ b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs @@ -32,6 +32,11 @@ use uefi_specs::hyperv::debug_level::DEBUG_FLAG_NAMES; use uefi_specs::hyperv::debug_level::DEBUG_WARN; use zerocopy::FromBytes; +/// Default amount of logs emitted per period +/// See `tracelimit::PERIOD_MS` and `tracelimit::EVENTS_PER_PERIOD` +/// for more details +const LOGS_PER_PERIOD: u32 = 150; + /// 8-byte alignment for every entry const ALIGNMENT: usize = 8; @@ -108,21 +113,24 @@ pub fn handle_efi_diagnostics_log<'a>(log: EfiDiagnosticsLog<'a>) { let phase_str = phase_to_string(log.phase); match log.debug_level { - DEBUG_WARN => tracing::warn!( + DEBUG_WARN => tracelimit::warn_ratelimited!( + limit: LOGS_PER_PERIOD, debug_level = %debug_level_str, ticks = log.ticks, phase = %phase_str, log_message = log.message, "EFI log entry" ), - DEBUG_ERROR => tracing::error!( + DEBUG_ERROR => tracelimit::error_ratelimited!( + limit: LOGS_PER_PERIOD, debug_level = %debug_level_str, ticks = log.ticks, phase = %phase_str, log_message = log.message, "EFI log entry" ), - _ => tracing::info!( + _ => tracelimit::info_ratelimited!( + limit: LOGS_PER_PERIOD, debug_level = %debug_level_str, ticks = log.ticks, phase = %phase_str, From f9a9c401a0afe7a680054f2e7b17143db50b43b8 Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Tue, 15 Jul 2025 16:13:36 -0700 Subject: [PATCH 03/10] Add runtime support for ratelimiting --- support/tracelimit/src/lib.rs | 101 ++++++++++++++---- vm/devices/firmware/firmware_uefi/src/lib.rs | 10 +- .../firmware_uefi/src/service/diagnostics.rs | 31 +++--- 3 files changed, 107 insertions(+), 35 deletions(-) diff --git a/support/tracelimit/src/lib.rs b/support/tracelimit/src/lib.rs index 5383af6104..d8acae606f 100644 --- a/support/tracelimit/src/lib.rs +++ b/support/tracelimit/src/lib.rs @@ -17,6 +17,8 @@ #![forbid(unsafe_code)] use parking_lot::Mutex; +use std::collections::HashMap; +use std::sync::LazyLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Instant; @@ -98,6 +100,42 @@ impl RateLimiter { } } +// Runtime rate limiter support +// +// NOTE: We can't use static rate limiters when users pass runtime variables to macros +// The solution is to use a global cache of rate limiters created on-demand +// +// HashMap: Cache rate limiters to preserve state across calls +// Mutex: Thread-safe access to the HashMap +// LazyLock: Initialize HashMap on first use +// +// The keys are (period, limit, call_site) to ensure that each macro location gets +// independent rate limiting. +static RUNTIME_RATE_LIMITERS: LazyLock>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// Get or create a runtime rate limiter for the given parameters and call site. +/// This will reuse existing rate limiters if they exist, otherwise create a new one. +#[doc(hidden)] +pub fn get_runtime_rate_limiter( + period_ms: u32, + events_per_period: u32, + call_site: &'static str, +) -> Result, RateLimited> { + if DISABLE_RATE_LIMITING.load(Ordering::Relaxed) { + return Ok(None); + } + + let key = (period_ms, events_per_period, call_site); + let mut limiters = RUNTIME_RATE_LIMITERS.lock(); + + let limiter = limiters + .entry(key) + .or_insert_with(|| RateLimiter::new(period_ms, events_per_period)); + + limiter.event() +} + /// As [`tracing::error!`], but rate limited. /// /// Can be called with optional parameters to customize rate limiting: @@ -117,8 +155,11 @@ macro_rules! error_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $limit); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $period, + $limit, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -126,8 +167,11 @@ macro_rules! error_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $crate::EVENTS_PER_PERIOD); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $period, + $crate::EVENTS_PER_PERIOD, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -135,8 +179,11 @@ macro_rules! error_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($crate::PERIOD_MS, $limit); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $crate::PERIOD_MS, + $limit, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -171,8 +218,11 @@ macro_rules! warn_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $limit); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $period, + $limit, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -180,8 +230,11 @@ macro_rules! warn_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $crate::EVENTS_PER_PERIOD); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $period, + $crate::EVENTS_PER_PERIOD, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -189,8 +242,11 @@ macro_rules! warn_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($crate::PERIOD_MS, $limit); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $crate::PERIOD_MS, + $limit, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -225,8 +281,11 @@ macro_rules! info_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $limit); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $period, + $limit, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -234,8 +293,11 @@ macro_rules! info_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($period, $crate::EVENTS_PER_PERIOD); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $period, + $crate::EVENTS_PER_PERIOD, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -243,8 +305,11 @@ macro_rules! info_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new($crate::PERIOD_MS, $limit); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter( + $crate::PERIOD_MS, + $limit, + concat!(file!(), ":", line!(), ":", column!()) + ) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } diff --git a/vm/devices/firmware/firmware_uefi/src/lib.rs b/vm/devices/firmware/firmware_uefi/src/lib.rs index b6257ebb83..9f4e306f6c 100644 --- a/vm/devices/firmware/firmware_uefi/src/lib.rs +++ b/vm/devices/firmware/firmware_uefi/src/lib.rs @@ -70,6 +70,8 @@ use local_clock::InspectableLocalClock; use pal_async::local::block_on; use platform::logger::UefiLogger; use platform::nvram::VsmConfig; +use service::diagnostics::DEFAULT_LOGS_PER_PERIOD; +use service::diagnostics::GUEST_LOGS_PER_PERIOD; use std::convert::TryInto; use std::ops::RangeInclusive; use std::task::Context; @@ -269,13 +271,15 @@ impl UefiDevice { tracelimit::info_ratelimited!(?addr, data, "set gpa for diagnostics"); self.service.diagnostics.set_gpa(data) } - UefiCommand::PROCESS_EFI_DIAGNOSTICS => self.process_diagnostics(false, "guest"), + UefiCommand::PROCESS_EFI_DIAGNOSTICS => { + self.process_diagnostics(false, GUEST_LOGS_PER_PERIOD, "guest") + } _ => tracelimit::warn_ratelimited!(addr, data, "unknown uefi write"), } } fn inspect_extra(&mut self, _resp: &mut inspect::Response<'_>) { - self.process_diagnostics(true, "inspect_extra"); + self.process_diagnostics(true, DEFAULT_LOGS_PER_PERIOD, "inspect_extra"); } } @@ -318,7 +322,7 @@ impl PollDevice for UefiDevice { // Poll watchdog timeout events if let Poll::Ready(Ok(())) = self.watchdog_recv.poll_recv(cx) { tracing::info!("watchdog timeout received"); - self.process_diagnostics(true, "watchdog timeout"); + self.process_diagnostics(true, DEFAULT_LOGS_PER_PERIOD, "watchdog timeout"); } } } diff --git a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs index f29b51878b..f361374ea4 100644 --- a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs +++ b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs @@ -32,10 +32,11 @@ use uefi_specs::hyperv::debug_level::DEBUG_FLAG_NAMES; use uefi_specs::hyperv::debug_level::DEBUG_WARN; use zerocopy::FromBytes; -/// Default amount of logs emitted per period -/// See `tracelimit::PERIOD_MS` and `tracelimit::EVENTS_PER_PERIOD` -/// for more details -const LOGS_PER_PERIOD: u32 = 150; +/// Default number of logs emitted per period +pub const DEFAULT_LOGS_PER_PERIOD: u32 = 500; + +/// Guest specific amount of logs emitted per period +pub const GUEST_LOGS_PER_PERIOD: u32 = 150; /// 8-byte alignment for every entry const ALIGNMENT: usize = 8; @@ -108,13 +109,13 @@ fn phase_to_string(phase: u16) -> &'static str { } /// Defines how we want EfiDiagnosticsLog entries to be handled. -pub fn handle_efi_diagnostics_log<'a>(log: EfiDiagnosticsLog<'a>) { +pub fn handle_efi_diagnostics_log<'a>(log: EfiDiagnosticsLog<'a>, limit: u32) { let debug_level_str = debug_level_to_string(log.debug_level); let phase_str = phase_to_string(log.phase); match log.debug_level { DEBUG_WARN => tracelimit::warn_ratelimited!( - limit: LOGS_PER_PERIOD, + limit: limit, debug_level = %debug_level_str, ticks = log.ticks, phase = %phase_str, @@ -122,7 +123,7 @@ pub fn handle_efi_diagnostics_log<'a>(log: EfiDiagnosticsLog<'a>) { "EFI log entry" ), DEBUG_ERROR => tracelimit::error_ratelimited!( - limit: LOGS_PER_PERIOD, + limit: limit, debug_level = %debug_level_str, ticks = log.ticks, phase = %phase_str, @@ -130,7 +131,7 @@ pub fn handle_efi_diagnostics_log<'a>(log: EfiDiagnosticsLog<'a>) { "EFI log entry" ), _ => tracelimit::info_ratelimited!( - limit: LOGS_PER_PERIOD, + limit: limit, debug_level = %debug_level_str, ticks = log.ticks, phase = %phase_str, @@ -470,12 +471,14 @@ impl UefiDevice { /// # Arguments /// * `allow_reprocess` - If true, allows processing even if already processed for guest /// * `context` - String to indicate who triggered the diagnostics processing - pub(crate) fn process_diagnostics(&mut self, allow_reprocess: bool, context: &str) { - if let Err(error) = self.service.diagnostics.process_diagnostics( - allow_reprocess, - &self.gm, - handle_efi_diagnostics_log, - ) { + pub(crate) fn process_diagnostics(&mut self, allow_reprocess: bool, limit: u32, context: &str) { + if let Err(error) = + self.service + .diagnostics + .process_diagnostics(allow_reprocess, &self.gm, |log| { + handle_efi_diagnostics_log(log, limit) + }) + { tracelimit::error_ratelimited!( error = &error as &dyn std::error::Error, context, From c9cd01012082dcc35a30ef1a40ecbaa73263ec1e Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Wed, 16 Jul 2025 10:44:17 -0700 Subject: [PATCH 04/10] Fix lifetime on handle_efi_diagnostics_log --- vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs index f361374ea4..3e30b5b22c 100644 --- a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs +++ b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs @@ -109,7 +109,7 @@ fn phase_to_string(phase: u16) -> &'static str { } /// Defines how we want EfiDiagnosticsLog entries to be handled. -pub fn handle_efi_diagnostics_log<'a>(log: EfiDiagnosticsLog<'a>, limit: u32) { +pub fn handle_efi_diagnostics_log(log: EfiDiagnosticsLog<'_>, limit: u32) { let debug_level_str = debug_level_to_string(log.debug_level); let phase_str = phase_to_string(log.phase); From bcb95e7e81cd439b0b68f3e1aaf41a565613a940 Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Wed, 16 Jul 2025 11:08:58 -0700 Subject: [PATCH 05/10] Remove hashmap in rate limiter --- support/tracelimit/src/lib.rs | 99 ++++++++++++----------------------- 1 file changed, 33 insertions(+), 66 deletions(-) diff --git a/support/tracelimit/src/lib.rs b/support/tracelimit/src/lib.rs index d8acae606f..5317075e7b 100644 --- a/support/tracelimit/src/lib.rs +++ b/support/tracelimit/src/lib.rs @@ -17,7 +17,6 @@ #![forbid(unsafe_code)] use parking_lot::Mutex; -use std::collections::HashMap; use std::sync::LazyLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; @@ -73,6 +72,20 @@ impl RateLimiter { } } + /// Reconfigure this rate limiter with new parameters. + /// Resets the rate limiting state if parameters changed. + pub fn reconfigure(&mut self, period_ms: u32, events_per_period: u32) { + if self.period_ms != period_ms || self.events_per_period != events_per_period { + self.period_ms = period_ms; + self.events_per_period = events_per_period; + // Reset state when parameters change + let state = self.state.get_mut(); + state.start = None; + state.events = 0; + state.missed = 0; + } + } + /// Returns `Ok(missed_events)` if this event should be logged. /// /// `missed_events` is `Some(n)` if there were any missed events or if this @@ -102,37 +115,27 @@ impl RateLimiter { // Runtime rate limiter support // -// NOTE: We can't use static rate limiters when users pass runtime variables to macros -// The solution is to use a global cache of rate limiters created on-demand +// NOTE: Rust does not allow using runtime values in static initializers, +// so we can't make the rate limiter itself static. // -// HashMap: Cache rate limiters to preserve state across calls -// Mutex: Thread-safe access to the HashMap -// LazyLock: Initialize HashMap on first use -// -// The keys are (period, limit, call_site) to ensure that each macro location gets -// independent rate limiting. -static RUNTIME_RATE_LIMITERS: LazyLock>> = - LazyLock::new(|| Mutex::new(HashMap::new())); +// Instead, we use a global `LazyLock` that holds a `Mutex`. +// This allows us to reconfigure the rate limiter at runtime while still +// providing a global instance that can be used in macros. +static RUNTIME_RATE_LIMITER: LazyLock> = + LazyLock::new(|| Mutex::new(RateLimiter::new_default())); -/// Get or create a runtime rate limiter for the given parameters and call site. -/// This will reuse existing rate limiters if they exist, otherwise create a new one. +/// Get the global runtime rate limiter, reconfiguring it if needed. #[doc(hidden)] pub fn get_runtime_rate_limiter( period_ms: u32, events_per_period: u32, - call_site: &'static str, ) -> Result, RateLimited> { if DISABLE_RATE_LIMITING.load(Ordering::Relaxed) { return Ok(None); } - let key = (period_ms, events_per_period, call_site); - let mut limiters = RUNTIME_RATE_LIMITERS.lock(); - - let limiter = limiters - .entry(key) - .or_insert_with(|| RateLimiter::new(period_ms, events_per_period)); - + let mut limiter = RUNTIME_RATE_LIMITER.lock(); + limiter.reconfigure(period_ms, events_per_period); limiter.event() } @@ -155,11 +158,7 @@ macro_rules! error_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $period, - $limit, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $limit) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -167,11 +166,7 @@ macro_rules! error_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $period, - $crate::EVENTS_PER_PERIOD, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $crate::EVENTS_PER_PERIOD) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -179,11 +174,7 @@ macro_rules! error_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $crate::PERIOD_MS, - $limit, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($crate::PERIOD_MS, $limit) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -218,11 +209,7 @@ macro_rules! warn_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $period, - $limit, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $limit) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -230,11 +217,7 @@ macro_rules! warn_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $period, - $crate::EVENTS_PER_PERIOD, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $crate::EVENTS_PER_PERIOD) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -242,11 +225,7 @@ macro_rules! warn_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $crate::PERIOD_MS, - $limit, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($crate::PERIOD_MS, $limit) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -281,11 +260,7 @@ macro_rules! info_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $period, - $limit, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $limit) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -293,11 +268,7 @@ macro_rules! info_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $period, - $crate::EVENTS_PER_PERIOD, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $crate::EVENTS_PER_PERIOD) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -305,11 +276,7 @@ macro_rules! info_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter( - $crate::PERIOD_MS, - $limit, - concat!(file!(), ":", line!(), ":", column!()) - ) { + if let Ok(missed_events) = $crate::get_runtime_rate_limiter($crate::PERIOD_MS, $limit) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } From 9aec5a185cdcd2c5a4e845f5303e00fad659ac7f Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Wed, 16 Jul 2025 13:50:03 -0700 Subject: [PATCH 06/10] Use inner mutex, change state in event_with_config(), additional feedback --- support/tracelimit/src/lib.rs | 122 +++++++++++++++++----------------- 1 file changed, 60 insertions(+), 62 deletions(-) diff --git a/support/tracelimit/src/lib.rs b/support/tracelimit/src/lib.rs index 5317075e7b..91c43fbeb1 100644 --- a/support/tracelimit/src/lib.rs +++ b/support/tracelimit/src/lib.rs @@ -17,7 +17,6 @@ #![forbid(unsafe_code)] use parking_lot::Mutex; -use std::sync::LazyLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Instant; @@ -25,9 +24,9 @@ use std::time::Instant; pub use tracing; #[doc(hidden)] -pub const PERIOD_MS: u32 = 5000; +const PERIOD_MS: u32 = 5000; #[doc(hidden)] -pub const EVENTS_PER_PERIOD: u32 = 10; +const EVENTS_PER_PERIOD: u32 = 10; static DISABLE_RATE_LIMITING: AtomicBool = AtomicBool::new(false); @@ -41,12 +40,12 @@ pub fn disable_rate_limiting(disabled: bool) { #[doc(hidden)] pub struct RateLimiter { - period_ms: u32, - events_per_period: u32, state: Mutex, } struct RateLimiterState { + period_ms: u32, + events_per_period: u32, start: Option, events: u32, missed: u64, @@ -62,9 +61,9 @@ impl RateLimiter { pub const fn new(period_ms: u32, events_per_period: u32) -> Self { Self { - period_ms, - events_per_period, state: Mutex::new(RateLimiterState { + period_ms, + events_per_period, start: None, events: 0, missed: 0, @@ -72,72 +71,74 @@ impl RateLimiter { } } - /// Reconfigure this rate limiter with new parameters. - /// Resets the rate limiting state if parameters changed. - pub fn reconfigure(&mut self, period_ms: u32, events_per_period: u32) { - if self.period_ms != period_ms || self.events_per_period != events_per_period { - self.period_ms = period_ms; - self.events_per_period = events_per_period; - // Reset state when parameters change - let state = self.state.get_mut(); - state.start = None; - state.events = 0; - state.missed = 0; - } + /// Returns `Ok(missed_events)` if this event should be logged. + /// + /// `missed_events` is `Some(n)` if there were any missed events or if this + /// event is the last one before rate limiting kicks in. + pub fn event(&self) -> Result, RateLimited> { + self.event_with_config(None, None) } /// Returns `Ok(missed_events)` if this event should be logged. + /// Optionally reconfigures the rate limiter if new parameters are provided. /// /// `missed_events` is `Some(n)` if there were any missed events or if this /// event is the last one before rate limiting kicks in. - pub fn event(&self) -> Result, RateLimited> { + pub fn event_with_config( + &self, + period_ms: Option, + events_per_period: Option, + ) -> Result, RateLimited> { if DISABLE_RATE_LIMITING.load(Ordering::Relaxed) { return Ok(None); } + let mut state = self.state.try_lock().ok_or(RateLimited)?; + + // Reconfigure if new parameters are provided + let mut reset_state = false; + if let Some(new_period) = period_ms { + if state.period_ms != new_period { + state.period_ms = new_period; + reset_state = true; + } + } + if let Some(new_events_per_period) = events_per_period { + if state.events_per_period != new_events_per_period { + state.events_per_period = new_events_per_period; + reset_state = true; + } + } + + // Reset state when parameters change + if reset_state { + state.start = None; + state.events = 0; + state.missed = 0; + } + let now = Instant::now(); + let period_ms = state.period_ms; let start = state.start.get_or_insert(now); let elapsed = now.duration_since(*start); - if elapsed.as_millis() > self.period_ms as u128 { + if elapsed.as_millis() > period_ms as u128 { *start = now; state.events = 0; } - if state.events >= self.events_per_period { + if state.events >= state.events_per_period { state.missed += 1; return Err(RateLimited); } state.events += 1; let missed = std::mem::take(&mut state.missed); - let missed = (missed != 0 || state.events == self.events_per_period).then_some(missed); + let missed = (missed != 0 || state.events == state.events_per_period).then_some(missed); Ok(missed) } } -// Runtime rate limiter support -// -// NOTE: Rust does not allow using runtime values in static initializers, -// so we can't make the rate limiter itself static. -// -// Instead, we use a global `LazyLock` that holds a `Mutex`. -// This allows us to reconfigure the rate limiter at runtime while still -// providing a global instance that can be used in macros. -static RUNTIME_RATE_LIMITER: LazyLock> = - LazyLock::new(|| Mutex::new(RateLimiter::new_default())); - -/// Get the global runtime rate limiter, reconfiguring it if needed. +// Global rate limiter instance #[doc(hidden)] -pub fn get_runtime_rate_limiter( - period_ms: u32, - events_per_period: u32, -) -> Result, RateLimited> { - if DISABLE_RATE_LIMITING.load(Ordering::Relaxed) { - return Ok(None); - } - - let mut limiter = RUNTIME_RATE_LIMITER.lock(); - limiter.reconfigure(period_ms, events_per_period); - limiter.event() -} +pub static RATE_LIMITER: RateLimiter = RateLimiter::new_default(); /// As [`tracing::error!`], but rate limited. /// @@ -158,7 +159,7 @@ macro_rules! error_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $limit) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), Some($limit)) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -166,7 +167,7 @@ macro_rules! error_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $crate::EVENTS_PER_PERIOD) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), None) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -174,7 +175,7 @@ macro_rules! error_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($crate::PERIOD_MS, $limit) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(None, Some($limit)) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -182,8 +183,7 @@ macro_rules! error_ratelimited { // Default case (no custom parameters) ($($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::RATE_LIMITER.event() { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -209,7 +209,7 @@ macro_rules! warn_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $limit) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), Some($limit)) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -217,7 +217,7 @@ macro_rules! warn_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $crate::EVENTS_PER_PERIOD) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), None) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -225,7 +225,7 @@ macro_rules! warn_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($crate::PERIOD_MS, $limit) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(None, Some($limit)) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -233,8 +233,7 @@ macro_rules! warn_ratelimited { // Default case (no custom parameters) ($($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::RATE_LIMITER.event() { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -260,7 +259,7 @@ macro_rules! info_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $limit) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), Some($limit)) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -268,7 +267,7 @@ macro_rules! info_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($period, $crate::EVENTS_PER_PERIOD) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), None) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -276,7 +275,7 @@ macro_rules! info_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::get_runtime_rate_limiter($crate::PERIOD_MS, $limit) { + if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(None, Some($limit)) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -284,8 +283,7 @@ macro_rules! info_ratelimited { // Default case (no custom parameters) ($($rest:tt)*) => { { - static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); - if let Ok(missed_events) = RATE_LIMITER.event() { + if let Ok(missed_events) = $crate::RATE_LIMITER.event() { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } From 2d26358317f7d87954402087dcdf7b06d82e9233 Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Wed, 16 Jul 2025 13:52:57 -0700 Subject: [PATCH 07/10] Remove unncessary doc hidden --- support/tracelimit/src/lib.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/support/tracelimit/src/lib.rs b/support/tracelimit/src/lib.rs index 91c43fbeb1..d6687d332a 100644 --- a/support/tracelimit/src/lib.rs +++ b/support/tracelimit/src/lib.rs @@ -23,9 +23,7 @@ use std::time::Instant; #[doc(hidden)] pub use tracing; -#[doc(hidden)] const PERIOD_MS: u32 = 5000; -#[doc(hidden)] const EVENTS_PER_PERIOD: u32 = 10; static DISABLE_RATE_LIMITING: AtomicBool = AtomicBool::new(false); From aab7e8edd82dfc8ae00c4ca0bf01279ddba8887e Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Wed, 16 Jul 2025 14:35:18 -0700 Subject: [PATCH 08/10] Remove shared rate limiter and create a static one for every callsite --- support/tracelimit/src/lib.rs | 40 +++++++++++++++++++++-------------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/support/tracelimit/src/lib.rs b/support/tracelimit/src/lib.rs index d6687d332a..0d3d0548c4 100644 --- a/support/tracelimit/src/lib.rs +++ b/support/tracelimit/src/lib.rs @@ -134,10 +134,6 @@ impl RateLimiter { } } -// Global rate limiter instance -#[doc(hidden)] -pub static RATE_LIMITER: RateLimiter = RateLimiter::new_default(); - /// As [`tracing::error!`], but rate limited. /// /// Can be called with optional parameters to customize rate limiting: @@ -157,7 +153,8 @@ macro_rules! error_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), Some($limit)) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), Some($limit)) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -165,7 +162,8 @@ macro_rules! error_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), None) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), None) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -173,7 +171,8 @@ macro_rules! error_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(None, Some($limit)) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(None, Some($limit)) { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -181,7 +180,8 @@ macro_rules! error_ratelimited { // Default case (no custom parameters) ($($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event() { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event() { $crate::tracing::error!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -207,7 +207,8 @@ macro_rules! warn_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), Some($limit)) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), Some($limit)) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -215,7 +216,8 @@ macro_rules! warn_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), None) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), None) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -223,7 +225,8 @@ macro_rules! warn_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(None, Some($limit)) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(None, Some($limit)) { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -231,7 +234,8 @@ macro_rules! warn_ratelimited { // Default case (no custom parameters) ($($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event() { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event() { $crate::tracing::warn!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -257,7 +261,8 @@ macro_rules! info_ratelimited { // With both period and limit (period: $period:expr, limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), Some($limit)) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), Some($limit)) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -265,7 +270,8 @@ macro_rules! info_ratelimited { // With period only (period: $period:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(Some($period), None) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(Some($period), None) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -273,7 +279,8 @@ macro_rules! info_ratelimited { // With limit only (limit: $limit:expr, $($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event_with_config(None, Some($limit)) { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event_with_config(None, Some($limit)) { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } @@ -281,7 +288,8 @@ macro_rules! info_ratelimited { // Default case (no custom parameters) ($($rest:tt)*) => { { - if let Ok(missed_events) = $crate::RATE_LIMITER.event() { + static RATE_LIMITER: $crate::RateLimiter = $crate::RateLimiter::new_default(); + if let Ok(missed_events) = RATE_LIMITER.event() { $crate::tracing::info!(dropped_ratelimited = missed_events, $($rest)*); } } From fdc52dacaed1a2f855436cf1136a5d2233bb1a0b Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Wed, 16 Jul 2025 15:03:28 -0700 Subject: [PATCH 09/10] Use only one events_per_period limit value, prevent reprocessing on watchdog timeout --- vm/devices/firmware/firmware_uefi/src/lib.rs | 11 +++++++---- .../firmware/firmware_uefi/src/service/diagnostics.rs | 7 ++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/vm/devices/firmware/firmware_uefi/src/lib.rs b/vm/devices/firmware/firmware_uefi/src/lib.rs index 9f4e306f6c..ce71b3139f 100644 --- a/vm/devices/firmware/firmware_uefi/src/lib.rs +++ b/vm/devices/firmware/firmware_uefi/src/lib.rs @@ -71,7 +71,6 @@ use pal_async::local::block_on; use platform::logger::UefiLogger; use platform::nvram::VsmConfig; use service::diagnostics::DEFAULT_LOGS_PER_PERIOD; -use service::diagnostics::GUEST_LOGS_PER_PERIOD; use std::convert::TryInto; use std::ops::RangeInclusive; use std::task::Context; @@ -272,13 +271,15 @@ impl UefiDevice { self.service.diagnostics.set_gpa(data) } UefiCommand::PROCESS_EFI_DIAGNOSTICS => { - self.process_diagnostics(false, GUEST_LOGS_PER_PERIOD, "guest") + self.process_diagnostics(false, DEFAULT_LOGS_PER_PERIOD, "guest") } _ => tracelimit::warn_ratelimited!(addr, data, "unknown uefi write"), } } fn inspect_extra(&mut self, _resp: &mut inspect::Response<'_>) { + // This is the only place where we allow reprocessing diagnostics, + // since this will be driven by the user themselves. self.process_diagnostics(true, DEFAULT_LOGS_PER_PERIOD, "inspect_extra"); } } @@ -321,8 +322,10 @@ impl PollDevice for UefiDevice { // Poll watchdog timeout events if let Poll::Ready(Ok(())) = self.watchdog_recv.poll_recv(cx) { - tracing::info!("watchdog timeout received"); - self.process_diagnostics(true, DEFAULT_LOGS_PER_PERIOD, "watchdog timeout"); + // NOTE: Do not allow reprocessing diagnostics here. + // UEFI programs the watchdog's configuration, so we should assume that + // this path could trigger multiple times. + self.process_diagnostics(false, DEFAULT_LOGS_PER_PERIOD, "watchdog timeout"); } } } diff --git a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs index 3e30b5b22c..42fecfd24b 100644 --- a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs +++ b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs @@ -32,11 +32,8 @@ use uefi_specs::hyperv::debug_level::DEBUG_FLAG_NAMES; use uefi_specs::hyperv::debug_level::DEBUG_WARN; use zerocopy::FromBytes; -/// Default number of logs emitted per period -pub const DEFAULT_LOGS_PER_PERIOD: u32 = 500; - -/// Guest specific amount of logs emitted per period -pub const GUEST_LOGS_PER_PERIOD: u32 = 150; +/// Default number of EfiDiagnosticsLogs emitted per period +pub const DEFAULT_LOGS_PER_PERIOD: u32 = 150; /// 8-byte alignment for every entry const ALIGNMENT: usize = 8; From 39acd254e6a39c5ee03570f2bdb66c04c9044c3b Mon Sep 17 00:00:00 2001 From: maheeraeron Date: Wed, 16 Jul 2025 16:02:22 -0700 Subject: [PATCH 10/10] Change context to triggered_by, output it on success --- .../firmware_uefi/src/service/diagnostics.rs | 33 ++++++++++++------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs index 42fecfd24b..bf9ddd5328 100644 --- a/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs +++ b/vm/devices/firmware/firmware_uefi/src/service/diagnostics.rs @@ -289,11 +289,13 @@ impl DiagnosticsServices { /// /// # Arguments /// * `allow_reprocess` - If true, allows processing even if already processed for guest + /// * `triggered_by` - String to indicate who triggered the diagnostics processing /// * `gm` - Guest memory to read diagnostics from /// * `log_handler` - Function to handle each parsed log entry fn process_diagnostics( &mut self, allow_reprocess: bool, + triggered_by: &str, gm: &GuestMemory, mut log_handler: F, ) -> Result<(), DiagnosticsError> @@ -456,7 +458,12 @@ impl DiagnosticsServices { } // Print summary statistics - tracelimit::info_ratelimited!(entries_processed, bytes_read, "processed EFI log entries"); + tracelimit::info_ratelimited!( + triggered_by, + entries_processed, + bytes_read, + "processed EFI log entries" + ); Ok(()) } @@ -467,18 +474,22 @@ impl UefiDevice { /// /// # Arguments /// * `allow_reprocess` - If true, allows processing even if already processed for guest - /// * `context` - String to indicate who triggered the diagnostics processing - pub(crate) fn process_diagnostics(&mut self, allow_reprocess: bool, limit: u32, context: &str) { - if let Err(error) = - self.service - .diagnostics - .process_diagnostics(allow_reprocess, &self.gm, |log| { - handle_efi_diagnostics_log(log, limit) - }) - { + /// * `triggered_by` - String to indicate who triggered the diagnostics processing + pub(crate) fn process_diagnostics( + &mut self, + allow_reprocess: bool, + limit: u32, + triggered_by: &str, + ) { + if let Err(error) = self.service.diagnostics.process_diagnostics( + allow_reprocess, + triggered_by, + &self.gm, + |log| handle_efi_diagnostics_log(log, limit), + ) { tracelimit::error_ratelimited!( error = &error as &dyn std::error::Error, - context, + triggered_by, "failed to process diagnostics buffer" ); }