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