move rate_limiter_set to its own module

This commit is contained in:
Chris Beck
2025-12-05 23:16:33 -07:00
parent 761877e35e
commit 5a12cd1bde
4 changed files with 83 additions and 74 deletions
+4 -1
View File
@@ -37,10 +37,13 @@ mod log_buffer;
mod log_handler;
use log_handler::{LogHandler, LogHandlerConfig};
mod rate_limiter_set;
pub use rate_limiter_set::{LimitResult, LimiterSet};
mod route;
pub use route::{Destination, Limit, Route};
pub use crate::rate_limiter::{LimitResult, Limiter, LimiterSet, RateThreshold};
pub use crate::rate_limiter::{Limiter, RateThreshold};
/// Configuration for the gateway.
#[derive(Conf, Debug)]
@@ -0,0 +1,76 @@
//! Rate limiter set for managing per-route rate limiting.
use crate::log_message::{LogMessage, Origin};
use crate::rate_limiter::Limiter;
use std::collections::HashMap;
/// Result of evaluating a limiter set.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LimitResult {
/// The event passed all limits (not rate-limited).
Passed,
/// The event was blocked by a per-origin limiter at the given index.
Limiter(usize),
/// The event was blocked by a global limiter at the given index.
GlobalLimiter(usize),
}
/// A set of limiters for a route, containing both per-origin and global limiters.
pub struct LimiterSet {
/// Factory to create limiters for new origins.
make_limiters: Box<dyn Fn() -> Vec<Limiter> + Send + Sync>,
/// Per-origin rate limiters, keyed by origin. Lazily created.
limiters: HashMap<Origin, Vec<Limiter>>,
/// Global rate limiters (shared across all origins).
global_limiters: Vec<Limiter>,
}
impl std::fmt::Debug for LimiterSet {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LimiterSet")
.field("limiters", &self.limiters)
.field("global_limiters", &self.global_limiters)
.finish_non_exhaustive()
}
}
impl LimiterSet {
/// Create a new limiter set with factories for per-origin and global limiters.
pub fn new(
make_limiters: impl Fn() -> Vec<Limiter> + Send + Sync + 'static,
global_limiters: Vec<Limiter>,
) -> Self {
Self {
make_limiters: Box::new(make_limiters),
limiters: HashMap::new(),
global_limiters,
}
}
/// Evaluate whether an event should pass all rate limits.
///
/// Returns [`LimitResult::Passed`] if the event passes all limits.
/// Returns [`LimitResult::Limiter(i)`] if blocked by per-origin limiter at index `i`.
/// Returns [`LimitResult::GlobalLimiter(i)`] if blocked by global limiter at index `i`.
pub fn evaluate(&mut self, log_msg: &LogMessage, origin: &Origin, ts_sec: i64) -> LimitResult {
// Get or create limiters for this origin
let origin_limiters = self
.limiters
.entry(origin.clone())
.or_insert_with(&self.make_limiters);
for (i, limiter) in origin_limiters.iter_mut().enumerate() {
if !limiter.evaluate(log_msg, ts_sec) {
return LimitResult::Limiter(i);
}
}
for (i, limiter) in self.global_limiters.iter_mut().enumerate() {
if !limiter.evaluate(log_msg, ts_sec) {
return LimitResult::GlobalLimiter(i);
}
}
LimitResult::Passed
}
}
+2 -1
View File
@@ -3,8 +3,9 @@
//! Routes define how log messages are processed based on filters, severity levels,
//! and destination overrides.
use super::rate_limiter_set::LimiterSet;
use crate::log_message::{Level, LogFilter};
use crate::rate_limiter::{Limiter, LimiterSet, RateThreshold};
use crate::rate_limiter::{Limiter, RateThreshold};
use serde::Deserialize;
/// A rate limit rule for suppressing repeated alerts.
+1 -72
View File
@@ -1,6 +1,6 @@
//! Rate limiting for log alerts.
use crate::log_message::{LogMessage, Origin};
use crate::log_message::LogMessage;
use serde::Deserialize;
use std::{
collections::HashMap,
@@ -87,77 +87,6 @@ pub enum Limiter {
SourceLocation(SourceLocationRateLimiter),
}
/// Result of evaluating a limiter set.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LimitResult {
/// The event passed all limits (not rate-limited).
Passed,
/// The event was blocked by a per-origin limiter at the given index.
Limiter(usize),
/// The event was blocked by a global limiter at the given index.
GlobalLimiter(usize),
}
/// A set of limiters for a route, containing both per-origin and global limiters.
pub struct LimiterSet {
/// Factory to create limiters for new origins.
make_limiters: Box<dyn Fn() -> Vec<Limiter> + Send + Sync>,
/// Per-origin rate limiters, keyed by origin. Lazily created.
limiters: HashMap<Origin, Vec<Limiter>>,
/// Global rate limiters (shared across all origins).
global_limiters: Vec<Limiter>,
}
impl std::fmt::Debug for LimiterSet {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LimiterSet")
.field("limiters", &self.limiters)
.field("global_limiters", &self.global_limiters)
.finish_non_exhaustive()
}
}
impl LimiterSet {
/// Create a new limiter set with factories for per-origin and global limiters.
pub fn new(
make_limiters: impl Fn() -> Vec<Limiter> + Send + Sync + 'static,
global_limiters: Vec<Limiter>,
) -> Self {
Self {
make_limiters: Box::new(make_limiters),
limiters: HashMap::new(),
global_limiters,
}
}
/// Evaluate whether an event should pass all rate limits.
///
/// Returns [`LimitResult::Passed`] if the event passes all limits.
/// Returns [`LimitResult::Limiter(i)`] if blocked by per-origin limiter at index `i`.
/// Returns [`LimitResult::GlobalLimiter(i)`] if blocked by global limiter at index `i`.
pub fn evaluate(&mut self, log_msg: &LogMessage, origin: &Origin, ts_sec: i64) -> LimitResult {
// Get or create limiters for this origin
let origin_limiters = self
.limiters
.entry(origin.clone())
.or_insert_with(&self.make_limiters);
for (i, limiter) in origin_limiters.iter_mut().enumerate() {
if !limiter.evaluate(log_msg, ts_sec) {
return LimitResult::Limiter(i);
}
}
for (i, limiter) in self.global_limiters.iter_mut().enumerate() {
if !limiter.evaluate(log_msg, ts_sec) {
return LimitResult::GlobalLimiter(i);
}
}
LimitResult::Passed
}
}
impl Limiter {
/// Evaluate whether an event should pass the rate limit.
///