use std::os::fd::{AsRawFd, IntoRawFd};
use std::pin::Pin;
use std::time::Instant;
use dope_core::driver::{DriverOps, FixedBuffers, PoolDriver};
use dope_core::{Fd, FdHandle, FixedFd};
use crate::protocol::client::ClientProtocol;
use dope_core::profile::Profile;
use dope_transport::Transport;
use super::super::Pool;
use super::super::core::PoolCore;
use super::Client;
use super::connect_source::ConnectSource;
use super::state::{ClientSlotPhase, ClientSlotState};
use crate::SlotId;
pub(super) enum ConnectPoll {
Pending,
Ready {
ok: bool,
tag: u32,
raw_fd: i32,
fixed_idx: u32,
},
}
impl<P, T, S> Client<P, T, S>
where
P: ClientProtocol,
T: Transport,
T::Addr: Clone,
S: ConnectSource<T>,
{
pub(super) fn spawn_connect<B: PoolDriver>(
&mut self,
core: &mut PoolCore<
T,
P::CodecLayer,
<P::CodecLayer as crate::CodecLayer>::Staging,
ClientSlotState<P, B>,
B,
>,
driver: &mut B::Driver,
addr: T::Addr,
tag: u32,
) -> bool {
let Some((slot_idx, _slot_id)) = core.reserve_slot() else {
return false;
};
let fd: FdHandle = match T::create_fd(&addr) {
Ok(fd) => fd,
Err(_) => {
core.release_reserved(slot_idx);
return false;
}
};
let _ = T::apply_stream_socket_opts(&fd, &self.socket_opts);
let fixed_idx = match driver.register_fixed(fd.as_fd()) {
Ok(Some(idx)) => idx,
_ => {
core.release_reserved(slot_idx);
return false;
}
};
let raw_fd = fd.as_fd().as_raw_fd();
let owned = fd.into_owned();
let _ = owned.into_raw_fd();
let sock_addr = match T::to_sock_addr(addr) {
Ok(a) => a,
Err(_) => {
driver.unregister_fixed(fixed_idx);
unsafe { libc::close(raw_fd) };
core.release_reserved(slot_idx);
return false;
}
};
let connect_req = match driver.try_submit_connect_raw(Fd::new(raw_fd), sock_addr) {
Ok(req) => req,
Err((_, _)) => {
driver.unregister_fixed(fixed_idx);
unsafe { libc::close(raw_fd) };
core.release_reserved(slot_idx);
return false;
}
};
let user = core.slots[slot_idx].user.get_mut();
user.upstream_tag = tag;
user.phase = ClientSlotPhase::Connecting {
fixed_idx,
raw_fd,
request: connect_req,
};
self.connecting_slots.push_back(slot_idx as u16);
true
}
}
impl<P, T, S, F, B> Pool<Client<P, T, S>, F, B>
where
P: ClientProtocol,
T: Transport,
T::Addr: Clone,
S: ConnectSource<T>,
F: Profile,
B: PoolDriver,
{
pub(super) fn poll_connect_requests(&mut self, driver: &mut B::Driver) {
let waker = std::task::Waker::noop();
let mut cx = std::task::Context::from_waker(waker);
let mut i = 0;
while i < self.direction.connecting_slots.len() {
let slot = self.direction.connecting_slots[i];
let slot_idx = slot as usize;
let outcome = {
let user = self.core.slots[slot_idx].user.get_mut();
let phase = std::mem::replace(&mut user.phase, ClientSlotPhase::Idle);
match phase {
ClientSlotPhase::Connecting {
fixed_idx,
raw_fd,
mut request,
} => match Pin::new(&mut request).poll(&mut cx) {
std::task::Poll::Pending => {
user.phase = ClientSlotPhase::Connecting {
fixed_idx,
raw_fd,
request,
};
Some(ConnectPoll::Pending)
}
std::task::Poll::Ready((res, _flags, _addr)) => Some(ConnectPoll::Ready {
ok: res.is_ok(),
tag: user.upstream_tag,
raw_fd,
fixed_idx,
}),
},
other => {
user.phase = other;
None
}
}
};
match outcome {
None => {
self.direction.connecting_slots.swap_remove_back(i);
}
Some(ConnectPoll::Pending) => {
i += 1;
}
Some(ConnectPoll::Ready {
ok,
tag,
raw_fd,
fixed_idx,
}) => {
self.direction.connecting_slots.swap_remove_back(i);
let now = Instant::now();
if ok {
let fixed_fd = FixedFd::new(raw_fd, fixed_idx);
if !self.core.activate_slot(driver, slot_idx, fixed_fd) {
driver.unregister_fixed(fixed_idx);
unsafe { libc::close(raw_fd) };
self.direction
.connect_source
.on_connect_outcome(tag, false, now);
self.core.release_reserved(slot_idx);
continue;
}
self.core.slots[slot_idx].user.get_mut().phase = ClientSlotPhase::Active;
self.direction
.connect_source
.on_connect_outcome(tag, true, now);
let conn_id = SlotId::from_parts(
slot_idx as u32,
self.core.slots[slot_idx].conn.generation,
);
let entry = &mut self.core.slots[slot_idx];
let user = entry.user.get_mut();
let codec_state = entry.codec_state.get_mut();
user.session
.on_connect(&mut self.direction.protocol, conn_id, codec_state);
self.try_submit_outbound(driver, slot_idx);
} else {
driver.unregister_fixed(fixed_idx);
unsafe { libc::close(raw_fd) };
self.direction
.connect_source
.on_connect_outcome(tag, false, now);
self.core.release_reserved(slot_idx);
}
}
}
}
}
pub(super) fn disconnect_and_free(&mut self, driver: &mut B::Driver, slot_idx: usize) {
let conn_id =
SlotId::from_parts(slot_idx as u32, self.core.slots[slot_idx].conn.generation);
let tag = {
let entry = &mut self.core.slots[slot_idx];
let user = entry.user.get_mut();
let codec_state = entry.codec_state.get_mut();
user.session
.on_disconnect(&mut self.direction.protocol, conn_id, codec_state);
user.upstream_tag
};
self.direction
.connect_source
.on_disconnect(tag, Instant::now());
self.core.free_slot(driver, slot_idx);
}
}