spacegate_kernel/helper_layers/balancer/
ip_hash.rs

1use crate::{extension::PeerAddr, SgRequest};
2
3use super::BalancePolicy;
4use std::{
5    hash::{DefaultHasher, Hash, Hasher},
6    marker::PhantomData,
7};
8
9/// A policy that selects an instance based on the hash of the IP address.
10#[derive(Debug, Clone)]
11pub struct IpHash<H = DefaultHasher> {
12    hasher: PhantomData<fn() -> H>,
13}
14
15impl Default for IpHash {
16    fn default() -> Self {
17        Self { hasher: PhantomData }
18    }
19}
20
21impl<S, H> BalancePolicy<S, SgRequest> for IpHash<H>
22where
23    H: Hasher + Default,
24{
25    fn pick<'s>(&self, instances: &'s [S], req: &SgRequest) -> Option<&'s S> {
26        if instances.is_empty() {
27            None
28        } else if instances.len() == 1 {
29            instances.first()
30        } else {
31            let mut hasher = H::default();
32            let ip = req.extensions().get::<PeerAddr>()?.0.ip();
33            ip.to_canonical().hash(&mut hasher);
34            let hash = hasher.finish();
35            let index = (hash % instances.len() as u64) as usize;
36            instances.get(index)
37        }
38    }
39}