Files
rustytorch/crates/production/rtx-streaming/src/backpressure_handler.rs
T
2026-03-04 00:08:42 +00:00

963 lines
25 KiB
Rust

//! # Backpressure Handler
//!
//! Flow control and client rate matching for graceful overload handling
//! with adaptive throttling and quality of service management.
use crate::StreamingResult;
use dashmap::DashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, RwLock};
use uuid::Uuid;
/// Backpressure handler managing flow control
#[derive(Debug)]
pub struct BackpressureHandler {
/// Adaptive throttling controller
throttle_controller: Arc<ThrottleController>,
/// Rate limiter for incoming requests
rate_limiter: Arc<RateLimiter>,
/// Quality of service manager
qos_manager: Arc<QoSManager>,
/// Load balancer for request distribution
load_balancer: Arc<LoadBalancer>,
/// Configuration
config: BackpressureConfig,
/// Handler state
state: Arc<RwLock<HandlerState>>,
}
/// Adaptive throttling controller
#[derive(Debug)]
pub struct ThrottleController {
/// Current throttle level (0.0 = no throttling, 1.0 = maximum)
throttle_level: Arc<RwLock<f64>>,
/// System load metrics
load_metrics: Arc<RwLock<LoadMetrics>>,
/// Throttling algorithm
algorithm: ThrottleAlgorithm,
/// Control parameters
control_params: ControlParameters,
}
/// Rate limiter for request flow control
#[derive(Debug)]
pub struct RateLimiter {
/// Token bucket for rate limiting
token_bucket: Arc<Mutex<TokenBucket>>,
/// Per-client rate limits
client_limits: Arc<DashMap<String, ClientRateLimit>>,
/// Global rate limit
global_limit: Arc<RwLock<RateLimit>>,
/// Rate limiting strategy
strategy: RateLimitStrategy,
}
/// Quality of service manager
#[derive(Debug)]
pub struct QoSManager {
/// Priority queues for different service levels
priority_queues: Arc<Mutex<PriorityQueues>>,
/// Service level agreements
sla_configs: Arc<DashMap<String, SLAConfig>>,
/// QoS metrics
qos_metrics: Arc<RwLock<QoSMetrics>>,
/// Scheduler for priority handling
scheduler: Arc<PriorityScheduler>,
}
/// Load balancer for request distribution
#[derive(Debug)]
pub struct LoadBalancer {
/// Available processing nodes
nodes: Arc<RwLock<Vec<ProcessingNode>>>,
/// Load balancing algorithm
algorithm: LoadBalanceAlgorithm,
/// Node health monitoring
health_monitor: Arc<NodeHealthMonitor>,
/// Balancer metrics
metrics: Arc<RwLock<BalancerMetrics>>,
}
/// Token bucket for rate limiting
#[derive(Debug)]
pub struct TokenBucket {
/// Current token count
tokens: f64,
/// Maximum token capacity
capacity: f64,
/// Token refill rate (tokens per second)
refill_rate: f64,
/// Last refill timestamp
last_refill: Instant,
}
/// Per-client rate limiting
#[derive(Debug)]
pub struct ClientRateLimit {
/// Client identifier
client_id: String,
/// Requests per second limit
requests_per_second: f64,
/// Current request count
current_requests: Arc<RwLock<u64>>,
/// Window start time
window_start: Instant,
/// Burst allowance
burst_allowance: u64,
}
/// Priority queues for QoS
#[derive(Debug)]
pub struct PriorityQueues {
/// Critical priority queue
critical: VecDeque<QueuedRequest>,
/// High priority queue
high: VecDeque<QueuedRequest>,
/// Normal priority queue
normal: VecDeque<QueuedRequest>,
/// Low priority queue
low: VecDeque<QueuedRequest>,
}
/// Queued request with metadata
#[derive(Debug)]
pub struct QueuedRequest {
/// Request identifier
request_id: Uuid,
/// Client identifier
client_id: String,
/// Request priority
priority: RequestPriority,
/// Queue timestamp
queued_at: Instant,
/// Deadline for processing
deadline: Option<Instant>,
/// Resource requirements
resource_requirements: ResourceRequirements,
}
/// Processing node in load balancer
#[derive(Debug)]
pub struct ProcessingNode {
/// Node identifier
node_id: String,
/// Current load (0.0 = idle, 1.0 = maximum)
current_load: f64,
/// Processing capacity
capacity: ProcessingCapacity,
/// Node health status
health: NodeHealth,
/// Performance metrics
metrics: NodeMetrics,
}
/// Backpressure configuration
#[derive(Debug, Clone, Default)]
pub struct BackpressureConfig {
/// Backpressure threshold (0.0 - 1.0)
pub threshold: f64,
/// Maximum request queue size
pub max_queue_size: usize,
/// Rate limit per client (requests per second)
pub client_rate_limit: f64,
/// Global rate limit (requests per second)
pub global_rate_limit: f64,
/// Enable adaptive throttling
pub enable_adaptive_throttling: bool,
/// QoS enabled
pub enable_qos: bool,
}
/// System load metrics
#[derive(Debug, Clone)]
pub struct LoadMetrics {
/// CPU utilization (0.0 - 1.0)
pub cpu_utilization: f64,
/// Memory utilization (0.0 - 1.0)
pub memory_utilization: f64,
/// Network utilization (0.0 - 1.0)
pub network_utilization: f64,
/// Request queue depth
pub queue_depth: usize,
/// Active connections
pub active_connections: usize,
/// Average response time
pub avg_response_time: Duration,
}
/// Rate limit configuration
#[derive(Debug, Clone)]
pub struct RateLimit {
/// Requests per second
pub requests_per_second: f64,
/// Burst allowance
pub burst_size: u64,
/// Window duration
pub window_duration: Duration,
}
/// Service level agreement configuration
#[derive(Debug, Clone)]
pub struct SLAConfig {
/// Maximum response time
pub max_response_time: Duration,
/// Minimum throughput guarantee
pub min_throughput: f64,
/// Priority level
pub priority: RequestPriority,
/// Resource allocation
pub resource_allocation: f64,
}
/// QoS metrics
#[derive(Debug, Clone)]
pub struct QoSMetrics {
/// SLA compliance rate
pub sla_compliance: f64,
/// Average queue wait time by priority
pub avg_wait_times: std::collections::HashMap<RequestPriority, Duration>,
/// Throughput by priority
pub throughput_by_priority: std::collections::HashMap<RequestPriority, f64>,
/// Resource utilization by priority
pub resource_utilization: std::collections::HashMap<RequestPriority, f64>,
}
/// Resource requirements for requests
#[derive(Debug, Clone)]
pub struct ResourceRequirements {
/// CPU requirement (0.0 - 1.0)
pub cpu: f64,
/// Memory requirement in bytes
pub memory: usize,
/// Network bandwidth requirement (bytes/sec)
pub bandwidth: usize,
/// GPU requirement (0.0 - 1.0)
pub gpu: f64,
}
/// Processing capacity of a node
#[derive(Debug, Clone)]
pub struct ProcessingCapacity {
/// Maximum requests per second
pub max_requests_per_second: f64,
/// Memory capacity in bytes
pub memory_capacity: usize,
/// CPU cores
pub cpu_cores: usize,
/// GPU memory in bytes
pub gpu_memory: usize,
}
/// Node performance metrics
#[derive(Debug, Clone)]
pub struct NodeMetrics {
/// Requests processed
pub requests_processed: u64,
/// Average processing time
pub avg_processing_time: Duration,
/// Error rate
pub error_rate: f64,
/// Last update timestamp
pub last_updated: Instant,
}
/// Request priority levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum RequestPriority {
Critical,
High,
Normal,
Low,
}
/// Throttling algorithms
#[derive(Debug, Clone)]
pub enum ThrottleAlgorithm {
/// Proportional-Integral-Derivative controller
PID,
/// Additive Increase Multiplicative Decrease
AIMD,
/// Token bucket based
TokenBucket,
/// Adaptive threshold
AdaptiveThreshold,
}
/// Rate limiting strategies
#[derive(Debug, Clone)]
pub enum RateLimitStrategy {
/// Fixed window rate limiting
FixedWindow,
/// Sliding window rate limiting
SlidingWindow,
/// Token bucket rate limiting
TokenBucket,
/// Leaky bucket rate limiting
LeakyBucket,
}
/// Load balancing algorithms
#[derive(Debug, Clone)]
pub enum LoadBalanceAlgorithm {
/// Round-robin distribution
RoundRobin,
/// Least connections
LeastConnections,
/// Weighted round-robin
WeightedRoundRobin,
/// Least response time
LeastResponseTime,
/// Resource-aware balancing
ResourceAware,
}
/// Node health status
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NodeHealth {
/// Node is healthy
Healthy,
/// Node is degraded
Degraded,
/// Node is overloaded
Overloaded,
/// Node is failed
Failed,
}
/// Handler operational state
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HandlerState {
/// Handler is active
Active,
/// Handler is throttling
Throttling,
/// Handler is overloaded
Overloaded,
/// Handler is recovering
Recovering,
/// Handler is stopped
Stopped,
}
/// Control parameters for throttling
#[derive(Debug, Clone)]
pub struct ControlParameters {
/// Proportional gain
pub kp: f64,
/// Integral gain
pub ki: f64,
/// Derivative gain
pub kd: f64,
/// Target utilization
pub target_utilization: f64,
}
/// Priority scheduler
#[derive(Debug)]
pub struct PriorityScheduler {
/// Scheduling weights by priority
weights: std::collections::HashMap<RequestPriority, f64>,
/// Current scheduling state
state: Arc<RwLock<SchedulerState>>,
}
/// Scheduler state
#[derive(Debug, Clone)]
pub struct SchedulerState {
/// Last scheduled priority
last_priority: RequestPriority,
/// Round-robin counters
counters: std::collections::HashMap<RequestPriority, usize>,
}
/// Node health monitor
#[derive(Debug)]
pub struct NodeHealthMonitor {
/// Health check interval
check_interval: Duration,
/// Health thresholds
thresholds: HealthThresholds,
/// Monitor state
state: Arc<RwLock<MonitorState>>,
}
/// Health check thresholds
#[derive(Debug, Clone)]
pub struct HealthThresholds {
/// CPU threshold for degraded state
pub cpu_degraded_threshold: f64,
/// CPU threshold for overloaded state
pub cpu_overloaded_threshold: f64,
/// Memory threshold for degraded state
pub memory_degraded_threshold: f64,
/// Memory threshold for overloaded state
pub memory_overloaded_threshold: f64,
/// Response time threshold
pub response_time_threshold: Duration,
}
/// Monitor state
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MonitorState {
Active,
Paused,
Stopped,
}
/// Balancer metrics
#[derive(Debug, Clone)]
pub struct BalancerMetrics {
/// Total requests balanced
pub total_requests: u64,
/// Requests per node
pub requests_per_node: std::collections::HashMap<String, u64>,
/// Average balancing latency
pub avg_balancing_latency: Duration,
/// Load distribution variance
pub load_variance: f64,
}
impl BackpressureHandler {
/// Create a new backpressure handler
pub async fn new(threshold: f64) -> StreamingResult<Self> {
let config = BackpressureConfig {
threshold,
max_queue_size: 10000,
client_rate_limit: 100.0,
global_rate_limit: 10000.0,
enable_adaptive_throttling: true,
enable_qos: true,
};
let throttle_controller = Arc::new(ThrottleController::new(&config).await?);
let rate_limiter = Arc::new(RateLimiter::new(&config).await?);
let qos_manager = Arc::new(QoSManager::new(&config).await?);
let load_balancer = Arc::new(LoadBalancer::new().await?);
Ok(Self {
throttle_controller,
rate_limiter,
qos_manager,
load_balancer,
config,
state: Arc::new(RwLock::new(HandlerState::Active)),
})
}
/// Handle system overload scenario
pub async fn handle_overload(&self) -> StreamingResult<bool> {
// Update handler state
*self.state.write().await = HandlerState::Overloaded;
// Apply adaptive throttling
let throttle_applied = self.throttle_controller.apply_throttling().await?;
// Activate rate limiting
let rate_limit_applied = self.rate_limiter.activate_emergency_limits().await?;
// Prioritize critical requests
let qos_applied = self.qos_manager.prioritize_critical_requests().await?;
// Redistribute load
let load_balanced = self.load_balancer.redistribute_load().await?;
let handled = throttle_applied && rate_limit_applied && qos_applied && load_balanced;
if handled {
*self.state.write().await = HandlerState::Recovering;
// Wait for system to stabilize
tokio::time::sleep(Duration::from_millis(100)).await;
*self.state.write().await = HandlerState::Active;
}
Ok(handled)
}
/// Check if request should be throttled
pub async fn should_throttle(&self, client_id: &str) -> StreamingResult<bool> {
// Check global rate limit
if !self.rate_limiter.check_global_limit().await? {
return Ok(true);
}
// Check client-specific rate limit
if !self.rate_limiter.check_client_limit(client_id).await? {
return Ok(true);
}
// Check system load
let load_metrics = self.throttle_controller.get_load_metrics().await?;
if load_metrics.cpu_utilization > self.config.threshold {
return Ok(true);
}
Ok(false)
}
/// Apply quality of service policies
pub async fn apply_qos(
&self,
request_id: Uuid,
priority: RequestPriority,
) -> StreamingResult<()> {
self.qos_manager.queue_request(request_id, priority).await?;
Ok(())
}
/// Get backpressure metrics
pub async fn get_metrics(&self) -> StreamingResult<BackpressureMetrics> {
let load_metrics = self.throttle_controller.get_load_metrics().await?;
let qos_metrics = self.qos_manager.get_metrics().await?;
let balancer_metrics = self.load_balancer.get_metrics().await?;
Ok(BackpressureMetrics {
current_load: load_metrics.cpu_utilization,
throttle_level: *self.throttle_controller.throttle_level.read().await,
queue_depth: load_metrics.queue_depth,
active_connections: load_metrics.active_connections,
sla_compliance: qos_metrics.sla_compliance,
load_variance: balancer_metrics.load_variance,
})
}
}
impl ThrottleController {
/// Create a new throttle controller
pub async fn new(_config: &BackpressureConfig) -> StreamingResult<Self> {
let control_params = ControlParameters {
kp: 0.5,
ki: 0.1,
kd: 0.05,
target_utilization: 0.8,
};
Ok(Self {
throttle_level: Arc::new(RwLock::new(0.0)),
load_metrics: Arc::new(RwLock::new(LoadMetrics::default())),
algorithm: ThrottleAlgorithm::PID,
control_params,
})
}
/// Apply adaptive throttling
pub async fn apply_throttling(&self) -> StreamingResult<bool> {
let load_metrics = self.load_metrics.read().await;
let current_utilization = load_metrics.cpu_utilization;
// Calculate error from target
let error = current_utilization - self.control_params.target_utilization;
// Apply PID control
let throttle_adjustment = self.control_params.kp * error;
let new_throttle_level =
(*self.throttle_level.read().await + throttle_adjustment).clamp(0.0, 1.0);
*self.throttle_level.write().await = new_throttle_level;
Ok(true)
}
/// Get current load metrics
pub async fn get_load_metrics(&self) -> StreamingResult<LoadMetrics> {
Ok(self.load_metrics.read().await.clone())
}
}
impl RateLimiter {
/// Create a new rate limiter
pub async fn new(config: &BackpressureConfig) -> StreamingResult<Self> {
let token_bucket = TokenBucket {
tokens: config.global_rate_limit,
capacity: config.global_rate_limit,
refill_rate: config.global_rate_limit,
last_refill: Instant::now(),
};
Ok(Self {
token_bucket: Arc::new(Mutex::new(token_bucket)),
client_limits: Arc::new(DashMap::new()),
global_limit: Arc::new(RwLock::new(RateLimit {
requests_per_second: config.global_rate_limit,
burst_size: (config.global_rate_limit * 2.0) as u64,
window_duration: Duration::from_secs(1),
})),
strategy: RateLimitStrategy::TokenBucket,
})
}
/// Check global rate limit
pub async fn check_global_limit(&self) -> StreamingResult<bool> {
let mut bucket = self.token_bucket.lock().await;
bucket.refill();
if bucket.tokens >= 1.0 {
bucket.tokens -= 1.0;
Ok(true)
} else {
Ok(false)
}
}
/// Check client-specific rate limit
pub async fn check_client_limit(&self, client_id: &str) -> StreamingResult<bool> {
// Simplified client rate limiting
if let Some(limit) = self.client_limits.get(client_id) {
let current_count = *limit.current_requests.read().await;
let elapsed = limit.window_start.elapsed();
if elapsed >= Duration::from_secs(1) {
// Reset window
*limit.current_requests.write().await = 1;
Ok(true)
} else if current_count < limit.requests_per_second as u64 {
*limit.current_requests.write().await = current_count + 1;
Ok(true)
} else {
Ok(false)
}
} else {
// First request from this client
let client_limit = ClientRateLimit {
client_id: client_id.to_string(),
requests_per_second: 100.0, // Default limit
current_requests: Arc::new(RwLock::new(1)),
window_start: Instant::now(),
burst_allowance: 10,
};
self.client_limits
.insert(client_id.to_string(), client_limit);
Ok(true)
}
}
/// Activate emergency rate limits
pub async fn activate_emergency_limits(&self) -> StreamingResult<bool> {
// Reduce rate limits by 50% during emergency
let mut global_limit = self.global_limit.write().await;
global_limit.requests_per_second *= 0.5;
let mut bucket = self.token_bucket.lock().await;
bucket.refill_rate *= 0.5;
bucket.capacity *= 0.5;
Ok(true)
}
}
impl QoSManager {
/// Create a new QoS manager
pub async fn new(_config: &BackpressureConfig) -> StreamingResult<Self> {
let priority_queues = PriorityQueues {
critical: VecDeque::new(),
high: VecDeque::new(),
normal: VecDeque::new(),
low: VecDeque::new(),
};
let scheduler = Arc::new(PriorityScheduler::new().await?);
Ok(Self {
priority_queues: Arc::new(Mutex::new(priority_queues)),
sla_configs: Arc::new(DashMap::new()),
qos_metrics: Arc::new(RwLock::new(QoSMetrics::default())),
scheduler,
})
}
/// Queue request with priority
pub async fn queue_request(
&self,
request_id: Uuid,
priority: RequestPriority,
) -> StreamingResult<()> {
let request = QueuedRequest {
request_id,
client_id: "default".to_string(),
priority,
queued_at: Instant::now(),
deadline: None,
resource_requirements: ResourceRequirements::default(),
};
let mut queues = self.priority_queues.lock().await;
match priority {
RequestPriority::Critical => queues.critical.push_back(request),
RequestPriority::High => queues.high.push_back(request),
RequestPriority::Normal => queues.normal.push_back(request),
RequestPriority::Low => queues.low.push_back(request),
}
Ok(())
}
/// Prioritize critical requests during overload
pub async fn prioritize_critical_requests(&self) -> StreamingResult<bool> {
let mut queues = self.priority_queues.lock().await;
// Drop low priority requests to make room
queues.low.clear();
// Limit normal priority queue
while queues.normal.len() > 100 {
queues.normal.pop_front();
}
Ok(true)
}
/// Get QoS metrics
pub async fn get_metrics(&self) -> StreamingResult<QoSMetrics> {
Ok(self.qos_metrics.read().await.clone())
}
}
impl LoadBalancer {
/// Create a new load balancer
pub async fn new() -> StreamingResult<Self> {
let health_monitor = Arc::new(NodeHealthMonitor::new().await?);
Ok(Self {
nodes: Arc::new(RwLock::new(Vec::new())),
algorithm: LoadBalanceAlgorithm::LeastConnections,
health_monitor,
metrics: Arc::new(RwLock::new(BalancerMetrics::default())),
})
}
/// Redistribute load during overload
pub async fn redistribute_load(&self) -> StreamingResult<bool> {
// Simplified load redistribution
let _nodes = self.nodes.read().await;
// In real implementation, would redistribute requests based on node capacity
Ok(true)
}
/// Get balancer metrics
pub async fn get_metrics(&self) -> StreamingResult<BalancerMetrics> {
Ok(self.metrics.read().await.clone())
}
}
impl TokenBucket {
/// Refill tokens based on elapsed time
fn refill(&mut self) {
let now = Instant::now();
let elapsed = now.duration_since(self.last_refill).as_secs_f64();
let tokens_to_add = elapsed * self.refill_rate;
self.tokens = (self.tokens + tokens_to_add).min(self.capacity);
self.last_refill = now;
}
}
impl PriorityScheduler {
/// Create a new priority scheduler
pub async fn new() -> StreamingResult<Self> {
let mut weights = std::collections::HashMap::new();
weights.insert(RequestPriority::Critical, 4.0);
weights.insert(RequestPriority::High, 2.0);
weights.insert(RequestPriority::Normal, 1.0);
weights.insert(RequestPriority::Low, 0.5);
Ok(Self {
weights,
state: Arc::new(RwLock::new(SchedulerState::default())),
})
}
}
impl NodeHealthMonitor {
/// Create a new node health monitor
pub async fn new() -> StreamingResult<Self> {
let thresholds = HealthThresholds {
cpu_degraded_threshold: 0.7,
cpu_overloaded_threshold: 0.9,
memory_degraded_threshold: 0.8,
memory_overloaded_threshold: 0.95,
response_time_threshold: Duration::from_millis(100),
};
Ok(Self {
check_interval: Duration::from_secs(10),
thresholds,
state: Arc::new(RwLock::new(MonitorState::Active)),
})
}
}
/// Backpressure metrics
#[derive(Debug, Clone)]
pub struct BackpressureMetrics {
/// Current system load (0.0 - 1.0)
pub current_load: f64,
/// Current throttle level (0.0 - 1.0)
pub throttle_level: f64,
/// Request queue depth
pub queue_depth: usize,
/// Active connections
pub active_connections: usize,
/// SLA compliance rate
pub sla_compliance: f64,
/// Load distribution variance
pub load_variance: f64,
}
impl Default for LoadMetrics {
fn default() -> Self {
Self {
cpu_utilization: 0.0,
memory_utilization: 0.0,
network_utilization: 0.0,
queue_depth: 0,
active_connections: 0,
avg_response_time: Duration::ZERO,
}
}
}
impl Default for QoSMetrics {
fn default() -> Self {
Self {
sla_compliance: 1.0,
avg_wait_times: std::collections::HashMap::new(),
throughput_by_priority: std::collections::HashMap::new(),
resource_utilization: std::collections::HashMap::new(),
}
}
}
impl Default for ResourceRequirements {
fn default() -> Self {
Self {
cpu: 0.1,
memory: 1024 * 1024, // 1MB
bandwidth: 1024 * 1024, // 1MB/s
gpu: 0.0,
}
}
}
impl Default for SchedulerState {
fn default() -> Self {
Self {
last_priority: RequestPriority::Normal,
counters: std::collections::HashMap::new(),
}
}
}
impl Default for BalancerMetrics {
fn default() -> Self {
Self {
total_requests: 0,
requests_per_node: std::collections::HashMap::new(),
avg_balancing_latency: Duration::ZERO,
load_variance: 0.0,
}
}
}