Skip to main content

s2n_quic_dc/stream/send/
shared.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{
5    stream::{
6        packet_number,
7        send::{
8            application::transmission, buffer, error::Error, flow, path, queue::Queue,
9            state::Transmission,
10        },
11        shared::ShutdownKind,
12    },
13    task::waker::worker::Waker as WorkerWaker,
14};
15use core::{
16    fmt,
17    sync::atomic::{AtomicU64, Ordering},
18};
19use crossbeam_queue::SegQueue;
20use s2n_quic_core::recovery::bandwidth::Bandwidth;
21use tracing::trace;
22
23#[derive(Debug)]
24pub struct Message {
25    /// The event being submitted to the worker
26    pub event: Event,
27}
28
29#[derive(Debug)]
30pub enum Event {
31    Shutdown { queue: Queue, kind: ShutdownKind },
32}
33
34pub struct State {
35    pub flow: flow::non_blocking::State,
36    pub packet_number: packet_number::Counter,
37    pub path: path::State,
38    pub worker_waker: WorkerWaker,
39    bandwidth: AtomicU64,
40    /// A channel sender for pushing transmission information to the worker task
41    ///
42    /// We use an unbounded sender since we already rely on flow control to apply backpressure
43    worker_queue: SegQueue<Message>,
44    pub application_transmission_queue: transmission::Queue<buffer::Segment>,
45    pub segment_alloc: buffer::Allocator,
46}
47
48impl fmt::Debug for State {
49    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
50        f.debug_struct("send::shared::State")
51            .field("flow", &self.flow)
52            .field("packet_number", &self.packet_number)
53            .field("path", &self.path)
54            .finish()
55    }
56}
57
58impl State {
59    #[inline]
60    pub fn new(
61        flow: flow::non_blocking::State,
62        path: path::Info,
63        bandwidth: Option<Bandwidth>,
64    ) -> Self {
65        let path = path::State::new(path);
66        let bandwidth = bandwidth.map(|v| v.serialize()).unwrap_or(u64::MAX).into();
67        Self {
68            flow,
69            packet_number: Default::default(),
70            path,
71            bandwidth,
72            // this will get set once the waker spawns
73            worker_waker: Default::default(),
74            worker_queue: Default::default(),
75            application_transmission_queue: Default::default(),
76            segment_alloc: Default::default(),
77        }
78    }
79
80    #[inline]
81    pub fn bandwidth(&self) -> Bandwidth {
82        Bandwidth::deserialize(self.bandwidth.load(Ordering::Relaxed))
83    }
84
85    #[inline]
86    pub fn set_bandwidth(&self, value: Bandwidth) {
87        self.bandwidth.store(value.serialize(), Ordering::Relaxed);
88    }
89
90    #[inline]
91    pub fn pop_worker_message(&self) -> Option<Message> {
92        self.worker_queue.pop()
93    }
94
95    #[inline]
96    pub fn push_to_worker(&self, transmissions: Vec<Transmission>) -> Result<(), Error> {
97        trace!(event = "transmission", len = transmissions.len());
98        self.application_transmission_queue
99            .push_batch(transmissions);
100
101        self.worker_waker.wake();
102
103        Ok(())
104    }
105
106    pub fn on_prune(&self) {
107        self.shutdown(Default::default(), ShutdownKind::Pruned);
108    }
109
110    #[inline]
111    pub fn shutdown(&self, queue: Queue, kind: ShutdownKind) {
112        trace!(event = "shutdown", queue = queue.accepted_len(), ?kind);
113        let message = Message {
114            event: Event::Shutdown { queue, kind },
115        };
116        self.worker_queue.push(message);
117        self.worker_waker.wake();
118    }
119}