veridian_kernel/services/lb/
l7.rs

1//! L7 (Application Layer) Load Balancer
2//!
3//! Provides HTTP-aware load balancing with path/host-based routing,
4//! rate limiting, and sticky sessions.
5
6#![allow(dead_code)]
7
8use alloc::{collections::BTreeMap, string::String, vec::Vec};
9
10use super::l4::{Backend, LbAlgorithm};
11
12// ---------------------------------------------------------------------------
13// HTTP Route
14// ---------------------------------------------------------------------------
15
16/// An HTTP routing rule.
17#[derive(Debug, Clone)]
18pub struct HttpRoute {
19    /// Path prefix to match (e.g., "/api/v1").
20    pub path_prefix: String,
21    /// Host header to match (empty = any).
22    pub host: String,
23    /// Required headers (all must match).
24    pub headers: BTreeMap<String, String>,
25    /// Name of the backend group to route to.
26    pub backend_group: String,
27}
28
29impl HttpRoute {
30    /// Check if a request matches this route.
31    pub fn matches(&self, path: &str, host: &str, headers: &BTreeMap<String, String>) -> bool {
32        // Check path prefix
33        if !path.starts_with(&self.path_prefix) {
34            return false;
35        }
36        // Check host (empty = any)
37        if !self.host.is_empty() && self.host != host {
38            return false;
39        }
40        // Check required headers
41        for (key, value) in &self.headers {
42            match headers.get(key) {
43                Some(v) if v == value => {}
44                _ => return false,
45            }
46        }
47        true
48    }
49}
50
51// ---------------------------------------------------------------------------
52// Backend Group
53// ---------------------------------------------------------------------------
54
55/// A named group of backends.
56#[derive(Debug, Clone)]
57pub struct BackendGroup {
58    /// Group name.
59    pub name: String,
60    /// Backends in this group.
61    pub backends: Vec<Backend>,
62    /// Load balancing algorithm for this group.
63    pub algorithm: LbAlgorithm,
64    /// Round-robin counter.
65    rr_counter: u64,
66}
67
68impl BackendGroup {
69    /// Create a new backend group.
70    pub fn new(name: String, algorithm: LbAlgorithm) -> Self {
71        BackendGroup {
72            name,
73            backends: Vec::new(),
74            algorithm,
75            rr_counter: 0,
76        }
77    }
78
79    /// Select a healthy backend.
80    pub fn select(&mut self) -> Option<(String, u16)> {
81        let healthy: Vec<usize> = self
82            .backends
83            .iter()
84            .enumerate()
85            .filter(|(_, b)| b.healthy)
86            .map(|(i, _)| i)
87            .collect();
88
89        if healthy.is_empty() {
90            return None;
91        }
92
93        let idx = match self.algorithm {
94            LbAlgorithm::RoundRobin | LbAlgorithm::WeightedRoundRobin => {
95                let i = (self.rr_counter as usize) % healthy.len();
96                self.rr_counter += 1;
97                healthy[i]
98            }
99            LbAlgorithm::LeastConnections => {
100                let mut min_idx = healthy[0];
101                let mut min_conn = self.backends[healthy[0]].active_connections;
102                for &h in &healthy[1..] {
103                    if self.backends[h].active_connections < min_conn {
104                        min_conn = self.backends[h].active_connections;
105                        min_idx = h;
106                    }
107                }
108                min_idx
109            }
110            _ => {
111                let i = (self.rr_counter as usize) % healthy.len();
112                self.rr_counter += 1;
113                healthy[i]
114            }
115        };
116
117        self.backends[idx].active_connections += 1;
118        self.backends[idx].total_requests += 1;
119        Some((self.backends[idx].address.clone(), self.backends[idx].port))
120    }
121}
122
123// ---------------------------------------------------------------------------
124// L7 Rule
125// ---------------------------------------------------------------------------
126
127/// A set of L7 routing rules.
128#[derive(Debug, Clone)]
129pub struct L7Rule {
130    /// Ordered routes (first match wins).
131    pub routes: Vec<HttpRoute>,
132    /// Default backend group (if no route matches).
133    pub default_backend: String,
134}
135
136// ---------------------------------------------------------------------------
137// Rate Limiting
138// ---------------------------------------------------------------------------
139
140/// Token bucket rate limiter.
141#[derive(Debug, Clone)]
142pub struct RateLimit {
143    /// Maximum requests per second (integer).
144    pub requests_per_second: u32,
145    /// Burst size.
146    pub burst: u32,
147    /// Current available tokens.
148    pub current_tokens: u32,
149    /// Tick when tokens were last refilled.
150    pub last_refill_tick: u64,
151}
152
153impl RateLimit {
154    /// Create a new rate limiter.
155    pub fn new(requests_per_second: u32, burst: u32) -> Self {
156        RateLimit {
157            requests_per_second,
158            burst,
159            current_tokens: burst,
160            last_refill_tick: 0,
161        }
162    }
163
164    /// Try to consume a token. Returns true if allowed.
165    pub fn allow(&mut self, current_tick: u64) -> bool {
166        self.refill(current_tick);
167        if self.current_tokens > 0 {
168            self.current_tokens -= 1;
169            true
170        } else {
171            false
172        }
173    }
174
175    /// Refill tokens based on elapsed time.
176    fn refill(&mut self, current_tick: u64) {
177        if current_tick <= self.last_refill_tick {
178            return;
179        }
180        let elapsed = current_tick - self.last_refill_tick;
181        let new_tokens = elapsed.saturating_mul(self.requests_per_second as u64);
182        // Cap at burst, but clamp to u32 first
183        let clamped = if new_tokens > self.burst as u64 {
184            self.burst
185        } else {
186            new_tokens as u32
187        };
188        self.current_tokens = self.current_tokens.saturating_add(clamped).min(self.burst);
189        self.last_refill_tick = current_tick;
190    }
191}
192
193// ---------------------------------------------------------------------------
194// Sticky Sessions
195// ---------------------------------------------------------------------------
196
197/// Sticky session manager (cookie-based affinity).
198#[derive(Debug, Clone)]
199pub struct StickySession {
200    /// Cookie name used for session affinity.
201    pub cookie_name: String,
202    /// Session ID -> backend index mapping.
203    pub backend_map: BTreeMap<String, usize>,
204}
205
206impl StickySession {
207    /// Create a new sticky session manager.
208    pub fn new(cookie_name: String) -> Self {
209        StickySession {
210            cookie_name,
211            backend_map: BTreeMap::new(),
212        }
213    }
214
215    /// Get the backend index for a session.
216    pub fn get_backend(&self, session_id: &str) -> Option<usize> {
217        self.backend_map.get(session_id).copied()
218    }
219
220    /// Set the backend for a session.
221    pub fn set_backend(&mut self, session_id: String, backend_idx: usize) {
222        self.backend_map.insert(session_id, backend_idx);
223    }
224
225    /// Remove a session.
226    pub fn remove_session(&mut self, session_id: &str) {
227        self.backend_map.remove(session_id);
228    }
229
230    /// Get the number of active sessions.
231    pub fn session_count(&self) -> usize {
232        self.backend_map.len()
233    }
234}
235
236// ---------------------------------------------------------------------------
237// L7 Error
238// ---------------------------------------------------------------------------
239
240/// L7 load balancer error.
241#[derive(Debug, Clone, PartialEq, Eq)]
242pub enum L7Error {
243    /// No route matched.
244    NoRouteMatch,
245    /// Backend group not found.
246    BackendGroupNotFound(String),
247    /// No healthy backend.
248    NoHealthyBackend,
249    /// Rate limit exceeded.
250    RateLimitExceeded,
251}
252
253// ---------------------------------------------------------------------------
254// L7 Load Balancer
255// ---------------------------------------------------------------------------
256
257/// L7 Load Balancer implementation.
258#[derive(Debug)]
259pub struct L7LoadBalancer {
260    /// Routing rules.
261    rules: Vec<L7Rule>,
262    /// Backend groups.
263    backend_groups: BTreeMap<String, BackendGroup>,
264    /// Rate limiters keyed by identifier (e.g., client IP).
265    rate_limiters: BTreeMap<String, RateLimit>,
266    /// Sticky session manager.
267    sticky_sessions: StickySession,
268    /// Default rate limit config.
269    default_rps: u32,
270    /// Default burst.
271    default_burst: u32,
272}
273
274impl Default for L7LoadBalancer {
275    fn default() -> Self {
276        Self::new()
277    }
278}
279
280impl L7LoadBalancer {
281    /// Create a new L7 load balancer.
282    pub fn new() -> Self {
283        L7LoadBalancer {
284            rules: Vec::new(),
285            backend_groups: BTreeMap::new(),
286            rate_limiters: BTreeMap::new(),
287            sticky_sessions: StickySession::new(String::from("VERIDIAN_SESSION")),
288            default_rps: 100,
289            default_burst: 200,
290        }
291    }
292
293    /// Add a routing rule.
294    pub fn add_rule(&mut self, rule: L7Rule) {
295        self.rules.push(rule);
296    }
297
298    /// Add a backend group.
299    pub fn add_backend_group(&mut self, group: BackendGroup) {
300        self.backend_groups.insert(group.name.clone(), group);
301    }
302
303    /// Route an HTTP request to a backend.
304    pub fn route_request(
305        &mut self,
306        path: &str,
307        host: &str,
308        headers: &BTreeMap<String, String>,
309    ) -> Result<(String, u16), L7Error> {
310        // Find matching route
311        let mut target_group = None;
312
313        for rule in &self.rules {
314            for route in &rule.routes {
315                if route.matches(path, host, headers) {
316                    target_group = Some(route.backend_group.clone());
317                    break;
318                }
319            }
320            if target_group.is_some() {
321                break;
322            }
323            // Use default backend from rule if no route matched
324            if target_group.is_none() && !rule.default_backend.is_empty() {
325                target_group = Some(rule.default_backend.clone());
326            }
327        }
328
329        let group_name = target_group.ok_or(L7Error::NoRouteMatch)?;
330
331        let group = self
332            .backend_groups
333            .get_mut(&group_name)
334            .ok_or_else(|| L7Error::BackendGroupNotFound(group_name.clone()))?;
335
336        group.select().ok_or(L7Error::NoHealthyBackend)
337    }
338
339    /// Check rate limit for a client.
340    pub fn check_rate_limit(&mut self, client_id: &str, current_tick: u64) -> Result<(), L7Error> {
341        let limiter = self
342            .rate_limiters
343            .entry(String::from(client_id))
344            .or_insert_with(|| RateLimit::new(self.default_rps, self.default_burst));
345
346        if limiter.allow(current_tick) {
347            Ok(())
348        } else {
349            Err(L7Error::RateLimitExceeded)
350        }
351    }
352
353    /// Get sticky backend for a session.
354    pub fn get_sticky_backend(&self, session_id: &str) -> Option<usize> {
355        self.sticky_sessions.get_backend(session_id)
356    }
357
358    /// Set sticky backend for a session.
359    pub fn set_sticky_backend(&mut self, session_id: String, backend_idx: usize) {
360        self.sticky_sessions.set_backend(session_id, backend_idx);
361    }
362
363    /// Get the number of backend groups.
364    pub fn backend_group_count(&self) -> usize {
365        self.backend_groups.len()
366    }
367
368    /// Get the number of routing rules.
369    pub fn rule_count(&self) -> usize {
370        self.rules.len()
371    }
372}
373
374// ---------------------------------------------------------------------------
375// Tests
376// ---------------------------------------------------------------------------
377
378#[cfg(test)]
379mod tests {
380    #[allow(unused_imports)]
381    use alloc::string::ToString;
382    #[allow(unused_imports)]
383    use alloc::vec;
384
385    use super::*;
386
387    fn make_lb() -> L7LoadBalancer {
388        let mut lb = L7LoadBalancer::new();
389
390        let mut group = BackendGroup::new(String::from("api-backends"), LbAlgorithm::RoundRobin);
391        group
392            .backends
393            .push(Backend::new(String::from("10.0.0.1"), 8080, 1));
394        group
395            .backends
396            .push(Backend::new(String::from("10.0.0.2"), 8080, 1));
397        lb.add_backend_group(group);
398
399        let mut default_group =
400            BackendGroup::new(String::from("default-backends"), LbAlgorithm::RoundRobin);
401        default_group
402            .backends
403            .push(Backend::new(String::from("10.0.1.1"), 80, 1));
404        lb.add_backend_group(default_group);
405
406        lb.add_rule(L7Rule {
407            routes: vec![HttpRoute {
408                path_prefix: String::from("/api/"),
409                host: String::new(),
410                headers: BTreeMap::new(),
411                backend_group: String::from("api-backends"),
412            }],
413            default_backend: String::from("default-backends"),
414        });
415
416        lb
417    }
418
419    #[test]
420    fn test_route_by_path() {
421        let mut lb = make_lb();
422        let headers = BTreeMap::new();
423        let (addr, port) = lb
424            .route_request("/api/v1/pods", "example.com", &headers)
425            .unwrap();
426        assert_eq!(port, 8080);
427        assert!(addr.starts_with("10.0.0."));
428    }
429
430    #[test]
431    fn test_route_default_backend() {
432        let mut lb = make_lb();
433        let headers = BTreeMap::new();
434        let (addr, port) = lb
435            .route_request("/other/page", "example.com", &headers)
436            .unwrap();
437        assert_eq!(addr, "10.0.1.1");
438        assert_eq!(port, 80);
439    }
440
441    #[test]
442    fn test_route_with_host_match() {
443        let mut lb = L7LoadBalancer::new();
444        let mut group = BackendGroup::new(String::from("host-group"), LbAlgorithm::RoundRobin);
445        group
446            .backends
447            .push(Backend::new(String::from("10.0.0.5"), 443, 1));
448        lb.add_backend_group(group);
449
450        lb.add_rule(L7Rule {
451            routes: vec![HttpRoute {
452                path_prefix: String::from("/"),
453                host: String::from("api.example.com"),
454                headers: BTreeMap::new(),
455                backend_group: String::from("host-group"),
456            }],
457            default_backend: String::new(),
458        });
459
460        let headers = BTreeMap::new();
461        assert!(lb
462            .route_request("/", "other.example.com", &headers)
463            .is_err());
464        let result = lb.route_request("/", "api.example.com", &headers);
465        assert!(result.is_ok());
466    }
467
468    #[test]
469    fn test_rate_limit_allow() {
470        let mut lb = make_lb();
471        assert!(lb.check_rate_limit("client-1", 100).is_ok());
472    }
473
474    #[test]
475    fn test_rate_limit_exceeded() {
476        let mut lb = L7LoadBalancer::new();
477        lb.default_rps = 1;
478        lb.default_burst = 1;
479
480        lb.check_rate_limit("client-1", 100).unwrap();
481        assert_eq!(
482            lb.check_rate_limit("client-1", 100),
483            Err(L7Error::RateLimitExceeded)
484        );
485    }
486
487    #[test]
488    fn test_rate_limit_refill() {
489        let mut lb = L7LoadBalancer::new();
490        lb.default_rps = 1;
491        lb.default_burst = 2;
492
493        lb.check_rate_limit("c1", 100).unwrap();
494        lb.check_rate_limit("c1", 100).unwrap();
495        assert!(lb.check_rate_limit("c1", 100).is_err());
496        // After 1 tick, should get 1 more token
497        assert!(lb.check_rate_limit("c1", 101).is_ok());
498    }
499
500    #[test]
501    fn test_sticky_session() {
502        let mut lb = make_lb();
503        lb.set_sticky_backend(String::from("sess-123"), 0);
504        assert_eq!(lb.get_sticky_backend("sess-123"), Some(0));
505        assert_eq!(lb.get_sticky_backend("sess-999"), None);
506    }
507
508    #[test]
509    fn test_http_route_matching() {
510        let route = HttpRoute {
511            path_prefix: String::from("/api/"),
512            host: String::from("example.com"),
513            headers: BTreeMap::new(),
514            backend_group: String::from("test"),
515        };
516        let empty_headers = BTreeMap::new();
517        assert!(route.matches("/api/v1", "example.com", &empty_headers));
518        assert!(!route.matches("/web/", "example.com", &empty_headers));
519        assert!(!route.matches("/api/v1", "other.com", &empty_headers));
520    }
521
522    #[test]
523    fn test_http_route_header_matching() {
524        let mut required = BTreeMap::new();
525        required.insert(String::from("x-version"), String::from("v2"));
526        let route = HttpRoute {
527            path_prefix: String::from("/"),
528            host: String::new(),
529            headers: required,
530            backend_group: String::from("test"),
531        };
532
533        let mut headers = BTreeMap::new();
534        headers.insert(String::from("x-version"), String::from("v2"));
535        assert!(route.matches("/", "", &headers));
536
537        let empty_headers = BTreeMap::new();
538        assert!(!route.matches("/", "", &empty_headers));
539    }
540
541    #[test]
542    fn test_backend_group_select() {
543        let mut group = BackendGroup::new(String::from("test"), LbAlgorithm::RoundRobin);
544        group
545            .backends
546            .push(Backend::new(String::from("10.0.0.1"), 80, 1));
547        let result = group.select();
548        assert!(result.is_some());
549    }
550
551    #[test]
552    fn test_no_route_match() {
553        let mut lb = L7LoadBalancer::new();
554        let headers = BTreeMap::new();
555        assert_eq!(
556            lb.route_request("/", "", &headers),
557            Err(L7Error::NoRouteMatch)
558        );
559    }
560}
veridian_kernel/services/lb/l7.rs

veridian_kernel/services/lb/
l7.rs
