hyper/client/
dispatch.rs

1use std::task::{Context, Poll};
2#[cfg(feature = "http2")]
3use std::{future::Future, pin::Pin};
4
5#[cfg(feature = "http2")]
6use http::{Request, Response};
7#[cfg(feature = "http2")]
8use http_body::Body;
9#[cfg(feature = "http2")]
10use pin_project_lite::pin_project;
11use tokio::sync::{mpsc, oneshot};
12
13#[cfg(feature = "http2")]
14use crate::{body::Incoming, proto::h2::client::ResponseFutMap};
15
16pub(crate) type RetryPromise<T, U> = oneshot::Receiver<Result<U, TrySendError<T>>>;
17pub(crate) type Promise<T> = oneshot::Receiver<Result<T, crate::Error>>;
18
19/// An error when calling `try_send_request`.
20///
21/// There is a possibility of an error occurring on a connection in-between the
22/// time that a request is queued and when it is actually written to the IO
23/// transport. If that happens, it is safe to return the request back to the
24/// caller, as it was never fully sent.
25#[derive(Debug)]
26pub struct TrySendError<T> {
27    pub(crate) error: crate::Error,
28    pub(crate) message: Option<T>,
29}
30
31pub(crate) fn channel<T, U>() -> (Sender<T, U>, Receiver<T, U>) {
32    let (tx, rx) = mpsc::unbounded_channel();
33    let (giver, taker) = want::new();
34    let tx = Sender {
35        #[cfg(feature = "http1")]
36        buffered_once: false,
37        giver,
38        inner: tx,
39    };
40    let rx = Receiver { inner: rx, taker };
41    (tx, rx)
42}
43
44/// A bounded sender of requests and callbacks for when responses are ready.
45///
46/// While the inner sender is unbounded, the Giver is used to determine
47/// if the Receiver is ready for another request.
48pub(crate) struct Sender<T, U> {
49    /// One message is always allowed, even if the Receiver hasn't asked
50    /// for it yet. This boolean keeps track of whether we've sent one
51    /// without notice.
52    #[cfg(feature = "http1")]
53    buffered_once: bool,
54    /// The Giver helps watch that the Receiver side has been polled
55    /// when the queue is empty. This helps us know when a request and
56    /// response have been fully processed, and a connection is ready
57    /// for more.
58    giver: want::Giver,
59    /// Actually bounded by the Giver, plus `buffered_once`.
60    inner: mpsc::UnboundedSender<Envelope<T, U>>,
61}
62
63/// An unbounded version.
64///
65/// Cannot poll the Giver, but can still use it to determine if the Receiver
66/// has been dropped. However, this version can be cloned.
67#[cfg(feature = "http2")]
68pub(crate) struct UnboundedSender<T, U> {
69    /// Only used for `is_closed`, since mpsc::UnboundedSender cannot be checked.
70    giver: want::SharedGiver,
71    inner: mpsc::UnboundedSender<Envelope<T, U>>,
72}
73
74impl<T, U> Sender<T, U> {
75    #[cfg(feature = "http1")]
76    pub(crate) fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
77        self.giver
78            .poll_want(cx)
79            .map_err(|_| crate::Error::new_closed())
80    }
81
82    #[cfg(feature = "http1")]
83    pub(crate) fn is_ready(&self) -> bool {
84        self.giver.is_wanting()
85    }
86
87    #[cfg(feature = "http1")]
88    pub(crate) fn is_closed(&self) -> bool {
89        self.giver.is_canceled()
90    }
91
92    #[cfg(feature = "http1")]
93    fn can_send(&mut self) -> bool {
94        if self.giver.give() || !self.buffered_once {
95            // If the receiver is ready *now*, then of course we can send.
96            //
97            // If the receiver isn't ready yet, but we don't have anything
98            // in the channel yet, then allow one message.
99            self.buffered_once = true;
100            true
101        } else {
102            false
103        }
104    }
105
106    #[cfg(feature = "http1")]
107    pub(crate) fn try_send(&mut self, val: T) -> Result<RetryPromise<T, U>, T> {
108        if !self.can_send() {
109            return Err(val);
110        }
111        let (tx, rx) = oneshot::channel();
112        self.inner
113            .send(Envelope(Some((val, Callback::Retry(Some(tx))))))
114            .map(move |_| rx)
115            .map_err(|mut e| (e.0).0.take().expect("envelope not dropped").0)
116    }
117
118    #[cfg(feature = "http1")]
119    pub(crate) fn send(&mut self, val: T) -> Result<Promise<U>, T> {
120        if !self.can_send() {
121            return Err(val);
122        }
123        let (tx, rx) = oneshot::channel();
124        self.inner
125            .send(Envelope(Some((val, Callback::NoRetry(Some(tx))))))
126            .map(move |_| rx)
127            .map_err(|mut e| (e.0).0.take().expect("envelope not dropped").0)
128    }
129
130    #[cfg(feature = "http2")]
131    pub(crate) fn unbound(self) -> UnboundedSender<T, U> {
132        UnboundedSender {
133            giver: self.giver.shared(),
134            inner: self.inner,
135        }
136    }
137}
138
139#[cfg(feature = "http2")]
140impl<T, U> UnboundedSender<T, U> {
141    pub(crate) fn is_ready(&self) -> bool {
142        !self.giver.is_canceled()
143    }
144
145    pub(crate) fn is_closed(&self) -> bool {
146        self.giver.is_canceled()
147    }
148
149    pub(crate) fn try_send(&mut self, val: T) -> Result<RetryPromise<T, U>, T> {
150        let (tx, rx) = oneshot::channel();
151        self.inner
152            .send(Envelope(Some((val, Callback::Retry(Some(tx))))))
153            .map(move |_| rx)
154            .map_err(|mut e| (e.0).0.take().expect("envelope not dropped").0)
155    }
156
157    pub(crate) fn send(&mut self, val: T) -> Result<Promise<U>, T> {
158        let (tx, rx) = oneshot::channel();
159        self.inner
160            .send(Envelope(Some((val, Callback::NoRetry(Some(tx))))))
161            .map(move |_| rx)
162            .map_err(|mut e| (e.0).0.take().expect("envelope not dropped").0)
163    }
164}
165
166#[cfg(feature = "http2")]
167impl<T, U> Clone for UnboundedSender<T, U> {
168    fn clone(&self) -> Self {
169        UnboundedSender {
170            giver: self.giver.clone(),
171            inner: self.inner.clone(),
172        }
173    }
174}
175
176pub(crate) struct Receiver<T, U> {
177    inner: mpsc::UnboundedReceiver<Envelope<T, U>>,
178    taker: want::Taker,
179}
180
181impl<T, U> Receiver<T, U> {
182    pub(crate) fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<(T, Callback<T, U>)>> {
183        match self.inner.poll_recv(cx) {
184            Poll::Ready(item) => {
185                Poll::Ready(item.map(|mut env| env.0.take().expect("envelope not dropped")))
186            }
187            Poll::Pending => {
188                self.taker.want();
189                Poll::Pending
190            }
191        }
192    }
193
194    #[cfg(feature = "http1")]
195    pub(crate) fn close(&mut self) {
196        self.taker.cancel();
197        self.inner.close();
198    }
199
200    #[cfg(feature = "http1")]
201    pub(crate) fn try_recv(&mut self) -> Option<(T, Callback<T, U>)> {
202        match crate::common::task::now_or_never(self.inner.recv()) {
203            Some(Some(mut env)) => env.0.take(),
204            _ => None,
205        }
206    }
207}
208
209impl<T, U> Drop for Receiver<T, U> {
210    fn drop(&mut self) {
211        // Notify the giver about the closure first, before dropping
212        // the mpsc::Receiver.
213        self.taker.cancel();
214    }
215}
216
217struct Envelope<T, U>(Option<(T, Callback<T, U>)>);
218
219impl<T, U> Drop for Envelope<T, U> {
220    fn drop(&mut self) {
221        if let Some((val, cb)) = self.0.take() {
222            cb.send(Err(TrySendError {
223                error: crate::Error::new_canceled().with("connection closed"),
224                message: Some(val),
225            }));
226        }
227    }
228}
229
230pub(crate) enum Callback<T, U> {
231    #[allow(unused)]
232    Retry(Option<oneshot::Sender<Result<U, TrySendError<T>>>>),
233    NoRetry(Option<oneshot::Sender<Result<U, crate::Error>>>),
234}
235
236impl<T, U> Drop for Callback<T, U> {
237    fn drop(&mut self) {
238        match self {
239            Callback::Retry(tx) => {
240                if let Some(tx) = tx.take() {
241                    let _ = tx.send(Err(TrySendError {
242                        error: dispatch_gone(),
243                        message: None,
244                    }));
245                }
246            }
247            Callback::NoRetry(tx) => {
248                if let Some(tx) = tx.take() {
249                    let _ = tx.send(Err(dispatch_gone()));
250                }
251            }
252        }
253    }
254}
255
256#[cold]
257fn dispatch_gone() -> crate::Error {
258    // FIXME(nox): What errors do we want here?
259    crate::Error::new_user_dispatch_gone().with(if std::thread::panicking() {
260        "user code panicked"
261    } else {
262        "runtime dropped the dispatch task"
263    })
264}
265
266impl<T, U> Callback<T, U> {
267    #[cfg(feature = "http2")]
268    pub(crate) fn is_canceled(&self) -> bool {
269        match *self {
270            Callback::Retry(Some(ref tx)) => tx.is_closed(),
271            Callback::NoRetry(Some(ref tx)) => tx.is_closed(),
272            _ => unreachable!(),
273        }
274    }
275
276    pub(crate) fn poll_canceled(&mut self, cx: &mut Context<'_>) -> Poll<()> {
277        match *self {
278            Callback::Retry(Some(ref mut tx)) => tx.poll_closed(cx),
279            Callback::NoRetry(Some(ref mut tx)) => tx.poll_closed(cx),
280            _ => unreachable!(),
281        }
282    }
283
284    pub(crate) fn send(mut self, val: Result<U, TrySendError<T>>) {
285        match self {
286            Callback::Retry(ref mut tx) => {
287                let _ = tx.take().unwrap().send(val);
288            }
289            Callback::NoRetry(ref mut tx) => {
290                let _ = tx.take().unwrap().send(val.map_err(|e| e.error));
291            }
292        }
293    }
294}
295
296impl<T> TrySendError<T> {
297    /// Take the message from this error.
298    ///
299    /// The message will not always have been recovered. If an error occurs
300    /// after the message has been serialized onto the connection, it will not
301    /// be available here.
302    pub fn take_message(&mut self) -> Option<T> {
303        self.message.take()
304    }
305
306    /// Returns a reference to the recovered message.
307    ///
308    /// The message will not always have been recovered. If an error occurs
309    /// after the message has been serialized onto the connection, it will not
310    /// be available here.
311    pub fn message(&self) -> Option<&T> {
312        self.message.as_ref()
313    }
314
315    /// Consumes this to return the inner error.
316    pub fn into_error(self) -> crate::Error {
317        self.error
318    }
319
320    /// Returns a reference to the inner error.
321    pub fn error(&self) -> &crate::Error {
322        &self.error
323    }
324}
325
326#[cfg(feature = "http2")]
327pin_project! {
328    pub struct SendWhen<B>
329    where
330        B: Body,
331        B: 'static,
332    {
333        #[pin]
334        pub(crate) when: ResponseFutMap<B>,
335        #[pin]
336        pub(crate) call_back: Option<Callback<Request<B>, Response<Incoming>>>,
337    }
338}
339
340#[cfg(feature = "http2")]
341impl<B> Future for SendWhen<B>
342where
343    B: Body + 'static,
344{
345    type Output = ();
346
347    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
348        let mut this = self.project();
349
350        let mut call_back = this.call_back.take().expect("polled after complete");
351
352        match Pin::new(&mut this.when).poll(cx) {
353            Poll::Ready(Ok(res)) => {
354                call_back.send(Ok(res));
355                Poll::Ready(())
356            }
357            Poll::Pending => {
358                // check if the callback is canceled
359                match call_back.poll_canceled(cx) {
360                    Poll::Ready(v) => v,
361                    Poll::Pending => {
362                        // Move call_back back to struct before return
363                        this.call_back.set(Some(call_back));
364                        return Poll::Pending;
365                    }
366                };
367                trace!("send_when canceled");
368                Poll::Ready(())
369            }
370            Poll::Ready(Err((error, message))) => {
371                call_back.send(Err(TrySendError { error, message }));
372                Poll::Ready(())
373            }
374        }
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    #[cfg(feature = "nightly")]
381    extern crate test;
382
383    use std::future::Future;
384    use std::pin::Pin;
385    use std::task::{Context, Poll};
386
387    use super::{channel, Callback, Receiver};
388
389    #[derive(Debug)]
390    struct Custom(#[allow(dead_code)] i32);
391
392    impl<T, U> Future for Receiver<T, U> {
393        type Output = Option<(T, Callback<T, U>)>;
394
395        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
396            self.poll_recv(cx)
397        }
398    }
399
400    /// Helper to check if the future is ready after polling once.
401    struct PollOnce<'a, F>(&'a mut F);
402
403    impl<F, T> Future for PollOnce<'_, F>
404    where
405        F: Future<Output = T> + Unpin,
406    {
407        type Output = Option<()>;
408
409        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
410            match Pin::new(&mut self.0).poll(cx) {
411                Poll::Ready(_) => Poll::Ready(Some(())),
412                Poll::Pending => Poll::Ready(None),
413            }
414        }
415    }
416
417    #[cfg(not(miri))]
418    #[tokio::test]
419    async fn drop_receiver_sends_cancel_errors() {
420        let _ = pretty_env_logger::try_init();
421
422        let (mut tx, mut rx) = channel::<Custom, ()>();
423
424        // must poll once for try_send to succeed
425        assert!(PollOnce(&mut rx).await.is_none(), "rx empty");
426
427        let promise = tx.try_send(Custom(43)).unwrap();
428        drop(rx);
429
430        let fulfilled = promise.await;
431        let err = fulfilled
432            .expect("fulfilled")
433            .expect_err("promise should error");
434        match (err.error.is_canceled(), err.message) {
435            (true, Some(_)) => (),
436            e => panic!("expected Error::Cancel(_), found {:?}", e),
437        }
438    }
439
440    #[cfg(not(miri))]
441    #[tokio::test]
442    async fn sender_checks_for_want_on_send() {
443        let (mut tx, mut rx) = channel::<Custom, ()>();
444
445        // one is allowed to buffer, second is rejected
446        let _ = tx.try_send(Custom(1)).expect("1 buffered");
447        tx.try_send(Custom(2)).expect_err("2 not ready");
448
449        assert!(PollOnce(&mut rx).await.is_some(), "rx once");
450
451        // Even though 1 has been popped, only 1 could be buffered for the
452        // lifetime of the channel.
453        tx.try_send(Custom(2)).expect_err("2 still not ready");
454
455        assert!(PollOnce(&mut rx).await.is_none(), "rx empty");
456
457        let _ = tx.try_send(Custom(2)).expect("2 ready");
458    }
459
460    #[cfg(feature = "http2")]
461    #[test]
462    fn unbounded_sender_doesnt_bound_on_want() {
463        let (tx, rx) = channel::<Custom, ()>();
464        let mut tx = tx.unbound();
465
466        let _ = tx.try_send(Custom(1)).unwrap();
467        let _ = tx.try_send(Custom(2)).unwrap();
468        let _ = tx.try_send(Custom(3)).unwrap();
469
470        drop(rx);
471
472        let _ = tx.try_send(Custom(4)).unwrap_err();
473    }
474
475    #[cfg(feature = "nightly")]
476    #[bench]
477    fn giver_queue_throughput(b: &mut test::Bencher) {
478        use crate::{body::Incoming, Request, Response};
479
480        let rt = tokio::runtime::Builder::new_current_thread()
481            .build()
482            .unwrap();
483        let (mut tx, mut rx) = channel::<Request<Incoming>, Response<Incoming>>();
484
485        b.iter(move || {
486            let _ = tx.send(Request::new(Incoming::empty())).unwrap();
487            rt.block_on(async {
488                loop {
489                    let poll_once = PollOnce(&mut rx);
490                    let opt = poll_once.await;
491                    if opt.is_none() {
492                        break;
493                    }
494                }
495            });
496        })
497    }
498
499    #[cfg(feature = "nightly")]
500    #[bench]
501    fn giver_queue_not_ready(b: &mut test::Bencher) {
502        let rt = tokio::runtime::Builder::new_current_thread()
503            .build()
504            .unwrap();
505        let (_tx, mut rx) = channel::<i32, ()>();
506        b.iter(move || {
507            rt.block_on(async {
508                let poll_once = PollOnce(&mut rx);
509                assert!(poll_once.await.is_none());
510            });
511        })
512    }
513
514    #[cfg(feature = "nightly")]
515    #[bench]
516    fn giver_queue_cancel(b: &mut test::Bencher) {
517        let (_tx, mut rx) = channel::<i32, ()>();
518
519        b.iter(move || {
520            rx.taker.cancel();
521        })
522    }
523}