dope_session/pool/client/
connect.rs1use std::os::fd::{AsRawFd, IntoRawFd};
2use std::pin::Pin;
3use std::time::Instant;
4
5use dope_core::driver::{DriverOps, FixedBuffers, PoolDriver};
6use dope_core::{Fd, FdHandle, FixedFd};
7
8use crate::protocol::client::ClientProtocol;
9use dope_core::profile::Profile;
10use dope_transport::Transport;
11
12use super::super::Pool;
13use super::super::core::PoolCore;
14use super::Client;
15use super::connect_source::ConnectSource;
16use super::state::{ClientSlotPhase, ClientSlotState};
17use crate::SlotId;
18
19pub(super) enum ConnectPoll {
20 Pending,
21 Ready {
22 ok: bool,
23 tag: u32,
24 raw_fd: i32,
25 fixed_idx: u32,
26 },
27}
28
29impl<P, T, S> Client<P, T, S>
30where
31 P: ClientProtocol,
32 T: Transport,
33 T::Addr: Clone,
34 S: ConnectSource<T>,
35{
36 pub(super) fn spawn_connect<B: PoolDriver>(
37 &mut self,
38 core: &mut PoolCore<T, P::CodecLayer, ClientSlotState<P, B>, B>,
39 driver: &mut B::Driver,
40 addr: T::Addr,
41 tag: u32,
42 ) -> bool {
43 let Some((slot_idx, _slot_id)) = core.reserve_slot() else {
44 return false;
45 };
46
47 let fd: FdHandle = match T::create_fd(&addr) {
48 Ok(fd) => fd,
49 Err(_) => {
50 core.release_reserved(slot_idx);
51 return false;
52 }
53 };
54 let _ = T::apply_stream_socket_opts(&fd, &self.socket_opts);
55 let fixed_idx = match driver.register_fixed(fd.as_fd()) {
56 Ok(Some(idx)) => idx,
57 _ => {
58 core.release_reserved(slot_idx);
59 return false;
60 }
61 };
62 let raw_fd = fd.as_fd().as_raw_fd();
63 let owned = fd.into_owned();
64 let _ = owned.into_raw_fd();
65
66 let sock_addr = match T::to_sock_addr(addr) {
67 Ok(a) => a,
68 Err(_) => {
69 driver.unregister_fixed(fixed_idx);
70
71 unsafe { libc::close(raw_fd) };
72 core.release_reserved(slot_idx);
73 return false;
74 }
75 };
76 let connect_req = match driver.try_submit_connect_raw(Fd::new(raw_fd), sock_addr) {
77 Ok(req) => req,
78 Err((_, _)) => {
79 driver.unregister_fixed(fixed_idx);
80 unsafe { libc::close(raw_fd) };
81 core.release_reserved(slot_idx);
82 return false;
83 }
84 };
85
86 let user = core.slots[slot_idx].user.get_mut();
87 user.upstream_tag = tag;
88 user.phase = ClientSlotPhase::Connecting {
89 fixed_idx,
90 raw_fd,
91 request: connect_req,
92 };
93
94 self.connecting_slots.push_back(slot_idx as u16);
95 true
96 }
97}
98
99impl<P, T, S, F, B> Pool<Client<P, T, S>, F, B>
100where
101 P: ClientProtocol,
102 T: Transport,
103 T::Addr: Clone,
104 S: ConnectSource<T>,
105 F: Profile,
106 B: PoolDriver,
107{
108 pub(super) fn poll_connect_requests(&mut self, driver: &mut B::Driver) {
109 let waker = std::task::Waker::noop();
110 let mut cx = std::task::Context::from_waker(waker);
111 let mut i = 0;
112 while i < self.direction.connecting_slots.len() {
113 let slot = self.direction.connecting_slots[i];
114 let slot_idx = slot as usize;
115
116 let outcome = {
117 let user = self.core.slots[slot_idx].user.get_mut();
118 let phase = std::mem::replace(&mut user.phase, ClientSlotPhase::Idle);
119 match phase {
120 ClientSlotPhase::Connecting {
121 fixed_idx,
122 raw_fd,
123 mut request,
124 } => match Pin::new(&mut request).poll(&mut cx) {
125 std::task::Poll::Pending => {
126 user.phase = ClientSlotPhase::Connecting {
127 fixed_idx,
128 raw_fd,
129 request,
130 };
131 Some(ConnectPoll::Pending)
132 }
133 std::task::Poll::Ready((res, _flags, _addr)) => Some(ConnectPoll::Ready {
134 ok: res.is_ok(),
135 tag: user.upstream_tag,
136 raw_fd,
137 fixed_idx,
138 }),
139 },
140 other => {
141 user.phase = other;
142 None
143 }
144 }
145 };
146
147 match outcome {
148 None => {
149 self.direction.connecting_slots.swap_remove_back(i);
150 }
151 Some(ConnectPoll::Pending) => {
152 i += 1;
153 }
154 Some(ConnectPoll::Ready {
155 ok,
156 tag,
157 raw_fd,
158 fixed_idx,
159 }) => {
160 self.direction.connecting_slots.swap_remove_back(i);
161 let now = Instant::now();
162 if ok {
163 let fixed_fd = FixedFd::new(raw_fd, fixed_idx);
164 if !self.core.activate_slot(driver, slot_idx, fixed_fd) {
165 driver.unregister_fixed(fixed_idx);
166 unsafe { libc::close(raw_fd) };
167 self.direction
168 .connect_source
169 .on_connect_outcome(tag, false, now);
170 self.core.release_reserved(slot_idx);
171 continue;
172 }
173 self.core.slots[slot_idx].user.get_mut().phase = ClientSlotPhase::Active;
174 self.direction
175 .connect_source
176 .on_connect_outcome(tag, true, now);
177 let conn_id = SlotId::from_parts(
178 slot_idx as u32,
179 self.core.slots[slot_idx].conn.generation,
180 );
181 let entry = &mut self.core.slots[slot_idx];
182 let user = entry.user.get_mut();
183 let codec_state = entry.codec_state.get_mut();
184 user.session
185 .on_connect(&mut self.direction.protocol, conn_id, codec_state);
186 self.try_submit_outbound(driver, slot_idx);
187 } else {
188 driver.unregister_fixed(fixed_idx);
189 unsafe { libc::close(raw_fd) };
190 self.direction
191 .connect_source
192 .on_connect_outcome(tag, false, now);
193 self.core.release_reserved(slot_idx);
194 }
195 }
196 }
197 }
198 }
199
200 pub(super) fn disconnect_and_free(&mut self, driver: &mut B::Driver, slot_idx: usize) {
201 let conn_id =
202 SlotId::from_parts(slot_idx as u32, self.core.slots[slot_idx].conn.generation);
203 let tag = {
204 let entry = &mut self.core.slots[slot_idx];
205 let user = entry.user.get_mut();
206 let codec_state = entry.codec_state.get_mut();
207 user.session
208 .on_disconnect(&mut self.direction.protocol, conn_id, codec_state);
209 user.upstream_tag
210 };
211 self.direction
212 .connect_source
213 .on_disconnect(tag, Instant::now());
214 self.core.free_slot(driver, slot_idx);
215 }
216}