Skip to main content

dope_session/pool/client/
write.rs

1use crate::CodecLayer;
2use crate::WriteBufStorage;
3use crate::protocol::client::ClientProtocol;
4use dope_core::driver::{DriverOps, PoolDriver};
5use dope_core::profile::Profile;
6use dope_transport::Transport;
7
8use super::super::Pool;
9use super::super::core::PushError;
10use super::super::slot::{ConnFlags, CoreSlot};
11use super::Client;
12use super::connect_source::ConnectSource;
13use super::state::CLIENT_BYTES_CAP;
14use crate::SlotId;
15
16impl<P, T, S, F, B> Pool<Client<P, T, S>, F, B>
17where
18    P: ClientProtocol,
19    T: Transport,
20    T::Addr: Clone,
21    S: ConnectSource<T>,
22    F: Profile,
23    B: PoolDriver,
24{
25    #[inline]
26    pub(super) fn try_submit_outbound(&mut self, driver: &mut B::Driver, slot_idx: usize) {
27        if !self.core.slots[slot_idx]
28            .conn
29            .flags
30            .contains(ConnFlags::LIVE)
31        {
32            return;
33        }
34        if self.core.slots[slot_idx]
35            .conn
36            .flags
37            .contains(ConnFlags::CLOSE_AFTER)
38        {
39            return;
40        }
41        if <P::CodecLayer as CodecLayer>::IS_PASSTHROUGH {
42            self.try_submit_outbound_plain(driver, slot_idx);
43        } else {
44            self.try_submit_outbound_codec(driver, slot_idx);
45        }
46    }
47
48    #[inline]
49    fn try_submit_outbound_plain(&mut self, driver: &mut B::Driver, slot_idx: usize) {
50        if self.core.slots[slot_idx].io.write.has_in_flight {
51            return;
52        }
53        let conn_fd;
54        let token_epoch: Option<(B::Token, u32)>;
55        let msg_ptr: *const libc::msghdr;
56        let close;
57        {
58            let external = self.core.slot_id_base + slot_idx;
59            let entry = &mut self.core.slots[slot_idx];
60            conn_fd = entry.conn.fd;
61            let user = entry.user.get_mut();
62            match user.session.fill_msghdr(CLIENT_BYTES_CAP) {
63                Some(m) => {
64                    msg_ptr = m as *const _;
65                    token_epoch = Some(entry.io.write.next_token::<B::Token>(external));
66                    close = false;
67                }
68                None => {
69                    msg_ptr = std::ptr::null();
70                    token_epoch = None;
71                    close = self.direction.protocol.wants_close(user.session.state());
72                }
73            }
74        }
75        if let Some((token, epoch)) = token_epoch {
76            if unsafe { driver.submit_send_msg_tagged(token, conn_fd, msg_ptr) } {
77                self.core.slots[slot_idx].io.write.mark_inflight(epoch);
78            }
79        } else if close {
80            self.disconnect_and_free(driver, slot_idx);
81        }
82    }
83
84    #[cold]
85    fn try_submit_outbound_codec(&mut self, driver: &mut B::Driver, slot_idx: usize) {
86        let cap = self.core.slots[slot_idx].buf.write_buf.as_slice().len();
87        let head = self.core.slots[slot_idx].conn.write_head as usize;
88        let room = cap.saturating_sub(head);
89
90        if room > 0 {
91            let CoreSlot {
92                conn, buf, user, ..
93            } = &mut self.core.slots[slot_idx];
94            let user = user.get_mut();
95            let req_cap = room.min(CLIENT_BYTES_CAP);
96            if let Some(msg) = user.session.fill_msghdr(req_cap) {
97                let iov_count = msg.msg_iovlen as usize;
98
99                let iov_slice: &[libc::iovec] =
100                    unsafe { std::slice::from_raw_parts(msg.msg_iov, iov_count) };
101                let mut copied = 0u32;
102                for iov in iov_slice {
103                    let len = iov.iov_len;
104                    if (head + copied as usize) + len > cap {
105                        break;
106                    }
107
108                    let src = unsafe { std::slice::from_raw_parts(iov.iov_base.cast::<u8>(), len) };
109                    let off = head + copied as usize;
110                    buf.write_buf.as_mut_slice()[off..off + len].copy_from_slice(src);
111                    copied += len as u32;
112                }
113                conn.write_head = (head + copied as usize).min(u16::MAX as usize) as u16;
114                user.session.ack_send(copied);
115            }
116        }
117
118        let close = {
119            let user = self.core.slots[slot_idx].user.get();
120            self.direction.protocol.wants_close(user.session.state())
121        };
122        if close {
123            self.core.slots[slot_idx]
124                .conn
125                .flags
126                .set(ConnFlags::CLOSE_AFTER, true);
127        }
128
129        self.core.encrypt_write_buf(slot_idx);
130        self.core.try_submit_codec(driver, slot_idx);
131    }
132
133    pub fn poke_outbound(
134        &mut self,
135        driver: &mut B::Driver,
136        conn_id: SlotId,
137    ) -> Result<(), PushError> {
138        let slot_idx = self.core.resolve_live_slot(conn_id)?;
139        self.try_submit_outbound(driver, slot_idx);
140        Ok(())
141    }
142
143    pub fn pump_outbound(&mut self, driver: &mut B::Driver) -> usize {
144        let mut count = 0usize;
145        for slot_idx in 0..self.core.slots.len() {
146            if !self.core.slots[slot_idx]
147                .conn
148                .flags
149                .contains(ConnFlags::LIVE)
150            {
151                continue;
152            }
153            self.try_submit_outbound(driver, slot_idx);
154            count += 1;
155        }
156        count
157    }
158
159    #[cold]
160    fn handle_write_event_codec(&mut self, driver: &mut B::Driver, slot_idx: usize, result: i32) {
161        match self.core.handle_codec_inflight(slot_idx, result) {
162            super::super::core::CodecAck::NoOp => {}
163            super::super::core::CodecAck::Err => self.disconnect_and_free(driver, slot_idx),
164            super::super::core::CodecAck::Progressed => {
165                self.core.try_submit_codec(driver, slot_idx);
166                self.try_submit_outbound(driver, slot_idx);
167                let close = self.core.slots[slot_idx]
168                    .conn
169                    .flags
170                    .contains(ConnFlags::CLOSE_AFTER)
171                    || self.core.codec_close_pending(slot_idx);
172                if self.core.codec_idle(slot_idx) && close {
173                    self.disconnect_and_free(driver, slot_idx);
174                }
175            }
176        }
177    }
178}
179
180impl<P, T, S, F, B> super::super::PoolWrite<B> for Pool<Client<P, T, S>, F, B>
181where
182    P: ClientProtocol,
183    T: Transport,
184    T::Addr: Clone,
185    S: ConnectSource<T>,
186    F: Profile,
187    B: PoolDriver,
188{
189    fn on_write_event(
190        &mut self,
191        driver: &mut B::Driver,
192        token: B::Token,
193        result: i32,
194        notif: bool,
195    ) {
196        let (slot_idx, result) = match self.core.classify_write(token, result, notif) {
197            super::super::core::WriteDispatch::Skipped => return,
198            super::super::core::WriteDispatch::Ready { slot_idx, result } => (slot_idx, result),
199        };
200        if <P::CodecLayer as CodecLayer>::IS_PASSTHROUGH {
201            if !self.core.slots[slot_idx]
202                .conn
203                .flags
204                .contains(ConnFlags::LIVE)
205            {
206                return;
207            }
208            if result < 0 {
209                self.disconnect_and_free(driver, slot_idx);
210                return;
211            }
212            let n = result as u32;
213            if n > 0 {
214                let user = self.core.slots[slot_idx].user.get_mut();
215                user.session.ack_send(n);
216            }
217            if self.core.slots[slot_idx]
218                .conn
219                .flags
220                .contains(ConnFlags::CLOSE_AFTER)
221            {
222                self.disconnect_and_free(driver, slot_idx);
223                return;
224            }
225            self.try_submit_outbound(driver, slot_idx);
226        } else {
227            self.handle_write_event_codec(driver, slot_idx, result);
228        }
229    }
230}