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 IPv4 addresses.
23    pub fn with_ipv4(interval: Duration) -> Self {
24        Self {
25            inner: IntervalLoadBalancer::new(
26                get_ipv4_list()
27                    .into_iter()
28                    .map(|v| {
29                        (
30                            interval,
31                            ClientBuilder::new().local_address(v).build().unwrap(),
32                        )
33                    })
34                    .collect(),
35            ),
36        }
37    }
38
39    /// Build a load balancer using all local IPv6 addresses.
40    pub fn with_ipv6(interval: Duration) -> Self {
41        Self {
42            inner: IntervalLoadBalancer::new(
43                get_ipv6_list()
44                    .into_iter()
45                    .map(|v| {
46                        (
47                            interval,
48                            ClientBuilder::new().local_address(v).build().unwrap(),
49                        )
50                    })
51                    .collect(),
52            ),
53        }
54    }
55
56    /// Build a load balancer using IPv4 addresses with a per-client timeout.
57    pub fn with_ipv4_timeout(interval: Duration, timeout: Duration) -> Self {
58        Self {
59            inner: IntervalLoadBalancer::new(
60                get_ipv4_list()
61                    .into_iter()
62                    .map(|v| {
63                        (
64                            interval,
65                            ClientBuilder::new()
66                                .local_address(v)
67                                .timeout(timeout)
68                                .build()
69                                .unwrap(),
70                        )
71                    })
72                    .collect(),
73            ),
74        }
75    }
76
77    /// Build a load balancer using IPv6 addresses with a per-client timeout.
78    pub fn with_ipv6_timeout(interval: Duration, timeout: Duration) -> Self {
79        Self {
80            inner: IntervalLoadBalancer::new(
81                get_ipv6_list()
82                    .into_iter()
83                    .map(|v| {
84                        (
85                            interval,
86                            ClientBuilder::new()
87                                .local_address(v)
88                                .timeout(timeout)
89                                .build()
90                                .unwrap(),
91                        )
92                    })
93                    .collect(),
94            ),
95        }
96    }
97
98    /// Update the internal load balancer.
99    pub async fn update<F, R>(&self, handle: F) -> anyhow::Result<()>
100    where
101        F: Fn(Arc<std::sync::RwLock<Vec<crate::interval::Entry<Client>>>>) -> R,
102        R: std::future::Future<Output = anyhow::Result<()>>,
103    {
104        self.inner.update(handle).await
105    }
106}
107
108impl LoadBalancer<Client> for IPClientLoadBalancer {
109    fn alloc(&self) -> impl std::future::Future<Output = Client> + Send {
110        LoadBalancer::alloc(&self.inner)
111    }
112
113    fn try_alloc(&self) -> Option<Client> {
114        LoadBalancer::try_alloc(&self.inner)
115    }
116}
117
118#[async_trait]
119impl BoxLoadBalancer<Client> for IPClientLoadBalancer {
120    async fn alloc(&self) -> Client {
121        LoadBalancer::alloc(self).await
122    }
123
124    fn try_alloc(&self) -> Option<Client> {
125        LoadBalancer::try_alloc(self)
126    }
127}
128
129/// Get all non-loopback IP addresses of the machine.
130pub fn get_ip_list() -> Vec<IpAddr> {
131    get_if_addrs()
132        .unwrap()
133        .into_iter()
134        .filter(|v| !v.is_loopback())
135        .map(|v| v.ip())
136        .collect::<Vec<_>>()
137}
138
139/// Get all non-loopback IPv4 addresses of the machine.
140pub fn get_ipv4_list() -> Vec<IpAddr> {
141    get_if_addrs()
142        .unwrap()
143        .into_iter()
144        .filter(|v| !v.is_loopback() && v.ip().is_ipv4())
145        .map(|v| v.ip())
146        .collect::<Vec<_>>()
147}
148
149/// Get all non-loopback IPv6 addresses of the machine.
150pub fn get_ipv6_list() -> Vec<IpAddr> {
151    get_if_addrs()
152        .unwrap()
153        .into_iter()
154        .filter(|v| !v.is_loopback() && v.ip().is_ipv6())
155        .map(|v| v.ip())
156        .collect::<Vec<_>>()
157}
load_balancer/ip.rs

load_balancer/
ip.rs
