compio_quic/
send_stream.rs

1use std::{
2    io,
3    task::{Context, Poll},
4};
5
6use compio_buf::{BufResult, IoBuf, bytes::Bytes};
7use compio_io::AsyncWrite;
8use futures_util::{future::poll_fn, ready};
9use quinn_proto::{ClosedStream, FinishError, StreamId, VarInt, Written};
10use thiserror::Error;
11
12use crate::{ConnectionError, ConnectionInner, sync::shared::Shared};
13
14/// A stream that can only be used to send data.
15///
16/// If dropped, streams that haven't been explicitly [`reset()`] will be
17/// implicitly [`finish()`]ed, continuing to (re)transmit previously written
18/// data until it has been fully acknowledged or the connection is closed.
19///
20/// # Cancellation
21///
22/// A `write` method is said to be *cancel-safe* when dropping its future before
23/// the future becomes ready will always result in no data being written to the
24/// stream. This is true of methods which succeed immediately when any progress
25/// is made, and is not true of methods which might need to perform multiple
26/// writes internally before succeeding. Each `write` method documents whether
27/// it is cancel-safe.
28///
29/// [`reset()`]: SendStream::reset
30/// [`finish()`]: SendStream::finish
31#[derive(Debug)]
32pub struct SendStream {
33    conn: Shared<ConnectionInner>,
34    stream: StreamId,
35    is_0rtt: bool,
36}
37
38impl SendStream {
39    pub(crate) fn new(conn: Shared<ConnectionInner>, stream: StreamId, is_0rtt: bool) -> Self {
40        Self {
41            conn,
42            stream,
43            is_0rtt,
44        }
45    }
46
47    /// Get the identity of this stream
48    pub fn id(&self) -> StreamId {
49        self.stream
50    }
51
52    /// Notify the peer that no more data will ever be written to this stream.
53    ///
54    /// It is an error to write to a stream after `finish()`ing it. [`reset()`]
55    /// may still be called after `finish` to abandon transmission of any stream
56    /// data that might still be buffered.
57    ///
58    /// To wait for the peer to receive all buffered stream data, see
59    /// [`stopped()`].
60    ///
61    /// May fail if [`finish()`] or  [`reset()`] was previously called.This
62    /// error is harmless and serves only to indicate that the caller may have
63    /// incorrect assumptions about the stream's state.
64    ///
65    /// [`reset()`]: Self::reset
66    /// [`stopped()`]: Self::stopped
67    /// [`finish()`]: Self::finish
68    pub fn finish(&mut self) -> Result<(), ClosedStream> {
69        let mut state = self.conn.state();
70        match state.conn.send_stream(self.stream).finish() {
71            Ok(()) => {
72                state.wake();
73                Ok(())
74            }
75            Err(FinishError::ClosedStream) => Err(ClosedStream::default()),
76            // Harmless. If the application needs to know about stopped streams
77            // at this point, it should call `stopped`.
78            Err(FinishError::Stopped(_)) => Ok(()),
79        }
80    }
81
82    /// Close the stream immediately.
83    ///
84    /// No new data can be written after calling this method. Locally buffered
85    /// data is dropped, and previously transmitted data will no longer be
86    /// retransmitted if lost. If an attempt has already been made to finish
87    /// the stream, the peer may still receive all written data.
88    ///
89    /// May fail if [`finish()`](Self::finish) or [`reset()`](Self::reset) was
90    /// previously called. This error is harmless and serves only to
91    /// indicate that the caller may have incorrect assumptions about the
92    /// stream's state.
93    pub fn reset(&mut self, error_code: VarInt) -> Result<(), ClosedStream> {
94        let mut state = self.conn.state();
95        if self.is_0rtt && !state.check_0rtt() {
96            return Ok(());
97        }
98        state.conn.send_stream(self.stream).reset(error_code)?;
99        state.wake();
100        Ok(())
101    }
102
103    /// Set the priority of the stream.
104    ///
105    /// Every stream has an initial priority of 0. Locally buffered data
106    /// from streams with higher priority will be transmitted before data
107    /// from streams with lower priority. Changing the priority of a stream
108    /// with pending data may only take effect after that data has been
109    /// transmitted. Using many different priority levels per connection may
110    /// have a negative impact on performance.
111    pub fn set_priority(&self, priority: i32) -> Result<(), ClosedStream> {
112        self.conn
113            .state()
114            .conn
115            .send_stream(self.stream)
116            .set_priority(priority)
117    }
118
119    /// Get the priority of the stream
120    pub fn priority(&self) -> Result<i32, ClosedStream> {
121        self.conn.state().conn.send_stream(self.stream).priority()
122    }
123
124    /// Completes when the peer stops the stream or reads the stream to
125    /// completion.
126    ///
127    /// Yields `Some` with the stop error code if the peer stops the stream.
128    /// Yields `None` if the local side [`finish()`](Self::finish)es the stream
129    /// and then the peer acknowledges receipt of all stream data (although not
130    /// necessarily the processing of it), after which the peer closing the
131    /// stream is no longer meaningful.
132    ///
133    /// For a variety of reasons, the peer may not send acknowledgements
134    /// immediately upon receiving data. As such, relying on `stopped` to
135    /// know when the peer has read a stream to completion may introduce
136    /// more latency than using an application-level response of some sort.
137    pub async fn stopped(&mut self) -> Result<Option<VarInt>, StoppedError> {
138        poll_fn(|cx| self.poll_stopped(cx)).await
139    }
140
141    /// Polls for the stream to be stopped or read to completion; see
142    /// [`stopped()`](Self::stopped).
143    pub fn poll_stopped(&mut self, cx: &mut Context) -> Poll<Result<Option<VarInt>, StoppedError>> {
144        let mut state = self.conn.state();
145        if self.is_0rtt && !state.check_0rtt() {
146            return Poll::Ready(Err(StoppedError::ZeroRttRejected));
147        }
148        match state.conn.send_stream(self.stream).stopped() {
149            Err(_) => Poll::Ready(Ok(None)),
150            Ok(Some(error_code)) => Poll::Ready(Ok(Some(error_code))),
151            Ok(None) => {
152                if let Some(e) = &state.error {
153                    return Poll::Ready(Err(e.clone().into()));
154                }
155                state.stopped.insert(self.stream, cx.waker().clone());
156                Poll::Pending
157            }
158        }
159    }
160
161    fn execute_poll_write<F, R>(&mut self, cx: &mut Context, f: F) -> Poll<Result<R, WriteError>>
162    where
163        F: FnOnce(quinn_proto::SendStream) -> Result<R, quinn_proto::WriteError>,
164    {
165        let mut state = self.conn.try_state()?;
166        if self.is_0rtt && !state.check_0rtt() {
167            return Poll::Ready(Err(WriteError::ZeroRttRejected));
168        }
169        match f(state.conn.send_stream(self.stream)) {
170            Ok(r) => {
171                state.wake();
172                Poll::Ready(Ok(r))
173            }
174            Err(e) => match e.try_into() {
175                Ok(e) => Poll::Ready(Err(e)),
176                Err(()) => {
177                    state.writable.insert(self.stream, cx.waker().clone());
178                    Poll::Pending
179                }
180            },
181        }
182    }
183
184    /// Write chunks to the stream.
185    ///
186    /// Yields the number of bytes and chunks written on success.
187    /// Congestion and flow control may cause this to be shorter than
188    /// `buf.len()`, indicating that only a prefix of `bufs` was written.
189    ///
190    /// This operation is cancel-safe.
191    pub async fn write_chunks(&mut self, bufs: &mut [Bytes]) -> Result<Written, WriteError> {
192        poll_fn(|cx| self.execute_poll_write(cx, |mut stream| stream.write_chunks(bufs))).await
193    }
194
195    /// Convenience method to write an entire list of chunks to the stream.
196    ///
197    /// This operation is *not* cancel-safe.
198    pub async fn write_all_chunks(&mut self, bufs: &mut [Bytes]) -> Result<(), WriteError> {
199        let mut chunks = 0;
200        poll_fn(|cx| {
201            loop {
202                if chunks == bufs.len() {
203                    return Poll::Ready(Ok(()));
204                }
205                let written = ready!(self.execute_poll_write(cx, |mut stream| {
206                    stream.write_chunks(&mut bufs[chunks..])
207                }))?;
208                chunks += written.chunks;
209            }
210        })
211        .await
212    }
213
214    /// Convert this stream into a [`futures_util`] compatible stream.
215    #[cfg(feature = "io-compat")]
216    pub fn into_compat(self) -> CompatSendStream {
217        CompatSendStream(self)
218    }
219}
220
221impl Drop for SendStream {
222    fn drop(&mut self) {
223        let mut state = self.conn.state();
224
225        // clean up any previously registered wakers
226        state.stopped.remove(&self.stream);
227        state.writable.remove(&self.stream);
228
229        if state.error.is_some() || (self.is_0rtt && !state.check_0rtt()) {
230            return;
231        }
232        match state.conn.send_stream(self.stream).finish() {
233            Ok(()) => state.wake(),
234            Err(FinishError::Stopped(reason)) => {
235                if state.conn.send_stream(self.stream).reset(reason).is_ok() {
236                    state.wake();
237                }
238            }
239            // Already finished or reset, which is fine.
240            Err(FinishError::ClosedStream) => {}
241        }
242    }
243}
244
245/// Errors that arise from writing to a stream
246#[derive(Debug, Error, Clone, PartialEq, Eq)]
247pub enum WriteError {
248    /// The peer is no longer accepting data on this stream
249    ///
250    /// Carries an application-defined error code.
251    #[error("sending stopped by peer: error {0}")]
252    Stopped(VarInt),
253    /// The connection was lost
254    #[error("connection lost")]
255    ConnectionLost(#[from] ConnectionError),
256    /// The stream has already been finished or reset
257    #[error("closed stream")]
258    ClosedStream,
259    /// This was a 0-RTT stream and the server rejected it
260    ///
261    /// Can only occur on clients for 0-RTT streams, which can be opened using
262    /// [`Connecting::into_0rtt()`].
263    ///
264    /// [`Connecting::into_0rtt()`]: crate::Connecting::into_0rtt()
265    #[error("0-RTT rejected")]
266    ZeroRttRejected,
267    /// Error when the stream is not ready, because it is still sending
268    /// data from a previous call
269    #[cfg(feature = "h3")]
270    #[error("stream not ready")]
271    NotReady,
272}
273
274impl TryFrom<quinn_proto::WriteError> for WriteError {
275    type Error = ();
276
277    fn try_from(value: quinn_proto::WriteError) -> Result<Self, Self::Error> {
278        use quinn_proto::WriteError::*;
279        match value {
280            Stopped(e) => Ok(Self::Stopped(e)),
281            ClosedStream => Ok(Self::ClosedStream),
282            Blocked => Err(()),
283        }
284    }
285}
286
287impl From<StoppedError> for WriteError {
288    fn from(x: StoppedError) -> Self {
289        match x {
290            StoppedError::ConnectionLost(e) => Self::ConnectionLost(e),
291            StoppedError::ZeroRttRejected => Self::ZeroRttRejected,
292        }
293    }
294}
295
296impl From<WriteError> for io::Error {
297    fn from(x: WriteError) -> Self {
298        use WriteError::*;
299        let kind = match x {
300            Stopped(_) | ZeroRttRejected => io::ErrorKind::ConnectionReset,
301            ConnectionLost(_) | ClosedStream => io::ErrorKind::NotConnected,
302            #[cfg(feature = "h3")]
303            NotReady => io::ErrorKind::Other,
304        };
305        Self::new(kind, x)
306    }
307}
308
309/// Errors that arise while monitoring for a send stream stop from the peer
310#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
311pub enum StoppedError {
312    /// The connection was lost
313    #[error("connection lost")]
314    ConnectionLost(#[from] ConnectionError),
315    /// This was a 0-RTT stream and the server rejected it
316    ///
317    /// Can only occur on clients for 0-RTT streams, which can be opened using
318    /// [`Connecting::into_0rtt()`].
319    ///
320    /// [`Connecting::into_0rtt()`]: crate::Connecting::into_0rtt()
321    #[error("0-RTT rejected")]
322    ZeroRttRejected,
323}
324
325impl From<StoppedError> for io::Error {
326    fn from(x: StoppedError) -> Self {
327        use StoppedError::*;
328        let kind = match x {
329            ZeroRttRejected => io::ErrorKind::ConnectionReset,
330            ConnectionLost(_) => io::ErrorKind::NotConnected,
331        };
332        Self::new(kind, x)
333    }
334}
335
336impl AsyncWrite for SendStream {
337    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
338        let res =
339            poll_fn(|cx| self.execute_poll_write(cx, |mut stream| stream.write(buf.as_init())))
340                .await
341                .map_err(Into::into);
342        BufResult(res, buf)
343    }
344
345    async fn flush(&mut self) -> io::Result<()> {
346        Ok(())
347    }
348
349    async fn shutdown(&mut self) -> io::Result<()> {
350        self.finish()?;
351        Ok(())
352    }
353}
354
355#[cfg(feature = "io-compat")]
356mod compat {
357    use std::{
358        ops::{Deref, DerefMut},
359        pin::Pin,
360    };
361
362    use compio_buf::IntoInner;
363
364    use super::*;
365
366    /// A [`futures_util`] compatible send stream.
367    pub struct CompatSendStream(pub(super) SendStream);
368
369    impl CompatSendStream {
370        /// Write bytes to the stream.
371        ///
372        /// Yields the number of bytes written on success. Congestion and flow
373        /// control may cause this to be shorter than `buf.len()`, indicating
374        /// that only a prefix of `buf` was written.
375        ///
376        /// This operation is cancel-safe.
377        pub async fn write(&mut self, buf: &[u8]) -> Result<usize, WriteError> {
378            poll_fn(|cx| self.execute_poll_write(cx, |mut stream| stream.write(buf))).await
379        }
380
381        /// Convenience method to write an entire buffer to the stream.
382        ///
383        /// This operation is *not* cancel-safe.
384        pub async fn write_all(&mut self, buf: &[u8]) -> Result<(), WriteError> {
385            let mut count = 0;
386            poll_fn(|cx| {
387                loop {
388                    if count == buf.len() {
389                        return Poll::Ready(Ok(()));
390                    }
391                    let n = ready!(
392                        self.execute_poll_write(cx, |mut stream| stream.write(&buf[count..]))
393                    )?;
394                    count += n;
395                }
396            })
397            .await
398        }
399    }
400
401    impl IntoInner for CompatSendStream {
402        type Inner = SendStream;
403
404        fn into_inner(self) -> Self::Inner {
405            self.0
406        }
407    }
408
409    impl Deref for CompatSendStream {
410        type Target = SendStream;
411
412        fn deref(&self) -> &Self::Target {
413            &self.0
414        }
415    }
416
417    impl DerefMut for CompatSendStream {
418        fn deref_mut(&mut self) -> &mut Self::Target {
419            &mut self.0
420        }
421    }
422
423    impl futures_util::AsyncWrite for CompatSendStream {
424        fn poll_write(
425            self: Pin<&mut Self>,
426            cx: &mut Context<'_>,
427            buf: &[u8],
428        ) -> Poll<io::Result<usize>> {
429            self.get_mut()
430                .execute_poll_write(cx, |mut stream| stream.write(buf))
431                .map_err(Into::into)
432        }
433
434        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
435            Poll::Ready(Ok(()))
436        }
437
438        fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
439            self.get_mut().finish()?;
440            Poll::Ready(Ok(()))
441        }
442    }
443}
444
445#[cfg(feature = "io-compat")]
446pub use compat::CompatSendStream;
447
448#[cfg(feature = "h3")]
449pub(crate) mod h3_impl {
450    use compio_buf::bytes::Buf;
451    use h3::quic::{self, StreamErrorIncoming, WriteBuf};
452
453    use super::*;
454
455    impl From<WriteError> for StreamErrorIncoming {
456        fn from(e: WriteError) -> Self {
457            use WriteError::*;
458            match e {
459                Stopped(code) => Self::StreamTerminated {
460                    error_code: code.into_inner(),
461                },
462                ConnectionLost(e) => Self::ConnectionErrorIncoming {
463                    connection_error: e.into(),
464                },
465
466                e => Self::Unknown(Box::new(e)),
467            }
468        }
469    }
470
471    /// A wrapper around `SendStream` that implements `quic::SendStream` and
472    /// `quic::SendStreamUnframed`.
473    pub struct SendStream<B> {
474        inner: super::SendStream,
475        buf: Option<WriteBuf<B>>,
476    }
477
478    impl<B> SendStream<B> {
479        pub(crate) fn new(conn: Shared<ConnectionInner>, stream: StreamId, is_0rtt: bool) -> Self {
480            Self {
481                inner: super::SendStream::new(conn, stream, is_0rtt),
482                buf: None,
483            }
484        }
485
486        /// Polls for the stream to be stopped or read to completion; see
487        /// [`SendStream::stopped()`](super::SendStream::stopped).
488        pub fn poll_stopped(
489            &mut self,
490            cx: &mut Context<'_>,
491        ) -> Poll<Result<Option<VarInt>, StoppedError>> {
492            self.inner.poll_stopped(cx)
493        }
494    }
495
496    impl<B> quic::SendStream<B> for SendStream<B>
497    where
498        B: Buf,
499    {
500        fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
501            if let Some(data) = &mut self.buf {
502                while data.has_remaining() {
503                    let n = ready!(
504                        self.inner
505                            .execute_poll_write(cx, |mut stream| stream.write(data.chunk()))
506                    )?;
507                    data.advance(n);
508                }
509            }
510            self.buf = None;
511            Poll::Ready(Ok(()))
512        }
513
514        fn send_data<T: Into<WriteBuf<B>>>(&mut self, data: T) -> Result<(), StreamErrorIncoming> {
515            if self.buf.is_some() {
516                return Err(WriteError::NotReady.into());
517            }
518            self.buf = Some(data.into());
519            Ok(())
520        }
521
522        fn poll_finish(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), StreamErrorIncoming>> {
523            Poll::Ready(
524                self.inner
525                    .finish()
526                    .map_err(|_| WriteError::ClosedStream.into()),
527            )
528        }
529
530        fn reset(&mut self, reset_code: u64) {
531            self.inner
532                .reset(reset_code.try_into().unwrap_or(VarInt::MAX))
533                .ok();
534        }
535
536        fn send_id(&self) -> quic::StreamId {
537            u64::from(self.inner.stream).try_into().unwrap()
538        }
539    }
540
541    impl<B> quic::SendStreamUnframed<B> for SendStream<B>
542    where
543        B: Buf,
544    {
545        fn poll_send<D: Buf>(
546            &mut self,
547            cx: &mut Context<'_>,
548            buf: &mut D,
549        ) -> Poll<Result<usize, StreamErrorIncoming>> {
550            // This signifies a bug in implementation
551            debug_assert!(
552                self.buf.is_none(),
553                "poll_send called while send stream is not ready"
554            );
555
556            let n = ready!(
557                self.inner
558                    .execute_poll_write(cx, |mut stream| stream.write(buf.chunk()))
559            )?;
560            buf.advance(n);
561            Poll::Ready(Ok(n))
562        }
563    }
564}
compio_quic/send_stream.rs

compio_quic/
send_stream.rs
