dope_session/pool/server/
write.rs

1use super::super::Pool;
2use super::super::core::{CodecAck, PushError, WriteDispatch};
3use super::super::slot::{ConnFlags, CoreSlot};
4use super::Server;
5use super::state::SlotEgress;
6use crate::CodecLayer;
7use crate::SlotId;
8use crate::WriteBufStorage;
9use crate::protocol::server::ServerProtocol;
10use dope_core::IoVec;
11use dope_core::driver::{DriverOps, PoolDriver};
12use dope_core::profile::Profile;
13use dope_transport::Transport;
14
15enum WriteEventOutcome {
16    Idle,
17    More,
18    Close,
19}
20
21impl<H: ServerProtocol, T: Transport, F: Profile, B: PoolDriver> super::super::PoolWrite<B>
22    for Pool<Server<H, T>, F, B>
23{
24    fn on_write_event(
25        &mut self,
26        driver: &mut B::Driver,
27        token: B::Token,
28        result: i32,
29        notif: bool,
30    ) {
31        let (slot_idx, result) = match self.core.classify_write(token, result, notif) {
32            WriteDispatch::Skipped => return,
33            WriteDispatch::Ready { slot_idx, result } => (slot_idx, result),
34        };
35        if <H::CodecLayer as CodecLayer>::IS_PASSTHROUGH {
36            if !self.core.slots[slot_idx]
37                .conn
38                .flags
39                .contains(ConnFlags::LIVE)
40            {
41                return;
42            }
43            if result < 0 {
44                self.core.free_slot(driver, slot_idx);
45                return;
46            }
47            let n = result as u32;
48            match self.advance_write_state(slot_idx, n) {
49                WriteEventOutcome::Close => self.core.free_slot(driver, slot_idx),
50                WriteEventOutcome::More => self.submit_slot_write(driver, slot_idx),
51                WriteEventOutcome::Idle => {}
52            }
53        } else {
54            self.handle_write_event_codec(driver, slot_idx, result);
55        }
56    }
57}
58
59impl<H: ServerProtocol, T: Transport, F: Profile, B: PoolDriver> Pool<Server<H, T>, F, B> {
60    fn handle_write_event_codec(&mut self, driver: &mut B::Driver, slot_idx: usize, result: i32) {
61        match self.core.handle_codec_inflight(slot_idx, result) {
62            CodecAck::NoOp => {}
63            CodecAck::Err => self.core.free_slot(driver, slot_idx),
64            CodecAck::Progressed => {
65                self.submit_slot_write(driver, slot_idx);
66                let close = self.core.slots[slot_idx]
67                    .conn
68                    .flags
69                    .contains(ConnFlags::CLOSE_AFTER)
70                    || self.core.codec_close_pending(slot_idx);
71                if self.core.codec_idle(slot_idx) && close {
72                    self.core.free_slot(driver, slot_idx);
73                }
74            }
75        }
76    }
77
78    #[inline(always)]
79    fn advance_write_state(&mut self, slot_idx: usize, n: u32) -> WriteEventOutcome {
80        let CoreSlot { conn, user, .. } = &mut self.core.slots[slot_idx];
81        let user = user.get_mut();
82
83        if user.egress.is_stream() {
84            if conn.write_tail < conn.write_inflight_end {
85                let new_tail = (conn.write_tail as u32).saturating_add(n);
86                conn.write_tail = new_tail.min(u16::MAX as u32) as u16;
87                if conn.write_tail >= conn.write_inflight_end {
88                    conn.write_head = 0;
89                    conn.write_tail = 0;
90                    conn.write_inflight_end = 0;
91                    return WriteEventOutcome::More;
92                }
93                return WriteEventOutcome::Idle;
94            }
95            let SlotEgress::Stream(stream_state) = &mut user.egress else {
96                return WriteEventOutcome::Close;
97            };
98            stream_state.advance(n);
99            if stream_state.eof && stream_state.cur_chunk.is_none() {
100                user.egress = SlotEgress::Idle;
101                if conn.flags.contains(ConnFlags::CLOSE_AFTER) {
102                    return WriteEventOutcome::Close;
103                }
104                return WriteEventOutcome::Idle;
105            }
106            return WriteEventOutcome::More;
107        }
108
109        if !conn.pending_body_ptr.is_null() {
110            let new_tail = conn.large_tail.saturating_add(n);
111            conn.large_tail = new_tail;
112            if conn.large_tail >= conn.large_inflight {
113                conn.write_head = 0;
114                conn.write_tail = 0;
115                conn.write_inflight_end = 0;
116                conn.pending_body_ptr = std::ptr::null();
117                conn.pending_body_len = 0;
118                conn.large_inflight = 0;
119                conn.large_tail = 0;
120                if conn.flags.contains(ConnFlags::CLOSE_AFTER) {
121                    return WriteEventOutcome::Close;
122                }
123            }
124            return WriteEventOutcome::Idle;
125        }
126
127        if conn.large_inflight > 0 {
128            let new_tail = conn.large_tail.saturating_add(n);
129            conn.large_tail = new_tail;
130            if conn.large_tail >= conn.large_inflight {
131                conn.large_inflight = 0;
132                conn.large_tail = 0;
133                if conn.flags.contains(ConnFlags::CLOSE_AFTER) {
134                    return WriteEventOutcome::Close;
135                }
136                if user.egress.gather_iov_count().is_some() {
137                    return WriteEventOutcome::More;
138                }
139            }
140            return WriteEventOutcome::Idle;
141        }
142
143        let new_tail = (conn.write_tail as u32).saturating_add(n);
144        conn.write_tail = new_tail.min(u16::MAX as u32) as u16;
145        if conn.write_tail >= conn.write_inflight_end {
146            conn.write_head = 0;
147            conn.write_tail = 0;
148            conn.write_inflight_end = 0;
149            if conn.flags.contains(ConnFlags::CLOSE_AFTER) {
150                return WriteEventOutcome::Close;
151            }
152        }
153        WriteEventOutcome::Idle
154    }
155
156    #[inline(always)]
157    fn try_submit_gather(&mut self, driver: &mut B::Driver, slot_idx: usize) -> bool {
158        let Some(iov_count) = self.core.slots[slot_idx]
159            .user
160            .get()
161            .egress
162            .gather_iov_count()
163        else {
164            return false;
165        };
166        if self.core.slots[slot_idx].io.write.has_in_flight {
167            return true;
168        }
169        let entry = &mut self.core.slots[slot_idx];
170        entry.buf.pending_msghdr.msg_iovlen = iov_count as _;
171        let msg_ptr: *const libc::msghdr = &entry.buf.pending_msghdr;
172        let submitted = unsafe { self.driver_submit_send_msg(driver, slot_idx, msg_ptr) };
173        if submitted {
174            self.core.slots[slot_idx].user.get_mut().egress = SlotEgress::Idle;
175        }
176        true
177    }
178
179    pub fn push_response(
180        &mut self,
181        driver: &mut B::Driver,
182        conn_id: SlotId,
183        bytes: &[u8],
184        close_after: bool,
185    ) -> Result<(), PushError> {
186        debug_assert!(
187            H::SUPPORTS_PARK,
188            "push_response requires ServerProtocol::SUPPORTS_PARK = true",
189        );
190        let slot_idx = self.core.resolve_live_slot(conn_id)?;
191        let entry = &mut self.core.slots[slot_idx];
192        let head = entry.conn.write_head as usize;
193        let room = entry.buf.write_buf.as_slice().len().saturating_sub(head);
194        if bytes.len() > room {
195            return Err(PushError::BufferFull);
196        }
197        entry.buf.write_buf.as_mut_slice()[head..head + bytes.len()].copy_from_slice(bytes);
198        entry.conn.write_head = (head + bytes.len()) as u16;
199        if close_after {
200            entry.conn.flags.set(ConnFlags::CLOSE_AFTER, true);
201        }
202        entry.conn.flags.set(ConnFlags::PARKED, false);
203        self.submit_slot_write(driver, slot_idx);
204        Ok(())
205    }
206
207    #[inline(always)]
208    pub(super) fn submit_slot_write(&mut self, driver: &mut B::Driver, slot_idx: usize) {
209        if !<H::CodecLayer as CodecLayer>::IS_PASSTHROUGH {
210            return self.submit_slot_write_codec(driver, slot_idx);
211        }
212        if self.try_submit_gather(driver, slot_idx) {
213            return;
214        }
215        self.submit_slot_write_plain(driver, slot_idx);
216    }
217
218    fn submit_slot_write_plain(&mut self, driver: &mut B::Driver, slot_idx: usize) {
219        if self.core.slots[slot_idx].user.get().egress.is_stream()
220            && self.core.slots[slot_idx].conn.pending_body_ptr.is_null()
221        {
222            self.encode_stream_chunks_to_buf(slot_idx);
223        }
224
225        let stream_in_progress = self.core.slots[slot_idx].user.get().egress.is_stream();
226        let CoreSlot { conn, buf, io, .. } = &mut self.core.slots[slot_idx];
227        let head = conn.write_head as usize;
228        let inflight_end = conn.write_inflight_end as usize;
229        if head <= inflight_end {
230            if stream_in_progress {
231                return self.submit_stream_write(driver, slot_idx);
232            }
233            return;
234        }
235        if io.write.has_in_flight {
236            return;
237        }
238
239        if !conn.pending_body_ptr.is_null() {
240            let hdr_ptr = buf.write_buf.as_mut_slice()[inflight_end..head].as_ptr();
241            let hdr_len = head - inflight_end;
242            let body_ptr = conn.pending_body_ptr;
243            let body_len = conn.pending_body_len as usize;
244            buf.pending_iovs[0] = IoVec::raw(hdr_ptr, hdr_len);
245            buf.pending_iovs[1] = IoVec::raw(body_ptr, body_len);
246            buf.pending_msghdr.msg_iovlen = 2;
247            let msg_ptr: *const libc::msghdr = &buf.pending_msghdr;
248            conn.large_inflight = (hdr_len as u32).saturating_add(body_len as u32);
249            conn.large_tail = 0;
250            conn.write_inflight_end = head as u16;
251            let _ = unsafe { self.driver_submit_send_msg(driver, slot_idx, msg_ptr) };
252        } else {
253            let bytes_ptr = buf.write_buf.as_mut_slice()[inflight_end..head].as_ptr();
254            let bytes_len = (head - inflight_end) as u32;
255            let _ = self.driver_submit_send(driver, slot_idx, bytes_ptr, bytes_len);
256            self.core.slots[slot_idx].conn.write_inflight_end = head as u16;
257        }
258    }
259
260    #[inline(always)]
261    pub(super) fn driver_submit_send(
262        &mut self,
263        driver: &mut B::Driver,
264        slot_idx: usize,
265        ptr: *const u8,
266        len: u32,
267    ) -> bool {
268        let external = self.core.slot_id_base + slot_idx;
269        let conn_fd = self.core.slots[slot_idx].conn.fd;
270        let write_slot = &mut self.core.slots[slot_idx].io.write;
271        let (token, epoch) = write_slot.next_token::<B::Token>(external);
272        if driver.submit_send_tagged(token, conn_fd, ptr, len) {
273            self.core.slots[slot_idx].io.write.mark_inflight(epoch);
274            true
275        } else {
276            false
277        }
278    }
279
280    #[inline(always)]
281    pub(super) unsafe fn driver_submit_send_msg(
282        &mut self,
283        driver: &mut B::Driver,
284        slot_idx: usize,
285        msg: *const libc::msghdr,
286    ) -> bool {
287        let external = self.core.slot_id_base + slot_idx;
288        let conn_fd = self.core.slots[slot_idx].conn.fd;
289        let write_slot = &mut self.core.slots[slot_idx].io.write;
290        let (token, epoch) = write_slot.next_token::<B::Token>(external);
291        if unsafe { driver.submit_send_msg_tagged(token, conn_fd, msg) } {
292            self.core.slots[slot_idx].io.write.mark_inflight(epoch);
293            true
294        } else {
295            false
296        }
297    }
298
299    fn submit_slot_write_codec(&mut self, driver: &mut B::Driver, slot_idx: usize) {
300        let stream_in_progress = self.core.slots[slot_idx].user.get().egress.is_stream();
301        if stream_in_progress && self.core.slots[slot_idx].conn.pending_body_ptr.is_null() {
302            self.encode_stream_chunks_to_buf(slot_idx);
303        }
304
305        self.core.encrypt_write_buf(slot_idx);
306
307        if !self.core.slots[slot_idx].conn.pending_body_ptr.is_null() {
308            let body_ptr = self.core.slots[slot_idx].conn.pending_body_ptr;
309            let body_len = self.core.slots[slot_idx].conn.pending_body_len as usize;
310
311            let body = unsafe { std::slice::from_raw_parts(body_ptr, body_len) };
312            <H::CodecLayer as CodecLayer>::process_outbound(
313                self.core.slots[slot_idx].codec_state.get_mut(),
314                body,
315            );
316            self.core.slots[slot_idx].conn.pending_body_ptr = std::ptr::null();
317            self.core.slots[slot_idx].conn.pending_body_len = 0;
318            self.core.slots[slot_idx].conn.large_inflight = 0;
319            self.core.slots[slot_idx].conn.large_tail = 0;
320        }
321
322        self.core.try_submit_codec(driver, slot_idx);
323    }
324}
dope_session/pool/server/write.rs

dope_session/pool/server/
write.rs
