use super::BalanceMake;
use rand::{rngs::SmallRng, Rng, SeedableRng};
use std::{fmt, marker::PhantomData};
use tower_layer::Layer;
#[derive(Clone)]
pub struct BalanceLayer<D, Req> {
rng: SmallRng,
_marker: PhantomData<fn(D, Req)>,
}
impl<D, Req> BalanceLayer<D, Req> {
pub fn new() -> Self {
Self {
rng: SmallRng::from_entropy(),
_marker: PhantomData,
}
}
pub fn from_rng<R: Rng>(rng: &mut R) -> Result<Self, rand::Error> {
let rng = SmallRng::from_rng(rng)?;
Ok(Self {
rng,
_marker: PhantomData,
})
}
}
impl<S, Req> Layer<S> for BalanceLayer<S, Req> {
type Service = BalanceMake<S, Req>;
fn layer(&self, make_discover: S) -> Self::Service {
BalanceMake::new(make_discover, self.rng.clone())
}
}
impl<D, Req> fmt::Debug for BalanceLayer<D, Req> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("BalanceLayer")
.field("rng", &self.rng)
.finish()
}
}