load_balancer/
ip.rs

1use crate::{BoxLoadBalancer, LoadBalancer, interval::IntervalLoadBalancer};
2use async_trait::async_trait;
3use get_if_addrs::get_if_addrs;
4use reqwest::{Client, ClientBuilder};
5use std::{net::IpAddr, sync::Arc, time::Duration};
6
7/// Load balancer for `reqwest::Client` instances bound to specific IP addresses.
8/// Uses interval-based allocation.
9#[derive(Clone)]
10pub struct IPClientLoadBalancer {
11    inner: IntervalLoadBalancer<Client>,
12}
13
14impl IPClientLoadBalancer {
15    /// Create a new interval-based load balancer with given clients.
16    pub fn new(entries: Vec<(Duration, Client)>) -> Self {
17        Self {
18            inner: IntervalLoadBalancer::new(entries),
19        }
20    }
21
22    /// Build a load balancer using all local IP addresses.
23    pub fn with_ip(ip: Vec<IpAddr>, interval: Duration) -> Self {
24        Self {
25            inner: IntervalLoadBalancer::new(
26                ip.into_iter()
27                    .map(|v| {
28                        (
29                            interval,
30                            ClientBuilder::new().local_address(v).build().unwrap(),
31                        )
32                    })
33                    .collect(),
34            ),
35        }
36    }
37
38    /// Build a load balancer using IP addresses with a per-client timeout.
39    pub fn with_timeout(ip: Vec<IpAddr>, interval: Duration, timeout: Duration) -> Self {
40        Self {
41            inner: IntervalLoadBalancer::new(
42                ip.into_iter()
43                    .map(|v| {
44                        (
45                            interval,
46                            ClientBuilder::new()
47                                .local_address(v)
48                                .timeout(timeout)
49                                .build()
50                                .unwrap(),
51                        )
52                    })
53                    .collect(),
54            ),
55        }
56    }
57
58    /// Update the internal load balancer.
59    pub async fn update<F, R>(&self, handle: F) -> anyhow::Result<()>
60    where
61        F: Fn(Arc<std::sync::RwLock<Vec<crate::interval::Entry<Client>>>>) -> R,
62        R: std::future::Future<Output = anyhow::Result<()>>,
63    {
64        self.inner.update(handle).await
65    }
66}
67
68impl LoadBalancer<Client> for IPClientLoadBalancer {
69    fn alloc(&self) -> impl std::future::Future<Output = Client> + Send {
70        LoadBalancer::alloc(&self.inner)
71    }
72
73    fn try_alloc(&self) -> Option<Client> {
74        LoadBalancer::try_alloc(&self.inner)
75    }
76}
77
78#[async_trait]
79impl BoxLoadBalancer<Client> for IPClientLoadBalancer {
80    async fn alloc(&self) -> Client {
81        LoadBalancer::alloc(self).await
82    }
83
84    fn try_alloc(&self) -> Option<Client> {
85        LoadBalancer::try_alloc(self)
86    }
87}
88
89/// Get all non-loopback IP addresses of the machine.
90pub fn get_ip_list() -> anyhow::Result<Vec<IpAddr>> {
91    Ok(get_if_addrs()?
92        .into_iter()
93        .filter(|v| !v.is_loopback())
94        .map(|v| v.ip())
95        .collect::<Vec<_>>())
96}
97
98/// Get all non-loopback IPv4 addresses of the machine.
99pub fn get_ipv4_list() -> anyhow::Result<Vec<IpAddr>> {
100    Ok(get_if_addrs()?
101        .into_iter()
102        .filter(|v| !v.is_loopback() && v.ip().is_ipv4())
103        .map(|v| v.ip())
104        .collect::<Vec<_>>())
105}
106
107/// Get all non-loopback IPv6 addresses of the machine.
108pub fn get_ipv6_list() -> anyhow::Result<Vec<IpAddr>> {
109    Ok(get_if_addrs()?
110        .into_iter()
111        .filter(|v| !v.is_loopback() && v.ip().is_ipv6())
112        .map(|v| v.ip())
113        .collect::<Vec<_>>())
114}
load_balancer/ip.rs

load_balancer/
ip.rs
