azoth-balancer 0.3.1

//! This module contains the endpoint selection strategies for the load balancer.
//!
//! The primary goal is to decouple the logic of *how* an endpoint is chosen
//! from the main `balancer` module, which is responsible for managing state.
//! This makes the selection algorithm swappable and easier to test in isolation.

use parking_lot::{Mutex, RwLockReadGuard};
use ratelimit_meter::DirectRateLimiter;
use std::{cmp::Ordering, collections::HashMap, num::NonZeroU32, sync::Arc};
use tracing::info;

use crate::{
    endpoint::{LoadBalancerError, RpcEndpoint},
    metrics::{
        ALL_ENDPOINTS_RATE_LIMITED, ENDPOINT_RATE_LIMIT_DEFERRED, PRIORITY_ENDPOINT_SELECTED,
    },
};

// --- Tier Weight Definitions ---
const TIER1_MIN_WEIGHT: u32 = 100;
const TIER2_MIN_WEIGHT: u32 = 50;
const TIER3_MIN_WEIGHT: u32 = 1;

/// Selects the best available endpoint based on a tiered-priority strategy.
///
/// Note: This implementation uses LRU (least recently used) sorting across all tiers to ensure fair load distribution.
/// Tier 1 ignores total_cost to prioritize high-weight (local) endpoints, unlike Tier 2, which considers cost for premium endpoints.
///
/// The selection process is strictly ordered by tiers, with fallback:
/// 1.  **Tier 1 (Local, weight >= 100):** Tries to find an available endpoint first.
///     - Sorted by `weight` (desc), then `last_selected` (asc - LRU).
///     - If all Tier 1 endpoints are unavailable or rate-limited, it falls back to Tier 2.
///
/// 2.  **Tier 2 (Premium, weight 50-99):** Tries if no suitable Tier 1 endpoint is found.
///     - Sorted by `weight` (desc), then `total_cost` (asc), then `last_selected` (asc - LRU).
///     - Falls back to Tier 3 if no suitable endpoint is found.
///
/// 3.  **Tier 3 (Free, weight 1-49):** The final fallback tier.
///     - Sorted by `weight` (desc), then `last_selected` (asc - LRU).
///     - Ignores cost for fair round-robin distribution.
pub fn select_best_endpoint<'a>(
    endpoints: &'a RwLockReadGuard<Vec<RpcEndpoint>>,
    rate_limiters: &RwLockReadGuard<HashMap<String, Arc<Mutex<DirectRateLimiter>>>>,
    batch_size: usize,
) -> Option<&'a RpcEndpoint> {
    let available_endpoints: Vec<&RpcEndpoint> =
        endpoints.iter().filter(|ep| ep.is_available() && ep.weight > 0).collect();

    if available_endpoints.is_empty() {
        return None;
    }

    // Collect endpoints by tier
    let mut tier1 = Vec::new();
    let mut tier2 = Vec::new();
    let mut tier3 = Vec::new();

    for ep in available_endpoints {
        match ep.weight {
            w if w >= TIER1_MIN_WEIGHT => tier1.push(ep),
            w if w >= TIER2_MIN_WEIGHT => tier2.push(ep),
            w if w >= TIER3_MIN_WEIGHT => tier3.push(ep),
            _ => {} // weight == 0, already filtered out
        }
    }

    let tiers = [&tier1, &tier2, &tier3];
    let use_cost_flags = [false, true, false]; // Tier 1 ignores cost, Tier 2 uses cost, Tier 3 ignores cost

    for (tier_index, (candidates, &use_cost)) in tiers.iter().zip(use_cost_flags.iter()).enumerate()
    {
        if candidates.is_empty() {
            continue;
        }

        let tier = (tier_index + 1) as u8;
        info!("Attempting selection from Tier {} ({} candidates)", tier, candidates.len());

        let mut sorted_candidates = (*candidates).clone();
        sorted_candidates.sort_by(|&a, &b| {
            b.weight
                .cmp(&a.weight)
                .then_with(|| {
                    if use_cost {
                        a.metrics.get_total_cost().cmp(&b.metrics.get_total_cost())
                    } else {
                        Ordering::Equal
                    }
                })
                .then_with(|| {
                    let time_a = a.metrics.last_selected.load(std::sync::atomic::Ordering::Relaxed);
                    let time_b = b.metrics.last_selected.load(std::sync::atomic::Ordering::Relaxed);
                    time_a.cmp(&time_b) // LRU for all tiers
                })
        });

        for (rank, &endpoint) in sorted_candidates.iter().enumerate() {
            if reserve_tokens(rate_limiters, &endpoint.url, batch_size).is_ok() {
                log_selection(endpoint, tier, rank);
                return Some(endpoint);
            }
            log_rate_limited(endpoint, tier, rank);
        }
    }

    ALL_ENDPOINTS_RATE_LIMITED.inc();
    None
}

/// A helper function to check and reserve tokens from the rate limiter for a given endpoint.
fn reserve_tokens(
    limiters: &RwLockReadGuard<HashMap<String, Arc<Mutex<DirectRateLimiter>>>>,
    endpoint_url: &str,
    batch_size: usize,
) -> Result<(), LoadBalancerError> {
    if batch_size > (u32::MAX as usize) {
        return Err(LoadBalancerError::UpstreamError("Batch size too large".to_string()));
    }

    if let Some(limiter) = limiters.get(endpoint_url) {
        if batch_size == 0 {
            return Ok(());
        }

        if let Some(n) = NonZeroU32::new(batch_size as u32) {
            if limiter.lock().check_n(n.get()).is_err() {
                return Err(LoadBalancerError::RateLimited(endpoint_url.to_string()));
            }
        }
    }
    Ok(())
}

fn log_selection(endpoint: &RpcEndpoint, tier: u8, rank: usize) {
    let priority_rank = (rank + 1).to_string();
    PRIORITY_ENDPOINT_SELECTED.with_label_values(&[&endpoint.name, &priority_rank]).inc();
    info!(
        endpoint = %endpoint.name,
        tier = tier,
        priority_rank = rank + 1,
        "Selected endpoint"
    );
}

fn log_rate_limited(endpoint: &RpcEndpoint, tier: u8, rank: usize) {
    ENDPOINT_RATE_LIMIT_DEFERRED.with_label_values(&[&endpoint.name]).inc();
    info!(
        endpoint = %endpoint.name,
        tier = tier,
        priority_rank = rank + 1,
        "Endpoint is rate-limited, trying next"
    );
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::endpoint::{EndpointMetrics, RpcEndpoint};
    use parking_lot::RwLock;
    use ratelimit_meter::DirectRateLimiter;
    use std::sync::atomic::{AtomicU32, AtomicU64};
    use std::time::{Duration, Instant};

    // Helper to create a test endpoint
    fn create_test_endpoint(
        url: &str,
        weight: u32,
        cost: u64,
        healthy: bool,
        cooldown_secs: Option<u64>,
        last_selected: u64,
    ) -> RpcEndpoint {
        RpcEndpoint {
            name: url.to_string(), // Use URL as name for tests
            url: url.to_string(),
            weight,
            metrics: Arc::new(EndpointMetrics {
                ema_latency_ms: AtomicU64::new(cost),
                ema_error_penalty_ms: AtomicU64::new(0),
                consecutive_failures: AtomicU32::new(0),
                last_selected: AtomicU64::new(last_selected),
            }),
            healthy,
            cooldown_until: cooldown_secs.map(|secs| Instant::now() + Duration::from_secs(secs)),
            last_check: Instant::now(),
            cooldown_attempts: 0,
            rate_limit_per_sec: 100,
            burst_size: 100,
        }
    }

    // Helper to create a rate limiter
    fn create_rate_limiter(burst: u32) -> Arc<Mutex<DirectRateLimiter>> {
        let capacity = NonZeroU32::new(burst).unwrap();
        Arc::new(Mutex::new(DirectRateLimiter::new(capacity, Duration::from_secs(1))))
    }

    fn run_selection<'a>(
        endpoints_lock: &'a RwLockReadGuard<Vec<RpcEndpoint>>,
        limiters_lock: &RwLockReadGuard<HashMap<String, Arc<Mutex<DirectRateLimiter>>>>,
        batch_size: usize,
    ) -> Option<&'a RpcEndpoint> {
        select_best_endpoint(endpoints_lock, limiters_lock, batch_size)
    }

    #[test]
    fn test_tier1_is_always_prioritized() {
        let endpoints = vec![
            create_test_endpoint("tier3.com", 10, 10, true, None, 0),
            create_test_endpoint("tier2.com", 60, 20, true, None, 0),
            create_test_endpoint("tier1.com", 100, 500, true, None, 0), // High cost, but Tier 1
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier1.com");
    }

    #[test]
    fn test_tier1_sorts_by_weight_then_lru_ignores_cost() {
        let endpoints = vec![
            create_test_endpoint("tier1-a.com", 110, 1200, true, None, 20), // Higher weight, worse cost, newer
            create_test_endpoint("tier1-b.com", 110, 1100, true, None, 10), // Higher weight, better cost, older
            create_test_endpoint("tier1-c.com", 100, 100, true, None, 5),   // Lower weight
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier1-b.com"); // Same weight, LRU wins (cost ignored)
    }

    #[test]
    fn test_tier1_lru_tiebreaker() {
        let endpoints = vec![
            create_test_endpoint("tier1-a.com", 110, 100, true, None, 20), // Same weight/cost, newer
            create_test_endpoint("tier1-b.com", 110, 100, true, None, 10), // Same weight/cost, older
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier1-b.com"); // LRU - older timestamp wins
    }

    #[test]
    fn test_tier1_rate_limited_falls_back_to_tier2() {
        let endpoints = vec![
            create_test_endpoint("tier1-rate-limited.com", 100, 50, true, None, 0),
            create_test_endpoint("tier2-available.com", 60, 20, true, None, 0),
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();

        let mut limiters = HashMap::new();
        let limiter1 = create_rate_limiter(1);
        limiter1.lock().check_n(1).unwrap(); // Use up all tokens
        limiters.insert("tier1-rate-limited.com".to_string(), limiter1);
        limiters.insert("tier2-available.com".to_string(), create_rate_limiter(1));

        let limiters_rwlock = RwLock::new(limiters);
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier2-available.com");
    }

    #[test]
    fn test_tier2_sorts_by_weight_then_cost_then_lru() {
        let endpoints = vec![
            create_test_endpoint("tier2-a.com", 60, 1200, true, None, 10),
            create_test_endpoint("tier2-b.com", 70, 1500, true, None, 5), // Higher weight, worse cost, older
            create_test_endpoint("tier2-c.com", 70, 1100, true, None, 15), // Higher weight, better cost, newer
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier2-c.com"); // Higher weight, then better cost (LRU not reached)
    }

    #[test]
    fn test_tier2_lru_tiebreaker() {
        let endpoints = vec![
            create_test_endpoint("tier2-a.com", 70, 1000, true, None, 20), // Same weight/cost, newer
            create_test_endpoint("tier2-b.com", 70, 1000, true, None, 10), // Same weight/cost, older
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier2-b.com"); // LRU - older timestamp wins
    }

    #[test]
    fn test_tier2_unavailable_falls_back_to_tier3() {
        let endpoints = vec![
            create_test_endpoint("tier2-unhealthy.com", 60, 20, false, None, 0),
            create_test_endpoint("tier3-best.com", 20, 100, true, None, 0),
            create_test_endpoint("tier3-worst.com", 10, 50, true, None, 0),
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier3-best.com");
    }

    #[test]
    fn test_tier3_sorts_by_weight_then_lru_ignores_cost() {
        let endpoints = vec![
            create_test_endpoint("tier3-a.com", 20, 500, true, None, 15), // Better weight, worse cost, newer
            create_test_endpoint("tier3-b.com", 10, 50, true, None, 5), // Worse weight, better cost, older
            create_test_endpoint("tier3-c.com", 20, 100, true, None, 8), // Better weight, better cost, older than a
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier3-c.com"); // Better weight, then LRU (ignores cost)
    }

    #[test]
    fn test_tier3_lru_rotation() {
        let endpoints = vec![
            create_test_endpoint("tier3-a.com", 20, 1000, true, None, 10), // Same weight, older, worse cost
            create_test_endpoint("tier3-b.com", 20, 10, true, None, 15), // Same weight, newer, better cost
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier3-a.com"); // LRU wins, cost ignored
    }

    // --- Edge Case Tests ---

    #[test]
    fn test_no_available_endpoints() {
        let endpoints = vec![
            create_test_endpoint("unhealthy.com", 100, 50, false, None, 0),
            create_test_endpoint("cooldown.com", 50, 100, true, Some(60), 0),
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1);
        assert!(selected.is_none());
    }

    #[test]
    fn test_all_endpoints_rate_limited() {
        let endpoints = vec![
            create_test_endpoint("tier1.com", 100, 50, true, None, 0),
            create_test_endpoint("tier2.com", 60, 100, true, None, 0),
            create_test_endpoint("tier3.com", 20, 200, true, None, 0),
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();

        let mut limiters = HashMap::new();
        for url in ["tier1.com", "tier2.com", "tier3.com"] {
            let limiter = create_rate_limiter(5);
            limiter.lock().check_n(5).unwrap(); // Use up all tokens
            limiters.insert(url.to_string(), limiter);
        }

        let limiters_rwlock = RwLock::new(limiters);
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1);
        assert!(selected.is_none());
    }

    #[test]
    fn test_empty_endpoints_list() {
        let endpoints: Vec<RpcEndpoint> = vec![];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1);
        assert!(selected.is_none());
    }

    // --- Additional Test Coverage from Original Version ---

    #[test]
    fn test_skip_cooldown_endpoint() {
        let endpoints = vec![
            create_test_endpoint("tier2-cooldown.com", 70, 50, true, Some(60), 0), // in cooldown
            create_test_endpoint("tier2-ok.com", 60, 100, true, None, 0), // next best in tier
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier2-ok.com");
    }

    #[test]
    fn test_batch_selects_highest_priority_when_available() {
        let endpoints = vec![
            create_test_endpoint("tier2-a.com", 60, 100, true, None, 0),
            create_test_endpoint("tier2-b.com", 70, 100, true, None, 0), // highest weight
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();

        let mut limiters = HashMap::new();
        limiters.insert("tier2-a.com".to_string(), create_rate_limiter(20));
        limiters.insert("tier2-b.com".to_string(), create_rate_limiter(20));
        let limiters_rwlock = RwLock::new(limiters);
        let limiters_lock = limiters_rwlock.read();

        // Batch of 10 should select tier2-b, as it has higher weight and enough capacity
        let selected = run_selection(&endpoints_lock, &limiters_lock, 10).unwrap();
        assert_eq!(selected.url, "tier2-b.com");
    }

    #[test]
    fn test_select_endpoint_with_no_rate_limiter_entry() {
        let endpoints = vec![create_test_endpoint("no-limiter.com", 100, 50, true, None, 0)];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new()); // Limiters map is empty
        let limiters_lock = limiters_rwlock.read();

        // Should be selected as it's considered to have no rate limit
        let selected = run_selection(&endpoints_lock, &limiters_lock, 100).unwrap();
        assert_eq!(selected.url, "no-limiter.com");
    }

    #[test]
    fn test_tier2_full_tie_break_is_stable() {
        let endpoints = vec![
            create_test_endpoint("tier2-c.com", 70, 100, true, None, 5),
            create_test_endpoint("tier2-a.com", 80, 50, true, None, 1), // Identical to B
            create_test_endpoint("tier2-b.com", 80, 50, true, None, 1), // Identical to A
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        // Because Rust's sort is stable, "a" should be chosen because it appeared first in the vec
        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected.url, "tier2-a.com");
    }

    // --- Helper Function Tests ---

    #[test]
    fn test_reserve_tokens_zero_batch_size() {
        let mut limiters = HashMap::new();
        let limiter = create_rate_limiter(1);
        limiter.lock().check_n(1).unwrap(); // Exhaust the limiter
        limiters.insert("e1.com".to_string(), limiter);
        let limiters_rwlock = RwLock::new(limiters);
        let limiters_lock = limiters_rwlock.read();

        // A batch of zero should always be Ok, even if the limiter is full.
        let result = reserve_tokens(&limiters_lock, "e1.com", 0);
        assert!(result.is_ok());
    }

    #[test]
    fn test_reserve_tokens_oversized_batch() {
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();
        let oversized_batch = (u32::MAX as usize) + 1;

        let result = reserve_tokens(&limiters_lock, "e1.com", oversized_batch);
        assert!(matches!(result, Err(LoadBalancerError::UpstreamError(_))));
    }

    #[test]
    fn test_reserve_tokens_rate_limited() {
        let mut limiters = HashMap::new();
        let limiter = create_rate_limiter(1);
        limiter.lock().check_n(1).unwrap(); // Exhaust the limiter
        limiters.insert("e1.com".to_string(), limiter);
        let limiters_rwlock = RwLock::new(limiters);
        let limiters_lock = limiters_rwlock.read();

        let result = reserve_tokens(&limiters_lock, "e1.com", 1);
        assert!(matches!(result, Err(LoadBalancerError::RateLimited(_))));
    }

    // --- Missing Tier 3 Tests ---

    #[test]
    fn test_tier3_ignores_cost_even_with_different_weights() {
        let endpoints = vec![
            create_test_endpoint("tier3-lowweight-lowcost.com", 10, 10, true, None, 5),
            create_test_endpoint("tier3-highweight-highcost.com", 30, 1000, true, None, 10),
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(
            selected.url, "tier3-highweight-highcost.com",
            "Tier 3 should prioritize weight over cost"
        );
    }

    #[test]
    fn test_tier3_mixed_weights_still_prioritize_weight() {
        let endpoints = vec![
            create_test_endpoint("tier3-weight25.com", 25, 1000, true, None, 1), // Higher weight, high cost, recent
            create_test_endpoint("tier3-weight20.com", 20, 10, true, None, 100), // Lower weight, low cost, older
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(
            selected.url, "tier3-weight25.com",
            "Tier 3 should prioritize higher weight even with worse cost/recency"
        );
    }

    // --- LRU Rotation Tests ---

    #[test]
    fn test_tier1_lru_rotation() {
        let endpoints = vec![
            create_test_endpoint("tier1-a.com", 100, 100, true, None, 10),
            create_test_endpoint("tier1-b.com", 100, 100, true, None, 5),
        ];
        let endpoints_rwlock = RwLock::new(endpoints);
        let endpoints_lock = endpoints_rwlock.read();
        let limiters_rwlock = RwLock::new(HashMap::new());
        let limiters_lock = limiters_rwlock.read();

        let selected1 = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected1.url, "tier1-b.com"); // Oldest timestamp
        selected1.metrics.last_selected.store(20, std::sync::atomic::Ordering::Relaxed);

        let selected2 = run_selection(&endpoints_lock, &limiters_lock, 1).unwrap();
        assert_eq!(selected2.url, "tier1-a.com"); // Now the oldest
    }
}