1use std::net::SocketAddr;
7
8use crate::client::HttpClient;
9use crate::error::HttpError;
10
11pub struct PoolConfig {
13 pub addr: SocketAddr,
15 pub host: String,
17 pub pool_size: usize,
19 pub protocol: Protocol,
21 pub connect_timeout_ms: u64,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum Protocol {
28 H2,
30 H1,
32 H1Plain,
34}
35
36enum Slot {
37 Connected(Box<HttpClient>),
38 Disconnected,
39}
40
41pub struct Pool {
43 addr: SocketAddr,
44 host: String,
45 protocol: Protocol,
46 slots: Vec<Slot>,
47 next: usize,
48 connect_timeout_ms: u64,
49}
50
51impl Pool {
52 pub fn new(config: PoolConfig) -> Self {
54 let mut slots = Vec::with_capacity(config.pool_size);
55 for _ in 0..config.pool_size {
56 slots.push(Slot::Disconnected);
57 }
58 Pool {
59 addr: config.addr,
60 host: config.host,
61 protocol: config.protocol,
62 slots,
63 next: 0,
64 connect_timeout_ms: config.connect_timeout_ms,
65 }
66 }
67
68 pub async fn connect_all(&mut self) -> Result<(), HttpError> {
70 for i in 0..self.slots.len() {
71 let client = self.do_connect().await?;
72 self.slots[i] = Slot::Connected(Box::new(client));
73 }
74 Ok(())
75 }
76
77 pub async fn client(&mut self) -> Result<&mut HttpClient, HttpError> {
83 let size = self.slots.len();
84 for _ in 0..size {
85 let idx = self.next;
86 self.next = (self.next + 1) % size;
87
88 match &self.slots[idx] {
89 Slot::Connected(_) => {
90 if let Slot::Connected(client) = &mut self.slots[idx] {
91 return Ok(client);
92 }
93 unreachable!();
94 }
95 Slot::Disconnected => {
96 if let Ok(client) = self.do_connect().await {
97 self.slots[idx] = Slot::Connected(Box::new(client));
98 if let Slot::Connected(client) = &mut self.slots[idx] {
99 return Ok(client);
100 }
101 unreachable!();
102 }
103 }
104 }
105 }
106 Err(HttpError::AllConnectionsFailed)
107 }
108
109 pub fn mark_disconnected(&mut self, idx: usize) {
111 if idx < self.slots.len() {
112 if let Slot::Connected(client) = &self.slots[idx] {
113 client.close();
114 }
115 self.slots[idx] = Slot::Disconnected;
116 }
117 }
118
119 pub fn close_all(&mut self) {
121 for slot in &mut self.slots {
122 if let Slot::Connected(client) = slot {
123 client.close();
124 }
125 *slot = Slot::Disconnected;
126 }
127 }
128
129 pub fn connected_count(&self) -> usize {
131 self.slots
132 .iter()
133 .filter(|s| matches!(s, Slot::Connected(_)))
134 .count()
135 }
136
137 pub fn pool_size(&self) -> usize {
139 self.slots.len()
140 }
141
142 async fn do_connect(&self) -> Result<HttpClient, HttpError> {
143 match self.protocol {
144 Protocol::H2 => {
145 if self.connect_timeout_ms > 0 {
146 HttpClient::connect_h2_with_timeout(
147 self.addr,
148 &self.host,
149 self.connect_timeout_ms,
150 )
151 .await
152 } else {
153 HttpClient::connect_h2(self.addr, &self.host).await
154 }
155 }
156 Protocol::H1 => HttpClient::connect_h1(self.addr, &self.host).await,
157 Protocol::H1Plain => HttpClient::connect_h1_plain(self.addr, &self.host).await,
158 }
159 }
160}