use std::net::SocketAddr;
use std::ops::{Deref, DerefMut};
use crate::client::HttpClient;
use crate::error::HttpError;
pub struct PoolConfig {
pub addr: SocketAddr,
pub host: String,
pub pool_size: usize,
pub protocol: Protocol,
pub connect_timeout_ms: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Protocol {
H2,
H1,
H1Plain,
}
enum Slot {
Connected(Box<HttpClient>),
Disconnected,
}
pub struct Pool {
addr: SocketAddr,
host: String,
protocol: Protocol,
slots: Vec<Slot>,
next: usize,
connect_timeout_ms: u64,
}
impl Pool {
pub fn new(config: PoolConfig) -> Self {
let mut slots = Vec::with_capacity(config.pool_size);
for _ in 0..config.pool_size {
slots.push(Slot::Disconnected);
}
Pool {
addr: config.addr,
host: config.host,
protocol: config.protocol,
slots,
next: 0,
connect_timeout_ms: config.connect_timeout_ms,
}
}
pub async fn connect_all(&mut self) -> Result<(), HttpError> {
for i in 0..self.slots.len() {
let client = self.do_connect().await?;
self.slots[i] = Slot::Connected(Box::new(client));
}
Ok(())
}
pub async fn client(&mut self) -> Result<PooledClient<'_>, HttpError> {
let size = self.slots.len();
for _ in 0..size {
let idx = self.next;
self.next = (self.next + 1) % size;
match &self.slots[idx] {
Slot::Connected(_) => {
return Ok(PooledClient { pool: self, idx });
}
Slot::Disconnected => {
if let Ok(client) = self.do_connect().await {
self.slots[idx] = Slot::Connected(Box::new(client));
return Ok(PooledClient { pool: self, idx });
}
}
}
}
Err(HttpError::AllConnectionsFailed)
}
pub fn mark_disconnected(&mut self, idx: usize) {
if idx < self.slots.len() {
if let Slot::Connected(client) = &self.slots[idx] {
client.close();
}
self.slots[idx] = Slot::Disconnected;
}
}
pub fn close_all(&mut self) {
for slot in &mut self.slots {
if let Slot::Connected(client) = slot {
client.close();
}
*slot = Slot::Disconnected;
}
}
pub fn connected_count(&self) -> usize {
self.slots
.iter()
.filter(|s| matches!(s, Slot::Connected(_)))
.count()
}
pub fn pool_size(&self) -> usize {
self.slots.len()
}
async fn do_connect(&self) -> Result<HttpClient, HttpError> {
match self.protocol {
Protocol::H2 => {
if self.connect_timeout_ms > 0 {
HttpClient::connect_h2_with_timeout(
self.addr,
&self.host,
self.connect_timeout_ms,
)
.await
} else {
HttpClient::connect_h2(self.addr, &self.host).await
}
}
Protocol::H1 => HttpClient::connect_h1(self.addr, &self.host).await,
Protocol::H1Plain => HttpClient::connect_h1_plain(self.addr, &self.host).await,
}
}
}
pub struct PooledClient<'a> {
pool: &'a mut Pool,
idx: usize,
}
impl PooledClient<'_> {
pub fn slot(&self) -> usize {
self.idx
}
}
impl Deref for PooledClient<'_> {
type Target = HttpClient;
fn deref(&self) -> &HttpClient {
match &self.pool.slots[self.idx] {
Slot::Connected(client) => client,
Slot::Disconnected => unreachable!("PooledClient over Disconnected slot"),
}
}
}
impl DerefMut for PooledClient<'_> {
fn deref_mut(&mut self) -> &mut HttpClient {
match &mut self.pool.slots[self.idx] {
Slot::Connected(client) => client,
Slot::Disconnected => unreachable!("PooledClient over Disconnected slot"),
}
}
}
impl Drop for PooledClient<'_> {
fn drop(&mut self) {
let should_recycle = match &self.pool.slots[self.idx] {
Slot::Connected(client) => client.peer_will_close(),
Slot::Disconnected => false,
};
if should_recycle {
self.pool.mark_disconnected(self.idx);
}
}
}