pick_fast 0.1.3

#![cfg_attr(docsrs, feature(doc_cfg))]

use std::{
  marker::PhantomData,
  ops::Deref,
  sync::atomic::{AtomicU32, Ordering},
};

/// 策略 Trait: 定义耗时与权重的转换逻辑
pub trait Rank {
  /// 计算权重
  /// - 输入: `latency_us` (微秒)
  /// - 输出: `weight` (权重值)
  fn calc(latency_us: u32) -> u32;

  /// 初始权重 (默认为 1，实现慢启动)
  fn init() -> u32 {
    1
  }
}

/// 默认策略: 倒数模型 (Inverse)
/// 逻辑: Weight = BASE / Latency
/// 场景: 适用于 N <= 256 的通用负载均衡
pub struct Inverse;

impl Rank for Inverse {
  #[inline(always)]
  fn calc(latency: u32) -> u32 {
    // 调整为 2^22 (约 4,194,304)
    // 计算: 256 * 2^22 = 2^30 (约10亿) < u32::MAX (约42亿)
    // 即使有 256 个节点且延时均为 1us，也不会溢出。
    const BASE: u32 = 1 << 22;

    // 精度说明:
    // 1ms (1000us) -> 权重 4194
    // 50ms(50000us)-> 权重 83
    // 1s  (10^6us) -> 权重 4
    // 精度足够区分不同节点节点的性能差异。

    BASE / latency.max(1)
  }
}

/// 节点: 包含数据和权重
pub struct Node<T> {
  /// 节点数据
  pub data: T,
  /// 节点权重 (原子)
  pub weight: AtomicU32,
}

impl<T> Node<T> {
  /// 创建节点
  pub fn new(data: T, weight: u32) -> Self {
    Self {
      data,
      weight: AtomicU32::new(weight),
    }
  }
}

impl<T> Deref for Node<T> {
  type Target = T;
  fn deref(&self) -> &Self::Target {
    &self.data
  }
}

/// 选中节点的句柄
/// 包含节点引用和索引
pub struct Handle<'a, T> {
  pub index: usize,
  pub node: &'a Node<T>,
}

impl<'a, T> Deref for Handle<'a, T> {
  type Target = T;
  fn deref(&self) -> &Self::Target {
    &self.node.data
  }
}

// ========================================================================
// 核心结构 (PickFast)
// ========================================================================

/// 极速加权负载均衡器
///
/// 泛型参数:
/// - `T`: 节点数据类型
/// - `M`: 权重模型 (默认为 `Inverse`)
pub struct PickFast<T, M = Inverse> {
  /// 节点列表
  pub li: Vec<Node<T>>,

  /// 总权重 (原子缓存)
  pub total: AtomicU32,

  _marker: PhantomData<M>,
}

unsafe impl<T: Sync, M> Sync for PickFast<T, M> {}
unsafe impl<T: Send, M> Send for PickFast<T, M> {}

impl<T, M: Rank> PickFast<T, M> {
  /// 创建一个新的选择器
  pub fn new(data: impl IntoIterator<Item = T>) -> Self {
    let init_val = M::init();
    let li: Vec<Node<T>> = data
      .into_iter()
      .map(|d| Node::new(d, init_val))
      .collect();

    let n = li.len();
    assert!(n > 0, "PickFast: node count must be > 0");

    // 如果 N 非常大，这里给个运行时提醒
    if n > 256 {
      log::warn!("PickFast n={n} is large, ensure Rank won't overflow u32");
    }

    // 初始化总权重
    let total = AtomicU32::new(init_val * (n as u32));

    Self {
      li,
      total,
      _marker: PhantomData,
    }
  }

  /// 节点数量
  #[inline(always)]
  pub fn len(&self) -> usize {
    self.li.len()
  }

  /// 是否为空
  #[inline(always)]
  pub fn is_empty(&self) -> bool {
    self.li.is_empty()
  }

  /// 极速挑选 (Pick)
  ///
  /// O(1) 获取总权重 -> O(N) 扫描
  #[inline(always)]
  pub fn pick(&self) -> Handle<'_, T> {
    let total_w = self.total.load(Ordering::Relaxed);

    if total_w == 0 {
      return Handle {
        index: 0,
        node: &self.li[0],
      };
    }

    // 随机目标
    let target = fastrand::u32(0..total_w);
    let mut sum = 0;

    // 扫描
    for (i, node) in self.li.iter().enumerate() {
      sum += node.weight.load(Ordering::Relaxed);
      if sum > target {
        return Handle { index: i, node };
      }
    }

    // 兜底 (处理并发更新时的微小窗口)
    let last = self.li.len() - 1;
    Handle {
      index: last,
      node: &self.li[last],
    }
  }

  /// 设定观测值 (Set)
  ///
  /// 传入观测值 (如耗时)，内部自动计算权重并平滑更新
  #[inline(always)]
  pub fn set(&self, index: usize, val: u32) {
    if index >= self.li.len() {
      return;
    }

    let target_w = M::calc(val);

    // CAS 更新单节点权重
    // EMA公式: New = (Old + Target) / 2
    let _ = self.li[index]
      .weight
      .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |old| {
        Some((old + target_w) >> 1)
      })
      .map(|prev| {
        // 修正总权重
        let new_w = (prev + target_w) >> 1;
        if new_w > prev {
          self.total.fetch_add(new_w - prev, Ordering::Relaxed);
        } else {
          self.total.fetch_sub(prev - new_w, Ordering::Relaxed);
        }
      });
  }
}