qbase/sid/
local_sid.rs

1use std::{
2    collections::VecDeque,
3    sync::{Arc, Mutex},
4    task::{Context, Poll, Waker},
5};
6
7use super::{Dir, Role, StreamId};
8use crate::{
9    frame::{
10        MaxStreamsFrame, StreamsBlockedFrame,
11        io::{ReceiveFrame, SendFrame},
12    },
13    net::tx::{ArcSendWakers, Signals},
14    sid::MAX_STREAMS_LIMIT,
15    varint::VarInt,
16};
17
18/// Local stream IDs management.
19#[derive(Debug)]
20struct LocalStreamIds<BLOCKED> {
21    /// Our role
22    role: Role,
23    max: [u64; 2],
24    unallocated: [u64; 2],
25    /// Used for waiting for the MaxStream frame notification from peer when we have exhausted the creation of stream IDs
26    wakers: [VecDeque<Waker>; 2],
27    /// The StreamsBlocked frames that will be sent to peer
28    blocked: BLOCKED,
29    tx_wakers: ArcSendWakers,
30}
31
32impl<BLOCKED> LocalStreamIds<BLOCKED>
33where
34    BLOCKED: SendFrame<StreamsBlockedFrame> + Clone + Send + 'static,
35{
36    /// Create a new [`LocalStreamIds`] with the given role,
37    /// and maximum number of streams that can be created in each [`Dir`].
38    fn new(
39        role: Role,
40        init_max_bi_streams: u64,
41        init_max_uni_streams: u64,
42        blocked: BLOCKED,
43        tx_wakers: ArcSendWakers,
44    ) -> Self {
45        debug_assert!(
46            role == Role::Client || (init_max_bi_streams == 0 && init_max_uni_streams == 0),
47            "Server cannot remember the parameters"
48        );
49        Self {
50            role,
51            max: [init_max_bi_streams, init_max_uni_streams],
52            unallocated: [0, 0],
53            wakers: [VecDeque::with_capacity(2), VecDeque::with_capacity(2)],
54            blocked,
55            tx_wakers,
56        }
57    }
58
59    /// Returns local role.
60    fn role(&self) -> Role {
61        self.role
62    }
63
64    /// Returns the number of opened streams in the `dir` direction.
65    fn opened_streams(&self, dir: Dir) -> u64 {
66        self.unallocated[dir as usize]
67    }
68
69    /// Receive the [`MaxStreamsFrame`](`crate::frame::MaxStreamsFrame`) from peer,
70    /// update the maximum stream ID that can be opened locally in the given direction.
71    fn recv_max_streams_frame(&mut self, frame: MaxStreamsFrame) {
72        let (dir, val) = match frame {
73            MaxStreamsFrame::Bi(max) => (Dir::Bi, max.into_u64()),
74            MaxStreamsFrame::Uni(max) => (Dir::Uni, max.into_u64()),
75        };
76        self.increase_limit(dir, val);
77    }
78
79    fn increase_limit(&mut self, dir: Dir, val: u64) {
80        assert!(val <= MAX_STREAMS_LIMIT);
81        let max_streams = &mut self.max[dir as usize];
82        // RFC9000: MAX_STREAMS frames that do not increase the stream limit MUST be ignored.
83        if *max_streams < val {
84            // The rejected 0rtt stream can be sent again, as if new data was written.
85            if *max_streams < self.unallocated[dir as usize] {
86                self.tx_wakers.wake_all_by(Signals::WRITTEN);
87            }
88            for waker in self.wakers[dir as usize].drain(..) {
89                waker.wake();
90            }
91            *max_streams = val;
92        }
93    }
94
95    fn poll_alloc_sid(&mut self, cx: &mut Context<'_>, dir: Dir) -> Poll<Option<StreamId>> {
96        let idx = dir as usize;
97        let max = self.max[idx];
98        let unallocated = self.unallocated[idx];
99        if unallocated > MAX_STREAMS_LIMIT {
100            Poll::Ready(None)
101        } else if unallocated < max {
102            self.unallocated[idx] += 1;
103            Poll::Ready(Some(StreamId::new(self.role, dir, unallocated)))
104        } else {
105            // waiting for MAX_STREAMS frame from peer
106            self.wakers[idx].push_back(cx.waker().clone());
107            // if Poll::Pending is returned, connection can send a STREAMS_BLOCKED frame to peer
108            self.blocked.send_frame([StreamsBlockedFrame::with(
109                dir,
110                VarInt::from_u64(max).expect("max_streams limit must be less than VARINT_MAX"),
111            )]);
112            Poll::Pending
113        }
114    }
115
116    pub fn revise_max_streams(
117        &mut self,
118        zero_rtt_rejected: bool,
119        max_stream_bidi: u64,
120        max_stream_uni: u64,
121    ) {
122        if zero_rtt_rejected {
123            self.max = [0, 0];
124        }
125        self.increase_limit(Dir::Bi, max_stream_bidi);
126        self.increase_limit(Dir::Uni, max_stream_uni);
127    }
128}
129
130/// Management of stream IDs that can ben allowed to use locally.
131///
132/// The maximum stream ID that can be created is limited by the
133/// [`MaxStreamsFrame`](`crate::frame::MaxStreamsFrame`) from the peer.
134///
135/// When the stream IDs in the `dir` direction are exhausted,
136/// a [`StreamsBlockedFrame`](`crate::frame::StreamsBlockedFrame`) will be sent to the peer.
137/// The generic parameter `BLOCKED` is the container of the [`StreamsBlockedFrame`]
138/// that will be sent to peer, it can be a channel, a queue, or a buffer,
139/// as long as it can send the [`StreamsBlockedFrame`] to peer.
140#[derive(Debug, Clone)]
141pub struct ArcLocalStreamIds<BLOCKED>(Arc<Mutex<LocalStreamIds<BLOCKED>>>);
142
143impl<BLOCKED> ArcLocalStreamIds<BLOCKED>
144where
145    BLOCKED: SendFrame<StreamsBlockedFrame> + Clone + Send + 'static,
146{
147    /// Create a new [`ArcLocalStreamIds`] with the given role,
148    /// and maximum number of streams that can be created in each direction,
149    /// the `blocked` contains the [`StreamsBlockedFrame`] that will be sent to peer.
150    pub fn new(
151        role: Role,
152        max_bidi: u64,
153        max_uni: u64,
154        blocked: BLOCKED,
155        tx_wakers: ArcSendWakers,
156    ) -> Self {
157        Self(Arc::new(Mutex::new(LocalStreamIds::new(
158            role, max_bidi, max_uni, blocked, tx_wakers,
159        ))))
160    }
161
162    /// Returns local role
163    pub fn role(&self) -> Role {
164        self.0.lock().unwrap().role()
165    }
166
167    /// Returns the number of opened streams in the `dir` direction.
168    ///
169    /// If `is_0rtt` is true, this will return the stream open in 0rtt phase.
170    ///
171    /// If `is_0rtt` is false, the return value will not be greater than max_streams,
172    /// that is, if 0rtt is rejected, the return value may be less than the number of open streams.
173    /// This is the number of streams that can actually be sent in the 1rtt space.
174    pub fn opened_streams(&self, dir: Dir) -> u64 {
175        self.0.lock().unwrap().opened_streams(dir)
176    }
177
178    /// Receive the [`MaxStreamsFrame`](`crate::frame::MaxStreamsFrame`) from peer,
179    /// and then update the maximum stream ID that can be allowed to use locally.
180    ///
181    /// The maximum stream ID that can be allowed to use is limited by peer.
182    /// Therefore, it mainly depends on the peer's attitude
183    /// and is subject to the [`MaxStreamsFrame`](`crate::frame::MaxStreamsFrame`)
184    /// received from peer.
185    pub fn recv_max_streams_frame(&self, frame: MaxStreamsFrame) {
186        self.0.lock().unwrap().recv_max_streams_frame(frame);
187    }
188
189    /// Asynchronously allocate the next new [`StreamId`] in the `dir` direction.
190    ///
191    /// Return a bool indicating whether the stream is opened in 0rtt phase.
192    ///
193    /// When the application layer wants to proactively open a new stream,
194    /// it needs to first apply to allocate the next unused [`StreamId`].
195    /// Note that streams on a QUIC connection usually have a maximum concurrency limit,
196    /// so when requesting a [`StreamId`], it may not be possible to obtain one due to
197    /// reaching the maximum concurrency limit.
198    /// However, this is temporary. When the active current streams end,
199    /// the peer will expand the maximum stream ID limit through a
200    /// [`MaxStreamsFrame`](`crate::frame::MaxStreamsFrame`),
201    /// allowing the allocation of the [`StreamId`] meanwhile.
202    ///
203    /// Return Pending when the stream IDs in the `dir` direction are exhausted,
204    /// until receiving the [`MaxStreamsFrame`](`crate::frame::MaxStreamsFrame`) from peer.
205    ///
206    /// Return None if the stream IDs in the `dir` direction finally exceed 2^60,
207    /// but it is very very hard to happen.
208    pub fn poll_alloc_sid(&self, cx: &mut Context<'_>, dir: Dir) -> Poll<Option<StreamId>> {
209        self.0.lock().unwrap().poll_alloc_sid(cx, dir)
210    }
211
212    pub fn revise_max_streams(
213        &self,
214        zero_rtt_rejected: bool,
215        max_stream_bidi: u64,
216        max_stream_uni: u64,
217    ) {
218        self.0.lock().unwrap().revise_max_streams(
219            zero_rtt_rejected,
220            max_stream_bidi,
221            max_stream_uni,
222        );
223    }
224}
225
226impl<BLOCKED> ReceiveFrame<MaxStreamsFrame> for ArcLocalStreamIds<BLOCKED>
227where
228    BLOCKED: SendFrame<StreamsBlockedFrame> + Clone + Send + 'static,
229{
230    type Output = ();
231
232    fn recv_frame(&self, frame: MaxStreamsFrame) -> Result<Self::Output, crate::error::Error> {
233        self.recv_max_streams_frame(frame);
234        Ok(())
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use derive_more::Deref;
241
242    use super::*;
243    use crate::util::ArcAsyncDeque;
244
245    #[derive(Clone, Deref, Default)]
246    struct StreamsBlockedFrameTx(ArcAsyncDeque<StreamsBlockedFrame>);
247
248    impl SendFrame<StreamsBlockedFrame> for StreamsBlockedFrameTx {
249        fn send_frame<I: IntoIterator<Item = StreamsBlockedFrame>>(&self, iter: I) {
250            (&self.0).extend(iter);
251        }
252    }
253
254    #[test]
255    fn test_stream_id_new() {
256        let sid = StreamId::new(Role::Client, Dir::Bi, 0);
257        assert_eq!(sid, StreamId(0));
258        assert_eq!(sid.role(), Role::Client);
259        assert_eq!(sid.dir(), Dir::Bi);
260    }
261
262    #[test]
263    fn test_recv_max_stream_frames() {
264        let local = ArcLocalStreamIds::new(
265            Role::Client,
266            0,
267            0,
268            StreamsBlockedFrameTx::default(),
269            ArcSendWakers::default(),
270        );
271        local.recv_max_streams_frame(MaxStreamsFrame::Bi(VarInt::from_u32(0)));
272        let waker = futures::task::noop_waker();
273        let mut cx = Context::from_waker(&waker);
274        assert_eq!(local.poll_alloc_sid(&mut cx, Dir::Bi), Poll::Pending,);
275        assert!(!local.0.lock().unwrap().wakers[0].is_empty());
276
277        local.recv_max_streams_frame(MaxStreamsFrame::Bi(VarInt::from_u32(1)));
278        let _ = local.0.lock().unwrap().wakers[0].pop_front();
279        assert_eq!(
280            local.poll_alloc_sid(&mut cx, Dir::Bi),
281            Poll::Ready(Some(StreamId(0)))
282        );
283        assert_eq!(local.poll_alloc_sid(&mut cx, Dir::Bi), Poll::Pending);
284        assert!(!local.0.lock().unwrap().wakers[0].is_empty());
285
286        local.recv_max_streams_frame(MaxStreamsFrame::Uni(VarInt::from_u32(2)));
287        assert_eq!(
288            local.poll_alloc_sid(&mut cx, Dir::Uni),
289            Poll::Ready(Some(StreamId(2)))
290        );
291        assert_eq!(
292            local.poll_alloc_sid(&mut cx, Dir::Uni),
293            Poll::Ready(Some(StreamId(6)))
294        );
295        assert_eq!(local.poll_alloc_sid(&mut cx, Dir::Uni), Poll::Pending);
296        assert!(!local.0.lock().unwrap().wakers[1].is_empty());
297    }
298}
qbase/sid/local_sid.rs

qbase/sid/
local_sid.rs
