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#[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
44pub(crate) struct Sender<T, U> {
49 #[cfg(feature = "http1")]
53 buffered_once: bool,
54 giver: want::Giver,
59 inner: mpsc::UnboundedSender<Envelope<T, U>>,
61}
62
63#[cfg(feature = "http2")]
68pub(crate) struct UnboundedSender<T, U> {
69 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 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 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 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 pub fn take_message(&mut self) -> Option<T> {
303 self.message.take()
304 }
305
306 pub fn message(&self) -> Option<&T> {
312 self.message.as_ref()
313 }
314
315 pub fn into_error(self) -> crate::Error {
317 self.error
318 }
319
320 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 match call_back.poll_canceled(cx) {
360 Poll::Ready(v) => v,
361 Poll::Pending => {
362 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 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 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 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 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}