loadwise-core 0.1.0

//! Load balancing strategies and the [`Strategy`] trait.
//!
//! Built-in strategies:
//!
//! | Strategy                  | Requires        | Algorithm                              |
//! |---------------------------|-----------------|----------------------------------------|
//! | [`RoundRobin`]            | —               | Sequential cycling                     |
//! | [`WeightedRoundRobin`]    | `Weighted+Node` | Nginx smooth weighted round-robin      |
//! | [`Random`]                | —               | Uniform random                         |
//! | [`WeightedRandom`]        | `Weighted`      | Proportional random                    |
//! | [`LeastLoad`]             | `LoadMetric`    | Pick lowest `load_score()`             |
//! | [`PowerOfTwoChoices`]     | `LoadMetric`    | Random 2, pick lower load              |
//! | [`ConsistentHash`]        | `Node`          | Virtual-node ring (cached)             |
//! | [`MostAvailable`]         | `RateMetric`    | Pick highest remaining, epsilon tie-break |
//!
//! Combinators: [`WithFallback`], [`FallbackChain`].

mod chain;
mod consistent_hash;
mod least_load;
mod most_available;
mod p2c;
mod random;
mod round_robin;
mod weighted_round_robin;

pub use chain::{FallbackChain, WithFallback};
pub use consistent_hash::ConsistentHash;
pub use least_load::LeastLoad;
pub use most_available::MostAvailable;
pub use p2c::PowerOfTwoChoices;
pub use random::{Random, WeightedRandom};
pub use round_robin::RoundRobin;
pub use weighted_round_robin::WeightedRoundRobin;

use std::collections::HashMap;
use bon::Builder;

/// Context provided to strategies during node selection.
///
/// # Examples
///
/// ```
/// # extern crate loadwise_core as loadwise;
/// use loadwise::SelectionContext;
///
/// // Empty context (most strategies work fine without it).
/// let ctx = SelectionContext::default();
///
/// // With a hash key for consistent hashing.
/// let ctx = SelectionContext::builder()
///     .hash_key(42)
///     .build();
///
/// // With estimated request cost and tags.
/// let ctx = SelectionContext::builder()
///     .estimated_cost(150.0)
///     .tags([("model".into(), "gpt-4".into())].into())
///     .build();
///
/// assert_eq!(ctx.tag("model"), Some("gpt-4"));
///
/// // Exclude specific indices (e.g., retry after failure).
/// let ctx = SelectionContext::builder()
///     .exclude(vec![0, 2])
///     .build();
/// assert!(ctx.is_excluded(0));
/// assert!(!ctx.is_excluded(1));
/// ```
#[derive(Debug, Clone, Default, Builder)]
pub struct SelectionContext {
    /// Estimated cost/weight of this request (e.g., token count).
    pub estimated_cost: Option<f64>,
    /// Hash key for consistent hashing strategies.
    pub hash_key: Option<u64>,
    /// Arbitrary tags for routing decisions. `None` by default to avoid allocation.
    pub tags: Option<HashMap<String, String>>,
    /// Indices to skip during selection (for retry scenarios).
    /// `None` by default to avoid allocation.
    pub exclude: Option<Vec<usize>>,
}

impl SelectionContext {
    /// Look up a tag value by key.
    pub fn tag(&self, key: &str) -> Option<&str> {
        self.tags.as_ref()?.get(key).map(|s| s.as_str())
    }

    /// Returns `true` if the given index should be skipped.
    ///
    /// Uses linear search — O(n) in the size of the exclude list.
    /// Fine for typical retry scenarios (1–3 entries); for larger lists
    /// consider pre-filtering candidates instead.
    pub fn is_excluded(&self, idx: usize) -> bool {
        self.exclude.as_ref().is_some_and(|e| e.contains(&idx))
    }
}

/// A load balancing strategy that selects a node from a list of candidates.
///
/// Returns the **index** of the selected node within the `candidates` slice.
/// Implementations must use interior mutability for any state (e.g., `AtomicUsize`).
///
/// # Examples
///
/// ```
/// # extern crate loadwise_core as loadwise;
/// use loadwise::{Strategy, SelectionContext};
/// use loadwise::strategy::RoundRobin;
///
/// let rr = RoundRobin::new();
/// let nodes = ["a", "b", "c"];
/// let ctx = SelectionContext::default();
///
/// assert_eq!(rr.select(&nodes, &ctx), Some(0));
/// assert_eq!(rr.select(&nodes, &ctx), Some(1));
/// assert_eq!(rr.select(&nodes, &ctx), Some(2));
/// assert_eq!(rr.select(&nodes, &ctx), Some(0)); // wraps around
/// ```
pub trait Strategy<N>: Send + Sync {
    /// Pick the next node. Returns `None` when `candidates` is empty or all
    /// candidates are excluded via [`SelectionContext::exclude`].
    fn select(&self, candidates: &[N], ctx: &SelectionContext) -> Option<usize>;
}

// ── internal helpers ──

use std::hash::{DefaultHasher, Hash, Hasher};
use crate::Node;

/// Order-sensitive fingerprint of a candidate set based on node IDs.
///
/// Used by strategies that cache internal state (e.g., hash ring, WRR weights)
/// to detect when the candidate list has changed.
///
/// NOTE: this is O(n) per call. For very large candidate lists, consider
/// letting the caller notify the strategy of changes instead of probing
/// every time. Good enough for typical pool sizes (10–100 nodes).
pub(crate) fn node_set_fingerprint<N: Node>(candidates: &[N]) -> u64 {
    let mut hasher = DefaultHasher::new();
    candidates.len().hash(&mut hasher);
    for node in candidates {
        node.id().hash(&mut hasher);
    }
    hasher.finish()
}