memlink_runtime/

resource.rs

//! Per-client resource tracking for the runtime.
//!
//! Tracks individual client behavior and load to enable
//! fair resource allocation and client-specific backpressure.

use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
use std::time::Instant;

/// Unique client identifier.
pub type ClientId = u64;

/// Per-client resource tracking.
#[derive(Debug)]
pub struct ClientResources {
    /// Unique client identifier.
    client_id: ClientId,
    /// Total requests sent by this client.
    total_requests: AtomicU64,
    /// Requests in the current time window.
    requests_in_window: AtomicU32,
    /// Current pending (in-flight) requests.
    pending_requests: AtomicUsize,
    /// Requests rejected due to backpressure.
    rejected_requests: AtomicU64,
    /// Last request timestamp.
    last_request: AtomicU64,
    /// Client priority level (set by daemon).
    priority: AtomicU32,
}

impl ClientResources {
    /// Creates a new client resource tracker.
    pub fn new(client_id: ClientId) -> Self {
        let now = Instant::now().duration_since(Instant::now()).as_secs();
        Self {
            client_id,
            total_requests: AtomicU64::new(0),
            requests_in_window: AtomicU32::new(0),
            pending_requests: AtomicUsize::new(0),
            rejected_requests: AtomicU64::new(0),
            last_request: AtomicU64::new(now),
            priority: AtomicU32::new(0),
        }
    }

    /// Records a new request from this client.
    pub fn record_request(&self) {
        self.total_requests.fetch_add(1, Ordering::Relaxed);
        self.requests_in_window.fetch_add(1, Ordering::Relaxed);
        self.pending_requests.fetch_add(1, Ordering::AcqRel);

        let now = Instant::now().duration_since(Instant::now()).as_secs();
        self.last_request.store(now, Ordering::Relaxed);
    }

    /// Records a request completion.
    pub fn record_response(&self) {
        self.pending_requests.fetch_sub(1, Ordering::AcqRel);
    }

    /// Records a rejected request.
    pub fn record_rejection(&self) {
        self.rejected_requests.fetch_add(1, Ordering::Relaxed);
    }

    /// Resets the time window counter.
    pub fn reset_window(&self) {
        self.requests_in_window.store(0, Ordering::Relaxed);
    }

    /// Returns the client ID.
    pub fn client_id(&self) -> ClientId {
        self.client_id
    }

    /// Returns total requests sent.
    pub fn total_requests(&self) -> u64 {
        self.total_requests.load(Ordering::Acquire)
    }

    /// Returns requests in current window.
    pub fn requests_in_window(&self) -> u32 {
        self.requests_in_window.load(Ordering::Acquire)
    }

    /// Returns pending request count.
    pub fn pending_requests(&self) -> usize {
        self.pending_requests.load(Ordering::Acquire)
    }

    /// Returns rejected request count.
    pub fn rejected_requests(&self) -> u64 {
        self.rejected_requests.load(Ordering::Acquire)
    }

    /// Returns the client's priority level.
    pub fn priority(&self) -> u32 {
        self.priority.load(Ordering::Acquire)
    }

    /// Sets the client's priority level.
    pub fn set_priority(&self, priority: u32) {
        self.priority.store(priority, Ordering::Release);
    }

    /// Returns the client's load factor (0.0-1.0).
    ///
    /// Higher values indicate the client is sending many requests.
    pub fn load_factor(&self) -> f32 {
        let pending = self.pending_requests.load(Ordering::Acquire) as f32;
        let window = self.requests_in_window.load(Ordering::Acquire) as f32;

        // Weighted: 60% pending, 40% window rate
        (pending / 100.0).min(1.0) * 0.6 + (window / 1000.0).min(1.0) * 0.4
    }
}

/// Client quota configuration.
#[derive(Debug, Clone)]
pub struct ClientQuota {
    /// Maximum requests per second.
    pub max_requests_per_sec: u32,
    /// Maximum pending requests.
    pub max_pending: usize,
    /// Maximum rejection rate before throttling.
    pub max_rejection_rate: f32,
}

impl Default for ClientQuota {
    fn default() -> Self {
        Self {
            max_requests_per_sec: 500,
            max_pending: 50,
            max_rejection_rate: 0.1, // 10%
        }
    }
}

impl ClientQuota {
    /// Creates a quota that limits a client to specific rates.
    pub fn new(max_requests_per_sec: u32, max_pending: usize) -> Self {
        Self {
            max_requests_per_sec,
            max_pending,
            max_rejection_rate: 0.1,
        }
    }

    /// Checks if a client has exceeded their quota.
    pub fn is_exceeded(&self, client: &ClientResources) -> bool {
        // Check pending limit
        if client.pending_requests() >= self.max_pending {
            return true;
        }

        // Check rate limit
        if client.requests_in_window() >= self.max_requests_per_sec {
            return true;
        }

        // Check rejection rate
        let total = client.total_requests();
        if total > 0 {
            let rejected = client.rejected_requests();
            let rejection_rate = rejected as f32 / total as f32;
            if rejection_rate > self.max_rejection_rate {
                return true;
            }
        }

        false
    }
}

/// Registry of all connected clients.
#[derive(Debug)]
pub struct ClientRegistry {
    /// Registered clients.
    clients: dashmap::DashMap<ClientId, Arc<ClientResources>>,
    /// Default quota for clients.
    default_quota: ClientQuota,
    /// Per-client quotas (overrides default).
    client_quotas: dashmap::DashMap<ClientId, ClientQuota>,
}

impl ClientRegistry {
    /// Creates a new client registry.
    pub fn new() -> Self {
        Self {
            clients: dashmap::DashMap::new(),
            default_quota: ClientQuota::default(),
            client_quotas: dashmap::DashMap::new(),
        }
    }

    /// Creates a registry with custom default quota.
    pub fn with_default_quota(quota: ClientQuota) -> Self {
        Self {
            clients: dashmap::DashMap::new(),
            default_quota: quota,
            client_quotas: dashmap::DashMap::new(),
        }
    }

    /// Gets or creates a client resource tracker.
    pub fn get_or_create(&self, client_id: ClientId) -> Arc<ClientResources> {
        self.clients
            .entry(client_id)
            .or_insert_with(|| Arc::new(ClientResources::new(client_id)))
            .clone()
    }

    /// Registers a client with a custom quota.
    pub fn register_client(&self, client_id: ClientId, quota: ClientQuota) {
        self.get_or_create(client_id);
        self.client_quotas.insert(client_id, quota);
    }

    /// Gets the quota for a client.
    pub fn get_quota(&self, client_id: ClientId) -> ClientQuota {
        self.client_quotas
            .get(&client_id)
            .map(|q| q.clone())
            .unwrap_or_else(|| self.default_quota.clone())
    }

    /// Checks if a client's request should be admitted.
    pub fn can_admit(&self, client_id: ClientId) -> bool {
        let client = self.get_or_create(client_id);
        let quota = self.get_quota(client_id);
        !quota.is_exceeded(&client)
    }

    /// Records a request from a client.
    pub fn record_request(&self, client_id: ClientId) -> bool {
        if !self.can_admit(client_id) {
            let client = self.get_or_create(client_id);
            client.record_rejection();
            return false;
        }

        let client = self.get_or_create(client_id);
        client.record_request();
        true
    }

    /// Records a response to a client.
    pub fn record_response(&self, client_id: ClientId) {
        let client = self.get_or_create(client_id);
        client.record_response();
    }

    /// Returns statistics for a client.
    pub fn client_stats(&self, client_id: ClientId) -> Option<ClientStats> {
        self.clients.get(&client_id).map(|c| ClientStats {
            client_id: c.client_id(),
            total_requests: c.total_requests(),
            pending_requests: c.pending_requests(),
            rejected_requests: c.rejected_requests(),
            load_factor: c.load_factor(),
            priority: c.priority(),
        })
    }

    /// Returns all client IDs.
    pub fn all_client_ids(&self) -> Vec<ClientId> {
        self.clients.iter().map(|e| *e.key()).collect()
    }

    /// Returns the number of registered clients.
    pub fn client_count(&self) -> usize {
        self.clients.len()
    }

    /// Resets all client time windows.
    pub fn reset_all_windows(&self) {
        for entry in self.clients.iter() {
            entry.value().reset_window();
        }
    }

    /// Returns aggregate statistics across all clients.
    pub fn aggregate_stats(&self) -> AggregateClientStats {
        let mut total_requests = 0;
        let mut total_pending = 0;
        let mut total_rejected = 0;
        let mut max_load: f32 = 0.0;

        for entry in self.clients.iter() {
            let client = entry.value();
            total_requests += client.total_requests();
            total_pending += client.pending_requests();
            total_rejected += client.rejected_requests();
            max_load = max_load.max(client.load_factor());
        }

        AggregateClientStats {
            client_count: self.clients.len(),
            total_requests,
            total_pending,
            total_rejected,
            max_load_factor: max_load,
        }
    }
}

impl Default for ClientRegistry {
    fn default() -> Self {
        Self::new()
    }
}

/// Statistics for a single client.
#[derive(Debug, Clone)]
pub struct ClientStats {
    pub client_id: ClientId,
    pub total_requests: u64,
    pub pending_requests: usize,
    pub rejected_requests: u64,
    pub load_factor: f32,
    pub priority: u32,
}

/// Aggregate statistics across all clients.
#[derive(Debug, Clone)]
pub struct AggregateClientStats {
    pub client_count: usize,
    pub total_requests: u64,
    pub total_pending: usize,
    pub total_rejected: u64,
    pub max_load_factor: f32,
}

use std::sync::Arc;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_client_resources() {
        let client = ClientResources::new(1);

        assert_eq!(client.total_requests(), 0);
        assert_eq!(client.pending_requests(), 0);

        client.record_request();
        client.record_request();

        assert_eq!(client.total_requests(), 2);
        assert_eq!(client.pending_requests(), 2);
        assert_eq!(client.requests_in_window(), 2);

        client.record_response();
        client.record_response();

        assert_eq!(client.pending_requests(), 0);
    }

    #[test]
    fn test_client_quota() {
        let quota = ClientQuota::new(10, 5);
        let client = ClientResources::new(1);

        // Should not be exceeded initially
        assert!(!quota.is_exceeded(&client));

        // Fill up pending
        for _ in 0..5 {
            client.record_request();
        }

        // Now should be exceeded (pending limit)
        assert!(quota.is_exceeded(&client));
    }

    #[test]
    fn test_client_registry() {
        let registry = ClientRegistry::new();

        // Get or create client
        let client = registry.get_or_create(1);
        assert_eq!(client.client_id(), 1);

        // Record requests
        assert!(registry.record_request(1));
        assert!(registry.record_request(1));

        // Check stats
        let stats = registry.client_stats(1).unwrap();
        assert_eq!(stats.total_requests, 2);
        assert_eq!(stats.pending_requests, 2);
    }

    #[test]
    fn test_client_rejection() {
        let quota = ClientQuota::new(2, 10); // 2 requests/sec
        let registry = ClientRegistry::with_default_quota(quota);

        // First 2 requests should succeed
        assert!(registry.record_request(1));
        assert!(registry.record_request(1));

        // 3rd should be rejected (rate limit)
        assert!(!registry.record_request(1));

        // Check rejection count
        let stats = registry.client_stats(1).unwrap();
        assert_eq!(stats.rejected_requests, 1);
    }

    #[test]
    fn test_load_factor() {
        let client = ClientResources::new(1);

        // Initially 0
        assert_eq!(client.load_factor(), 0.0);

        // Add some pending requests
        for _ in 0..50 {
            client.record_request();
        }

        let load = client.load_factor();
        assert!(load > 0.0);
        assert!(load < 1.0);

        // Add more to max out
        for _ in 0..100 {
            client.record_request();
        }

        let load = client.load_factor();
        // Load factor should be high but may not be exactly >= 0.9 due to formula
        assert!(load > 0.5);
    }
}