net-pool 0.5.1

use crate::backend::{Address, BackendState};
use crate::utils::bytes_to_hash_code;
use std::collections::BTreeMap;
use std::sync::RwLock;

/// 一致性hash负载策略
pub struct CHStrategy {
    inner: RwLock<Inner>,
}

struct Inner {
    /// usize:虚拟的后端个数
    virtual_count: usize,
    /// 排序后用来表示环形
    tree: BTreeMap<u64, u64>,
    /// HashSet<BackendState>: 真实的节点
    bss: super::backends::BackendStates,
}

impl CHStrategy {
    pub fn new(virtual_count: usize) -> Self {
        if virtual_count == 0 {
            return CHStrategy::default();
        }
        CHStrategy {
            inner: RwLock::new(Inner {
                virtual_count,
                tree: BTreeMap::new(),
                bss: super::backends::BackendStates::new(),
            }),
        }
    }

    fn sub_address(id: usize, addr: &Address) -> String {
        let mut flag = match addr {
            Address::Ori(ori) => ori.clone(),
            Address::Addr(addr) => crate::utils::socketaddr_to_ip_string(addr),
        };

        flag.push('#');
        flag.push_str(&id.to_string());
        flag
    }
}

impl Default for CHStrategy {
    fn default() -> Self {
        CHStrategy {
            inner: RwLock::new(Inner {
                virtual_count: 64,
                tree: BTreeMap::new(),
                bss: super::backends::BackendStates::new(),
            }),
        }
    }
}

impl super::Strategy for CHStrategy {
    fn contain(&self, addr: &Address) -> bool {
        let g = self.inner.read().unwrap();
        g.bss.contain(addr)
    }

    fn add_backend(&self, id: Option<u32>, addr: Address) {
        let src_code = addr.hash_code();
        let mut g = self.inner.write().unwrap();
        // 添加虚拟节点
        for idx in 0..g.virtual_count {
            let address = Self::sub_address(idx, &addr);
            let code = bytes_to_hash_code(address.as_bytes());
            g.tree.insert(code, src_code);
        }
        // 添加真实节点
        g.bss.add(id, addr);
    }

    fn remove_backend(&self, addr: &Address) -> bool {
        let mut g = self.inner.write().unwrap();
        // 删除虚拟节点
        for id in 0..g.virtual_count {
            let address = Self::sub_address(id, addr);
            let code = bytes_to_hash_code(address.as_bytes());
            if let None = g.tree.remove(&code) {
                return false;
            }
        }
        // 删除真实节点
        g.bss.remove(addr)
    }

    /// 获取一个后端节点
    fn get_backend(&self, key: &str) -> Option<BackendState> {
        let g = self.inner.read().unwrap();
        if g.tree.is_empty() {
            return None;
        }

        // 算hash值
        let code = bytes_to_hash_code(key.as_bytes());

        // 获取虚拟节点
        let mut o = g
            .tree
            .iter()
            .find(|(key, _)| if code <= **key { true } else { false });

        if o.is_none() {
            // 使用第一个虚拟节点
            o = g.tree.iter().min();
        };

        g.bss.get_by_code(*o.unwrap().1).map(|b| b.clone())
    }

    fn get_backend_by_id(&self, id: u32) -> Option<BackendState> {
        let g = self.inner.read().unwrap();
        g.bss.get_by_id(id).map(|b| b.clone())
    }

    fn get_backend_by_code(&self, code: u64) -> Option<BackendState> {
        let g = self.inner.read().unwrap();
        g.bss.get_by_code(code).map(|b| b.clone())
    }

    fn get_backends(&self) -> Vec<BackendState> {
        let g = self.inner.read().unwrap();
        let bss = g.bss.get_all();
        bss.iter().map(|bs| bs.clone()).collect()
    }
}

#[test]
fn test() {
    let s = CHStrategy::default();
    super::Strategy::add_backend(&s, None, Address::from(":10001"));
    super::Strategy::add_backend(&s, Some(2), Address::from(":10002"));
    println!("{:?}", super::Strategy::get_backend(&s, "123"));
    println!("{:?}", super::Strategy::get_backend(&s, "123"));
    println!("{:?}", super::Strategy::get_backend(&s, "123"));
    println!("{:?}", super::Strategy::get_backend(&s, "1234"));
    println!("{:?}", super::Strategy::get_backend(&s, "1234"));
    println!("{:?}", super::Strategy::get_backend_by_id(&s, 2));
    println!(
        "{:?}",
        super::Strategy::get_backend_by_code(&s, 15291072903040886511)
    );
    super::Strategy::remove_backend(&s, &Address::from(":10002"));
    println!("{:?}", super::Strategy::get_backend_by_id(&s, 2));
    println!("{:?}", super::Strategy::get_backend(&s, "1234"));
    println!(
        "{:?}",
        super::Strategy::get_backend_by_code(&s, 15291072903040886511)
    );
}