1use std::cell::{Cell, OnceCell, RefCell};
3use std::task::{Context, Poll, Waker};
4use std::{collections::VecDeque, fmt, future::Future, future::poll_fn, marker::PhantomData};
5use std::{mem, pin::Pin, rc::Rc, time::Instant};
6
7use ntex_bytes::{Buf, BytePages, Bytes};
8use ntex_error::Error;
9use ntex_h2::{self as h2, Capacity, OperationError, StreamRef, frame::Reason};
10use ntex_http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
11use ntex_util::{Stream, task::LocalWaker, time};
12
13use crate::client::{CallFailure, ClientError, NotGrpc, Status, transport as tr};
14use crate::server::ServerError;
15use crate::utils::{self, Data, Enc};
16use crate::{GrpcStatus, Message, consts};
17
18#[derive(thiserror::Error, Debug, Clone)]
23pub enum StreamError {
24 #[error("Stream is closed")]
26 Closed,
27 #[error("Stream is reset with {0:?}")]
29 Reset(Reason),
30 #[error("Connection is closed")]
32 Disconnected,
33 #[error("Grpc status {0:?}: {1:?}")]
36 Status(GrpcStatus, HeaderValue),
37 #[error("{0}")]
39 Client(Error<ClientError>),
40}
41
42impl StreamError {
43 pub fn status(&self) -> GrpcStatus {
45 match self {
46 StreamError::Closed => GrpcStatus::Cancelled,
47 StreamError::Reset(r) => GrpcStatus::from(*r),
48 StreamError::Disconnected => GrpcStatus::Unavailable,
49 StreamError::Status(st, _) => *st,
50 StreamError::Client(e) => e.status(),
51 }
52 }
53
54 fn from_send(err: &Error<OperationError>) -> Self {
55 match **err {
56 OperationError::RemoteReset(r) | OperationError::Closed(Some(r)) => {
57 StreamError::Reset(r)
58 }
59 OperationError::Connection(_)
60 | OperationError::Disconnected
61 | OperationError::Disconnecting => StreamError::Disconnected,
62 _ => StreamError::Closed,
63 }
64 }
65}
66
67impl From<StreamError> for ServerError {
68 fn from(err: StreamError) -> Self {
69 let (status, msg) = match err {
70 StreamError::Status(st, msg) => (st, msg),
71 StreamError::Reset(r) => (
72 GrpcStatus::from(r),
73 HeaderValue::try_from(format!("Stream reset with {r:?}"))
74 .unwrap_or_else(|_| HeaderValue::from_static("Stream reset")),
75 ),
76 StreamError::Disconnected => (
77 GrpcStatus::Unavailable,
78 HeaderValue::from_static("Connection is closed"),
79 ),
80 StreamError::Closed | StreamError::Client(_) => (
81 GrpcStatus::Internal,
82 HeaderValue::from_static("Stream is closed"),
83 ),
84 };
85 ServerError::new(status, msg, None)
86 }
87}
88
89impl From<StreamError> for Error<ClientError> {
93 fn from(err: StreamError) -> Self {
94 let err = match err {
95 StreamError::Client(e) => return e,
96 StreamError::Status(st, msg) => {
97 ClientError::InvalidRequest(CallFailure::new(st, tr::text(&msg)))
98 }
99 StreamError::Reset(r) => {
100 tr::reset_error(r, HeaderMap::new(), None, h2::StreamError::Reset(r))
101 }
102 StreamError::Disconnected => ClientError::Connection {
103 request_sent: true,
104 failure: CallFailure::new(GrpcStatus::Unavailable, "Connection is closed"),
105 },
106 StreamError::Closed => ClientError::StreamFailed(CallFailure::new(
107 GrpcStatus::Cancelled,
108 "Stream is closed",
109 )),
110 };
111 Error::from(err)
112 }
113}
114
115bitflags::bitflags! {
116 #[derive(Copy, Clone, Debug)]
117 struct Flags: u8 {
118 const CLOSED = 0b0000_0001;
120 const BUSY = 0b0000_0010;
122 const SENDING = 0b0000_0100;
124 const CLOSE_PENDING = 0b0000_1000;
126 const ATTACHED = 0b0001_0000;
128 const RESET = 0b0010_0000;
130 const CLAIMED = 0b0100_0000;
132 }
133}
134
135enum Backend {
136 Client(h2::client::SendStream),
137 Server(StreamRef),
138}
139
140impl Backend {
141 fn reset(&self, reason: Reason) {
142 match self {
143 Backend::Client(tx) => {
144 tx.reset(reason);
145 }
146 Backend::Server(stream) => {
147 stream.reset(reason);
148 }
149 }
150 }
151}
152
153#[derive(Copy, Clone)]
155pub(crate) struct SendConfig {
156 pub(crate) max_size: usize,
157 pub(crate) enc: Option<Enc>,
158 pub(crate) deadline: Option<Instant>,
160 pub(crate) service: &'static str,
161}
162
163pub(crate) struct SendShared {
164 backend: OnceCell<Backend>,
165 flags: Cell<Flags>,
166 handles: Cell<usize>,
168 ready_wakers: RefCell<Vec<Waker>>,
169 send_waiters: RefCell<VecDeque<Rc<LocalWaker>>>,
170 max_size: Cell<usize>,
171 enc: Cell<Option<Enc>>,
172 deadline: Cell<Option<Instant>>,
173 service: Cell<&'static str>,
174 headers: RefCell<Option<HeaderMap>>,
176 trailers: RefCell<HeaderMap>,
177 error: RefCell<Option<StreamError>>,
178 sent: Cell<usize>,
180}
181
182impl SendShared {
183 fn new() -> Self {
184 SendShared {
185 backend: OnceCell::new(),
186 flags: Cell::new(Flags::empty()),
187 handles: Cell::new(0),
188 ready_wakers: RefCell::new(Vec::new()),
189 send_waiters: RefCell::new(VecDeque::new()),
190 max_size: Cell::new(usize::MAX),
191 enc: Cell::new(None),
192 deadline: Cell::new(None),
193 service: Cell::new(""),
194 headers: RefCell::new(None),
195 trailers: RefCell::new(HeaderMap::new()),
196 error: RefCell::new(None),
197 sent: Cell::new(0),
198 }
199 }
200
201 pub(crate) fn server(
203 stream: StreamRef,
204 headers: HeaderMap,
205 max_size: usize,
206 enc: Option<Enc>,
207 ) -> Rc<Self> {
208 let shared = SendShared::new();
209 *shared.headers.borrow_mut() = Some(headers);
210 let _ = shared.attach(
211 Backend::Server(stream),
212 SendConfig {
213 max_size,
214 enc,
215 deadline: None,
216 service: "",
217 },
218 );
219 Rc::new(shared)
220 }
221
222 fn insert(&self, f: Flags) {
223 let mut flags = self.flags.get();
224 flags.insert(f);
225 self.flags.set(flags);
226 }
227
228 fn remove(&self, f: Flags) {
229 let mut flags = self.flags.get();
230 flags.remove(f);
231 self.flags.set(flags);
232 }
233
234 fn wake_sender(&self) {
235 if matches!(self.check_ready(), Poll::Ready(Ok(()))) {
236 let waker = self.send_waiters.borrow().front().and_then(|w| w.take());
237 if let Some(waker) = waker {
238 waker.wake();
239 }
240 }
241 }
242
243 fn wake_waiters(&self) {
244 if self.flags.get().contains(Flags::CLOSED) || self.error.borrow().is_some() {
245 let waiters = mem::take(&mut *self.send_waiters.borrow_mut());
246 for waiter in waiters {
247 waiter.wake();
248 }
249 } else {
250 self.wake_sender();
251 }
252 let wakers = mem::take(&mut *self.ready_wakers.borrow_mut());
253 for waker in wakers {
254 waker.wake();
255 }
256 }
257
258 fn set_error(&self, err: StreamError) {
260 let mut error = self.error.borrow_mut();
261 if error.is_none() {
262 *error = Some(err);
263 }
264 }
265
266 pub(crate) fn is_attached(&self) -> bool {
267 self.flags.get().contains(Flags::ATTACHED)
268 }
269
270 pub(crate) fn start_open(&self) -> Option<SendOpenGuard<'_>> {
272 if self
273 .flags
274 .get()
275 .intersects(Flags::ATTACHED | Flags::CLAIMED)
276 {
277 None
278 } else {
279 self.insert(Flags::CLAIMED);
280 Some(SendOpenGuard(self))
281 }
282 }
283
284 pub(crate) fn sent(&self) -> usize {
285 self.sent.get()
286 }
287
288 pub(crate) fn is_reset(&self) -> bool {
290 self.flags.get().contains(Flags::RESET)
291 }
292
293 fn reset(&self) {
294 if let Some(backend) = self.backend.get() {
295 backend.reset(Reason::CANCEL);
296 self.insert(Flags::RESET);
297 }
298 }
299
300 fn attach(&self, backend: Backend, cfg: SendConfig) -> Result<(), StreamError> {
302 if let Some(err) = self.error.borrow().clone() {
303 backend.reset(Reason::CANCEL);
304 return Err(err);
305 }
306 if self.backend.set(backend).is_err() {
307 return Err(StreamError::Closed);
308 }
309 self.max_size.set(cfg.max_size);
310 self.enc.set(cfg.enc);
311 self.deadline.set(cfg.deadline);
312 self.service.set(cfg.service);
313 self.insert(Flags::ATTACHED);
314 if self.flags.get().contains(Flags::CLOSED) {
315 self.half_close();
316 }
317 self.wake_waiters();
318 Ok(())
319 }
320
321 pub(crate) fn attach_client(
323 &self,
324 tx: h2::client::SendStream,
325 cfg: SendConfig,
326 ) -> Result<(), StreamError> {
327 self.attach(Backend::Client(tx), cfg)
328 }
329
330 fn half_close(&self) {
332 if let Some(Backend::Client(tx)) = self.backend.get() {
333 let mut cx = Context::from_waker(Waker::noop());
335 let fut = std::pin::pin!(tx.send_pages(Bytes::new(), true));
336 if let Poll::Ready(Err(err)) = fut.poll(&mut cx) {
337 log::trace!("Cannot close grpc stream: {err:?}");
338 }
339 }
340 }
341
342 fn check_ready(&self) -> Poll<Result<(), StreamError>> {
343 if let Some(err) = self.error.borrow().as_ref() {
344 return Poll::Ready(Err(err.clone()));
345 }
346 let flags = self.flags.get();
347 if flags.contains(Flags::CLOSED) {
348 Poll::Ready(Err(StreamError::Closed))
349 } else if flags.contains(Flags::ATTACHED) && !flags.contains(Flags::BUSY) {
350 Poll::Ready(Ok(()))
351 } else {
352 Poll::Pending
353 }
354 }
355
356 fn poll_ready(&self, cx: &Context<'_>) -> Poll<Result<(), StreamError>> {
357 let ready = self.check_ready();
358 if ready.is_pending() {
359 let mut wakers = self.ready_wakers.borrow_mut();
360 if !wakers.iter().any(|w| w.will_wake(cx.waker())) {
361 wakers.push(cx.waker().clone());
362 }
363 }
364 ready
365 }
366
367 async fn send<T: Message>(&self, msg: &T) -> Result<(), StreamError> {
368 let _guard = SendWait::new(self).await?;
369
370 let mut buf = BytePages::default();
371 msg.write(&mut buf);
372 let buf = utils::frame_message(buf, self.max_size.get(), self.enc.get())
373 .await
374 .map_err(|(st, msg)| StreamError::Status(st, msg))?;
375 let len = buf.len();
376
377 let Some(backend) = self.backend.get() else {
378 return Err(StreamError::Closed);
379 };
380 let res = match backend {
381 Backend::Client(tx) => self.send_client(tx, buf).await,
382 Backend::Server(stream) => {
383 let headers = self.headers.borrow_mut().take();
384 if let Some(hdrs) = headers
385 && let Err(err) = send_response_headers(stream, hdrs)
386 {
387 Err(StreamError::from_send(&err))
388 } else {
389 self.insert(Flags::SENDING);
390 let res = stream.send_pages(buf, false).await;
391 self.remove(Flags::SENDING);
392 res.map_err(|e| StreamError::from_send(&e))
393 }
394 }
395 };
396 match res {
397 Ok(()) => {
398 self.sent.set(self.sent.get() + len);
399 Ok(())
400 }
401 Err(err) => {
402 self.set_error(err.clone());
403 Err(err)
404 }
405 }
406 }
407
408 async fn send_client(
409 &self,
410 tx: &h2::client::SendStream,
411 buf: BytePages,
412 ) -> Result<(), StreamError> {
413 let deadline_exceeded = || {
414 tx.reset(Reason::CANCEL);
415 let err = tr::deadline_exceeded(tr::TIMEOUT_RAN_OUT, HeaderMap::new());
416 StreamError::Client(err.with_service(self.service.get()))
417 };
418 self.insert(Flags::SENDING);
419 let res = if let Some(deadline) = self.deadline.get() {
420 let remaining = deadline.saturating_duration_since(Instant::now());
421 if remaining.is_zero() {
422 Err(deadline_exceeded())
423 } else {
424 match time::timeout(remaining, tx.send_pages(buf, false)).await {
425 Ok(res) => res.map_err(|e| StreamError::from_send(&e)),
426 Err(()) => Err(deadline_exceeded()),
427 }
428 }
429 } else {
430 tx.send_pages(buf, false)
431 .await
432 .map_err(|e| StreamError::from_send(&e))
433 };
434 self.remove(Flags::SENDING);
435 res
436 }
437
438 fn close(&self) {
439 let flags = self.flags.get();
440 if flags.contains(Flags::CLOSED) {
441 return;
442 }
443 self.insert(Flags::CLOSED);
444 if flags.contains(Flags::ATTACHED) {
445 if flags.contains(Flags::BUSY) {
446 self.insert(Flags::CLOSE_PENDING);
447 } else {
448 self.half_close();
449 }
450 }
451 self.wake_waiters();
452 }
453
454 fn cancel(&self) {
455 self.insert(Flags::CLOSED);
456 self.set_error(StreamError::Closed);
457 self.reset();
458 self.wake_waiters();
459 }
460
461 pub(crate) fn finish(&self) -> Option<(Option<HeaderMap>, HeaderMap)> {
467 let flags = self.flags.get();
468 self.insert(Flags::CLOSED);
469 if flags.contains(Flags::SENDING) {
470 self.set_error(StreamError::Closed);
471 self.reset();
472 }
473 self.wake_waiters();
474 if self.error.borrow().is_some() {
475 return None;
476 }
477 Some((
478 self.headers.borrow_mut().take(),
479 mem::take(&mut *self.trailers.borrow_mut()),
480 ))
481 }
482}
483
484pub(crate) fn send_response_headers(
486 stream: &StreamRef,
487 headers: HeaderMap,
488) -> Result<(), Error<OperationError>> {
489 let res = stream
490 .send_response(StatusCode::OK, &headers, false)
491 .inspect_err(|err| {
492 log::debug!("{}: Cannot send response headers: {err:?}", stream.tag());
493 stream.reset(Reason::INTERNAL_ERROR);
494 });
495 ntex_h2::release_header_map(headers);
496 res
497}
498
499pub(crate) struct SendOpenGuard<'a>(&'a SendShared);
501
502impl SendOpenGuard<'_> {
503 pub(crate) fn fail(self, err: Error<ClientError>) {
504 self.0.set_error(StreamError::Client(err));
505 }
506}
507
508impl Drop for SendOpenGuard<'_> {
509 fn drop(&mut self) {
510 if !self.0.is_attached() {
511 self.0.cancel();
512 }
513 }
514}
515
516struct SendWait<'a> {
518 shared: &'a SendShared,
519 waiter: Option<Rc<LocalWaker>>,
520}
521
522impl<'a> SendWait<'a> {
523 fn new(shared: &'a SendShared) -> Self {
524 Self {
525 shared,
526 waiter: None,
527 }
528 }
529}
530
531impl<'a> Future for SendWait<'a> {
532 type Output = Result<SendGuard<'a>, StreamError>;
533
534 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
535 let this = self.get_mut();
536 match this.shared.check_ready() {
537 Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
538 Poll::Ready(Ok(())) => {
539 let mut waiters = this.shared.send_waiters.borrow_mut();
540 if waiters
541 .front()
542 .is_none_or(|front| this.waiter.as_ref().is_some_and(|w| Rc::ptr_eq(w, front)))
543 {
544 waiters.pop_front();
545 this.shared.insert(Flags::BUSY);
546 let guard = SendGuard(this.shared);
547 drop(waiters);
548 this.waiter = None;
549 return Poll::Ready(Ok(guard));
550 }
551 }
552 Poll::Pending => {}
553 }
554 let waiter = this.waiter.get_or_insert_with(|| {
555 let waiter = Rc::new(LocalWaker::new());
556 this.shared
557 .send_waiters
558 .borrow_mut()
559 .push_back(waiter.clone());
560 waiter
561 });
562 waiter.register(cx.waker());
563 Poll::Pending
564 }
565}
566
567impl Drop for SendWait<'_> {
568 fn drop(&mut self) {
569 if let Some(waiter) = &self.waiter {
570 let first = {
571 let mut waiters = self.shared.send_waiters.borrow_mut();
572 if let Some(idx) = waiters.iter().position(|w| Rc::ptr_eq(w, waiter)) {
573 waiters.remove(idx);
574 idx == 0
575 } else {
576 false
577 }
578 };
579 if first {
580 self.shared.wake_sender();
581 }
582 }
583 }
584}
585
586struct SendGuard<'a>(&'a SendShared);
588
589impl Drop for SendGuard<'_> {
590 fn drop(&mut self) {
591 let shared = self.0;
592 let flags = shared.flags.get();
593 if flags.contains(Flags::SENDING) {
594 shared.reset();
596 shared.set_error(StreamError::Closed);
597 shared.insert(Flags::CLOSED);
598 }
599 shared.remove(Flags::BUSY | Flags::SENDING);
600 if flags.contains(Flags::CLOSE_PENDING) {
601 shared.remove(Flags::CLOSE_PENDING);
602 if !flags.contains(Flags::SENDING) {
603 shared.half_close();
604 }
605 }
606 shared.wake_waiters();
607 }
608}
609
610pub struct SendStream<T> {
623 shared: Rc<SendShared>,
624 _t: PhantomData<fn(&T)>,
625}
626
627impl<T> SendStream<T> {
628 pub fn new() -> Self {
633 Self::from_shared(Rc::new(SendShared::new()))
634 }
635
636 pub(crate) fn from_shared(shared: Rc<SendShared>) -> Self {
637 shared.handles.set(shared.handles.get() + 1);
638 SendStream {
639 shared,
640 _t: PhantomData,
641 }
642 }
643
644 pub(crate) fn shared(&self) -> &Rc<SendShared> {
645 &self.shared
646 }
647
648 pub fn poll_ready(&self, cx: &Context<'_>) -> Poll<Result<(), StreamError>> {
654 self.shared.poll_ready(cx)
655 }
656
657 pub async fn ready(&self) -> Result<(), StreamError> {
659 poll_fn(|cx| self.shared.poll_ready(cx)).await
660 }
661
662 pub fn close(&self) {
667 self.shared.close();
668 }
669
670 pub fn cancel(&self) {
676 self.shared.cancel();
677 }
678
679 pub fn is_closed(&self) -> bool {
681 self.shared.flags.get().contains(Flags::CLOSED) || self.shared.error.borrow().is_some()
682 }
683
684 pub fn header(&self, name: HeaderName, value: HeaderValue) -> bool {
692 if let Some(hdrs) = self.shared.headers.borrow_mut().as_mut() {
693 hdrs.append(name, value);
694 true
695 } else {
696 false
697 }
698 }
699
700 pub fn trailer(&self, name: HeaderName, value: HeaderValue) {
704 self.shared.trailers.borrow_mut().append(name, value);
705 }
706}
707
708impl<T: Message> SendStream<T> {
709 pub async fn send(&self, msg: &T) -> Result<(), StreamError> {
720 self.shared.send(msg).await
721 }
722}
723
724impl<T> Default for SendStream<T> {
725 fn default() -> Self {
726 Self::new()
727 }
728}
729
730impl<T> Clone for SendStream<T> {
731 fn clone(&self) -> Self {
732 Self::from_shared(self.shared.clone())
733 }
734}
735
736impl<T> Drop for SendStream<T> {
737 fn drop(&mut self) {
738 let handles = self.shared.handles.get() - 1;
739 self.shared.handles.set(handles);
740 if handles == 0 {
741 self.shared.close();
742 }
743 }
744}
745
746impl<T> fmt::Debug for SendStream<T> {
747 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
748 f.debug_struct("SendStream")
749 .field("flags", &self.shared.flags.get())
750 .field("sent", &self.shared.sent.get())
751 .finish()
752 }
753}
754
755pub(crate) struct Decoder {
762 data: Data,
763 caps: VecDeque<Capacity>,
764 max_size: usize,
765 enc: Option<HeaderValue>,
766 unsupported: GrpcStatus,
768}
769
770impl Decoder {
771 pub(crate) fn new(max_size: usize, enc: Option<HeaderValue>, unsupported: GrpcStatus) -> Self {
772 Decoder {
773 data: Data::Empty,
774 caps: VecDeque::new(),
775 max_size,
776 enc,
777 unsupported,
778 }
779 }
780
781 fn push(&mut self, data: Bytes, cap: Capacity) {
782 self.data.push(data, self.max_size);
783 self.caps.push_back(cap);
784 if !self.has_message() {
785 self.caps.clear();
786 }
787 }
788
789 fn next_len(&self) -> Option<usize> {
791 let data = self.data.as_slice();
792 Some(u32::from_be_bytes(data.get(1..5)?.try_into().ok()?) as usize)
793 }
794
795 fn has_message(&self) -> bool {
796 self.next_len()
797 .is_some_and(|len| self.data.as_slice().len() - 5 >= len)
798 }
799
800 fn has_data(&self) -> bool {
801 !self.data.as_slice().is_empty()
802 }
803
804 fn next(&mut self) -> Result<Option<(u8, Bytes)>, usize> {
807 let Some(len) = self.next_len() else {
808 return Ok(None);
809 };
810 if len > self.max_size {
811 return Err(len);
812 }
813 if self.data.as_slice().len() - 5 < len {
814 return Ok(None);
815 }
816 let mut rest = self.data.get();
817 let mut msg = rest.split_to(5 + len);
818 let flag = msg[0];
819 msg.advance(5);
820 if !rest.is_empty() {
821 self.data = Data::Chunk(rest);
822 }
823 if !self.has_message() {
824 self.caps.clear();
825 }
826 Ok(Some((flag, msg)))
827 }
828}
829
830struct ServerRecvInner {
831 chunks: VecDeque<(Bytes, Capacity)>,
832 end: Option<Result<HeaderMap, StreamError>>,
834 ended: bool,
836 closed: bool,
838}
839
840pub(crate) struct ServerRecv {
842 inner: RefCell<ServerRecvInner>,
843 waker: LocalWaker,
844}
845
846impl ServerRecv {
847 pub(crate) fn new(eof: bool) -> Rc<Self> {
848 Rc::new(ServerRecv {
849 inner: RefCell::new(ServerRecvInner {
850 chunks: VecDeque::new(),
851 end: eof.then(|| Ok(HeaderMap::new())),
852 ended: eof,
853 closed: false,
854 }),
855 waker: LocalWaker::new(),
856 })
857 }
858
859 pub(crate) fn push_data(&self, data: Bytes, cap: Capacity) {
860 let mut inner = self.inner.borrow_mut();
861 if !inner.closed && !inner.ended {
862 inner.chunks.push_back((data, cap));
863 self.waker.wake();
864 }
865 }
866
867 pub(crate) fn push_end(&self, end: Result<HeaderMap, StreamError>) {
869 let mut inner = self.inner.borrow_mut();
870 if !inner.ended {
871 inner.ended = true;
872 if !inner.closed {
873 inner.end = Some(end);
874 }
875 self.waker.wake();
876 }
877 }
878
879 pub(crate) fn is_ended(&self) -> bool {
881 self.inner.borrow().ended
882 }
883
884 fn close(&self) {
885 let mut inner = self.inner.borrow_mut();
886 inner.closed = true;
887 inner.chunks.clear();
888 inner.end = None;
889 }
890
891 fn poll(&self, trailers: &mut HeaderMap, cx: &Context<'_>) -> Poll<Item> {
892 let mut inner = self.inner.borrow_mut();
893 if let Some((data, cap)) = inner.chunks.pop_front() {
894 Poll::Ready(Item::Data(data, cap))
895 } else if let Some(end) = inner.end.take() {
896 Poll::Ready(Item::End(end.map(|t| *trailers = t)))
897 } else if inner.closed {
898 Poll::Ready(Item::End(Err(StreamError::Closed)))
899 } else {
900 self.waker.register(cx.waker());
901 Poll::Pending
902 }
903 }
904}
905
906enum Item {
908 Data(Bytes, Capacity),
909 End(Result<(), StreamError>),
910}
911
912type ClientResult = Result<Option<(Bytes, Capacity)>, Error<ClientError>>;
913
914pub(crate) struct ClientRecv {
916 rx: Option<h2::client::RecvStream>,
918 deadline: Option<Instant>,
919 sleep: Option<time::Sleep>,
920 status: Option<StatusCode>,
921 other: Data,
923 pending_end: Option<Error<ClientError>>,
925 service: &'static str,
926 send: Option<Rc<SendShared>>,
928}
929
930impl ClientRecv {
931 pub(crate) fn new(
932 rx: h2::client::RecvStream,
933 deadline: Option<Instant>,
934 service: &'static str,
935 send: Option<Rc<SendShared>>,
936 ) -> Self {
937 let sleep = deadline.map(|d| {
938 time::sleep(time::Millis::from(
939 d.saturating_duration_since(Instant::now()),
940 ))
941 });
942 ClientRecv {
943 rx: Some(rx),
944 deadline,
945 sleep,
946 status: None,
947 other: Data::Empty,
948 pending_end: None,
949 service,
950 send,
951 }
952 }
953
954 fn close(&mut self) {
955 self.rx = None;
956 self.sleep = None;
957 }
958
959 fn poll(
960 &mut self,
961 headers: &mut HeaderMap,
962 trailers: &mut HeaderMap,
963 cx: &mut Context<'_>,
964 ) -> Poll<Item> {
965 let res = match self.poll_inner(headers, trailers, cx) {
966 Poll::Ready(res) => res,
967 Poll::Pending => return Poll::Pending,
968 };
969 Poll::Ready(match res {
970 Ok(Some((data, cap))) => Item::Data(data, cap),
971 Ok(None) => {
972 self.close();
973 Item::End(Ok(()))
974 }
975 Err(err) => {
976 self.close();
977 Item::End(Err(StreamError::Client(err.with_service(self.service))))
978 }
979 })
980 }
981
982 fn end(
984 &mut self,
985 hdrs: HeaderMap,
986 headers: &HeaderMap,
987 trailers: &mut HeaderMap,
988 ) -> ClientResult {
989 *trailers = hdrs.clone();
990 match tr::check_grpc_status(&hdrs) {
991 Some(Ok(GrpcStatus::Ok)) | None if !tr::is_grpc(self.status, headers) => {
992 Err(self.other_error(headers))
993 }
994 Some(Ok(GrpcStatus::Ok)) => Ok(None),
995 None => Err(tr::invalid_response(
996 GrpcStatus::Unknown,
997 tr::NO_GRPC_STATUS,
998 hdrs,
999 )),
1000 Some(Ok(st)) => Err(Error::from(ClientError::Status(Status::new(st, hdrs)))),
1001 Some(Err(msg)) => Err(tr::invalid_response(GrpcStatus::Unknown, msg, hdrs)),
1002 }
1003 }
1004
1005 fn poll_inner(
1006 &mut self,
1007 headers: &mut HeaderMap,
1008 trailers: &mut HeaderMap,
1009 cx: &mut Context<'_>,
1010 ) -> Poll<ClientResult> {
1011 if let Some(err) = self.pending_end.take() {
1012 return Poll::Ready(Err(err));
1013 }
1014 loop {
1015 let Some(rx) = self.rx.as_ref() else {
1016 return Poll::Ready(Err(StreamError::Closed.into()));
1017 };
1018 if self
1019 .sleep
1020 .as_ref()
1021 .is_some_and(|s| s.poll_elapsed(cx).is_ready())
1022 {
1023 return Poll::Ready(Err(tr::deadline_exceeded(
1024 tr::TIMEOUT_RAN_OUT,
1025 headers.clone(),
1026 )));
1027 }
1028 let Poll::Ready(msg) = rx.poll_recv(cx) else {
1029 return Poll::Pending;
1030 };
1031 let Some(msg) = msg else {
1032 if self.deadline.is_some_and(|d| Instant::now() >= d) {
1033 return Poll::Ready(Err(tr::deadline_exceeded(
1034 tr::TIMEOUT_RAN_OUT,
1035 headers.clone(),
1036 )));
1037 }
1038 if self.send.as_ref().is_some_and(|s| s.is_reset()) {
1039 let failure = CallFailure::new(GrpcStatus::Cancelled, tr::CANCELLED)
1040 .with_headers(headers.clone());
1041 return Poll::Ready(Err(Error::from(ClientError::StreamFailed(failure))));
1042 }
1043 let failure = CallFailure::new(GrpcStatus::Unavailable, tr::STREAM_CLOSED)
1044 .with_headers(headers.clone());
1045 return Poll::Ready(Err(Error::from(ClientError::Connection {
1046 request_sent: true,
1047 failure,
1048 })));
1049 };
1050 match msg.kind {
1051 h2::MessageKind::Headers {
1052 headers: hdrs,
1053 pseudo,
1054 eof,
1055 } => {
1056 if eof {
1057 *trailers = hdrs.clone();
1058 self.status = pseudo.status;
1059 return Poll::Ready(trailers_only(pseudo.status, hdrs));
1060 }
1061 *headers = hdrs;
1062 self.status = pseudo.status;
1063 }
1064 h2::MessageKind::Data(data, cap) => {
1065 if tr::is_grpc(self.status, headers) {
1066 return Poll::Ready(Ok(Some((data, cap))));
1067 }
1068 self.other.push(data, 0);
1069 if tr::limit_other_body(self.status, headers, &mut self.other) {
1070 return Poll::Ready(Err(self.other_error(headers)));
1071 }
1072 }
1073 h2::MessageKind::Eof(h2::StreamEof::Data(data, cap)) => {
1074 if tr::is_grpc(self.status, headers) {
1075 self.pending_end = Some(tr::invalid_response(
1076 GrpcStatus::Internal,
1077 tr::NO_TRAILERS,
1078 headers.clone(),
1079 ));
1080 return Poll::Ready(Ok(Some((data, cap))));
1081 }
1082 self.other.push(data, 0);
1083 tr::limit_other_body(self.status, headers, &mut self.other);
1084 return Poll::Ready(Err(self.other_error(headers)));
1085 }
1086 h2::MessageKind::Eof(h2::StreamEof::Trailers(hdrs)) => {
1087 return Poll::Ready(self.end(hdrs, headers, trailers));
1088 }
1089 h2::MessageKind::Eof(h2::StreamEof::Error(err)) => {
1090 return Poll::Ready(Err(
1091 err.map(|e| tr::stream_error(e, headers.clone(), self.deadline))
1092 ));
1093 }
1094 h2::MessageKind::Disconnect(err) => {
1095 return Poll::Ready(Err(tr::operation_error(
1096 err,
1097 true,
1098 headers.clone(),
1099 self.deadline,
1100 )));
1101 }
1102 }
1103 }
1104 }
1105
1106 fn other_error(&mut self, headers: &HeaderMap) -> Error<ClientError> {
1108 let body = self.other.get();
1109 Error::from(ClientError::NotGrpc(NotGrpc::new(
1110 self.status,
1111 headers.clone(),
1112 body,
1113 )))
1114 }
1115}
1116
1117fn trailers_only(status: Option<StatusCode>, hdrs: HeaderMap) -> ClientResult {
1120 if tr::is_grpc(status, &hdrs)
1121 && matches!(tr::check_grpc_status(&hdrs), Some(Ok(GrpcStatus::Ok)))
1122 {
1123 Ok(None)
1124 } else {
1125 Err(Error::from(tr::trailers_only(status, hdrs)))
1126 }
1127}
1128
1129const ERR_NO_MESSAGE: HeaderValue = HeaderValue::from_static("grpc: request without a message");
1130const ERR_TRUNCATED: HeaderValue = HeaderValue::from_static("grpc: request message is truncated");
1131const ERR_EXTRA_DATA: HeaderValue =
1132 HeaderValue::from_static("grpc: received data after the request message");
1133const ERR_DECODE: HeaderValue = HeaderValue::from_static("Cannot decode grpc message");
1134
1135enum Source {
1136 Server(Rc<ServerRecv>),
1137 Client(Box<ClientRecv>),
1138}
1139
1140type Decompress = Pin<Box<dyn Future<Output = Result<Bytes, (GrpcStatus, HeaderValue)>>>>;
1141
1142pub struct RecvStream<T> {
1158 source: Source,
1159 decoder: Decoder,
1160 decompress: Option<Decompress>,
1161 headers: HeaderMap,
1162 trailers: HeaderMap,
1163 end: Option<Result<(), StreamError>>,
1164 done: bool,
1165 res_size: usize,
1166 _t: PhantomData<fn() -> T>,
1167}
1168
1169impl<T> RecvStream<T> {
1170 pub(crate) fn server(
1171 recv: Rc<ServerRecv>,
1172 headers: HeaderMap,
1173 max_size: usize,
1174 enc: Option<HeaderValue>,
1175 ) -> Self {
1176 Self::with_source(
1177 Source::Server(recv),
1178 Decoder::new(max_size, enc, GrpcStatus::Unimplemented),
1179 headers,
1180 )
1181 }
1182
1183 pub(crate) fn client(recv: ClientRecv, max_size: usize) -> Self {
1184 Self::with_source(
1185 Source::Client(Box::new(recv)),
1186 Decoder::new(max_size, None, GrpcStatus::Internal),
1187 HeaderMap::new(),
1188 )
1189 }
1190
1191 fn with_source(source: Source, decoder: Decoder, headers: HeaderMap) -> Self {
1192 RecvStream {
1193 source,
1194 decoder,
1195 decompress: None,
1196 headers,
1197 trailers: HeaderMap::new(),
1198 end: None,
1199 done: false,
1200 res_size: 0,
1201 _t: PhantomData,
1202 }
1203 }
1204
1205 pub fn headers(&self) -> &HeaderMap {
1211 &self.headers
1212 }
1213
1214 pub fn trailers(&self) -> &HeaderMap {
1216 &self.trailers
1217 }
1218
1219 pub fn status(&self) -> Option<GrpcStatus> {
1221 if !self.done {
1222 return None;
1223 }
1224 Some(match &self.end {
1225 Some(Err(err)) => err.status(),
1226 Some(Ok(())) | None => GrpcStatus::Ok,
1227 })
1228 }
1229
1230 pub(crate) fn res_size(&self) -> usize {
1232 self.res_size
1233 }
1234
1235 fn is_client(&self) -> bool {
1236 matches!(self.source, Source::Client(_))
1237 }
1238
1239 fn service(&self) -> &'static str {
1240 match &self.source {
1241 Source::Client(recv) => recv.service,
1242 Source::Server(_) => "",
1243 }
1244 }
1245
1246 fn close_source(&mut self) {
1247 match &mut self.source {
1248 Source::Server(recv) => recv.close(),
1249 Source::Client(recv) => recv.close(),
1250 }
1251 }
1252
1253 fn fail(&mut self, err: StreamError) -> StreamError {
1255 self.done = true;
1256 self.decompress = None;
1257 self.end = Some(Err(err.clone()));
1258 self.close_source();
1259 self.decoder.data = Data::Empty;
1260 self.decoder.caps.clear();
1261 err
1262 }
1263
1264 fn status_error(&self, st: GrpcStatus, msg: HeaderValue) -> StreamError {
1265 if self.is_client() {
1266 let err = tr::invalid_response(st, tr::text(&msg), self.headers.clone());
1267 StreamError::Client(err.with_service(self.service()))
1268 } else {
1269 StreamError::Status(st, msg)
1270 }
1271 }
1272
1273 fn too_large(&self, len: usize) -> StreamError {
1275 let max_size = self.decoder.max_size;
1276 if self.is_client() {
1277 let msg = tr::too_large(len, max_size);
1278 let err =
1279 tr::invalid_response(GrpcStatus::ResourceExhausted, msg, self.headers.clone());
1280 StreamError::Client(err.with_service(self.service()))
1281 } else {
1282 let msg = format!("grpc: received message larger than max ({len} vs. {max_size})");
1283 let msg = HeaderValue::try_from(msg).unwrap_or_else(|_| {
1284 HeaderValue::from_static("grpc: received message larger than max")
1285 });
1286 StreamError::Status(GrpcStatus::ResourceExhausted, msg)
1287 }
1288 }
1289
1290 pub(crate) fn truncated(&self) -> StreamError {
1293 if self.is_client() {
1294 let err =
1295 tr::invalid_response(GrpcStatus::Internal, tr::TRUNCATED, self.trailers.clone());
1296 StreamError::Client(err.with_service(self.service()))
1297 } else {
1298 StreamError::Status(GrpcStatus::InvalidArgument, ERR_TRUNCATED)
1299 }
1300 }
1301
1302 fn poll_source(&mut self, cx: &mut Context<'_>) -> Poll<Item> {
1303 match &mut self.source {
1304 Source::Server(recv) => recv.poll(&mut self.trailers, cx),
1305 Source::Client(recv) => recv.poll(&mut self.headers, &mut self.trailers, cx),
1306 }
1307 }
1308}
1309
1310impl<T: Message> RecvStream<T> {
1311 fn decode(&mut self, mut msg: Bytes) -> Result<T, StreamError> {
1312 T::read(&mut msg).map_err(|e| {
1313 let err = if self.is_client() {
1314 let failure = CallFailure::new(GrpcStatus::Internal, e.to_string())
1315 .with_headers(self.headers.clone())
1316 .with_cause(e);
1317 let err = Error::from(ClientError::InvalidResponse(failure));
1318 StreamError::Client(err.with_service(self.service()))
1319 } else {
1320 StreamError::Status(GrpcStatus::InvalidArgument, ERR_DECODE)
1321 };
1322 self.fail(err)
1323 })
1324 }
1325
1326 pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<T, StreamError>>> {
1331 loop {
1332 if let Some(fut) = self.decompress.as_mut() {
1333 let res = match fut.as_mut().poll(cx) {
1334 Poll::Ready(res) => res,
1335 Poll::Pending => return Poll::Pending,
1336 };
1337 self.decompress = None;
1338 return Poll::Ready(Some(match res {
1339 Ok(msg) => self.decode(msg),
1340 Err((st, msg)) => Err(self.fail(self.status_error(st, msg))),
1341 }));
1342 }
1343 if self.done {
1344 return Poll::Ready(None);
1345 }
1346 match self.decoder.next() {
1347 Ok(Some((flag, msg))) => {
1348 self.res_size += msg.len() + 5;
1349 if flag == 0 {
1350 return Poll::Ready(Some(self.decode(msg)));
1351 }
1352 self.decompress = Some(Box::pin(utils::read_message(
1353 flag,
1354 self.decoder.enc.clone(),
1355 msg,
1356 self.decoder.max_size,
1357 self.decoder.unsupported,
1358 )));
1359 continue;
1360 }
1361 Ok(None) => {}
1362 Err(len) => return Poll::Ready(Some(Err(self.fail(self.too_large(len))))),
1363 }
1364 if let Some(end) = &self.end {
1365 let end = end.clone();
1366 return Poll::Ready(match end {
1367 Ok(()) if self.decoder.has_data() => Some(Err(self.fail(self.truncated()))),
1368 Ok(()) => {
1369 self.done = true;
1370 None
1371 }
1372 Err(err) => {
1373 self.done = true;
1374 Some(Err(err))
1375 }
1376 });
1377 }
1378 match self.poll_source(cx) {
1379 Poll::Ready(Item::Data(data, cap)) => {
1380 if self.is_client() && self.decoder.enc.is_none() {
1382 self.decoder.enc = self.headers.get(consts::GRPC_ENCODING).cloned();
1383 }
1384 self.decoder.push(data, cap);
1385 }
1386 Poll::Ready(Item::End(res)) => self.end = Some(res),
1387 Poll::Pending => return Poll::Pending,
1388 }
1389 }
1390 }
1391
1392 pub async fn recv(&mut self) -> Option<Result<T, StreamError>> {
1394 poll_fn(|cx| self.poll_recv(cx)).await
1395 }
1396
1397 #[doc(hidden)]
1399 pub async fn single(mut self) -> Result<T, ServerError> {
1400 let msg = match self.recv().await {
1401 Some(Ok(msg)) => msg,
1402 Some(Err(err)) => return Err(err.into()),
1403 None => {
1404 return Err(ServerError::new(
1405 GrpcStatus::InvalidArgument,
1406 ERR_NO_MESSAGE,
1407 None,
1408 ));
1409 }
1410 };
1411 match self.recv().await {
1412 None => Ok(msg),
1413 Some(Err(err)) => Err(err.into()),
1414 Some(Ok(_)) => Err(ServerError::new(GrpcStatus::Internal, ERR_EXTRA_DATA, None)),
1415 }
1416 }
1417}
1418
1419impl<T: Message> Stream for RecvStream<T> {
1420 type Item = Result<T, StreamError>;
1421
1422 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1423 self.get_mut().poll_recv(cx)
1424 }
1425}
1426
1427impl<T> Drop for RecvStream<T> {
1428 fn drop(&mut self) {
1429 if let Source::Server(recv) = &self.source {
1430 recv.close();
1431 }
1432 }
1433}
1434
1435impl<T> fmt::Debug for RecvStream<T> {
1436 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1437 f.debug_struct("RecvStream")
1438 .field("headers", &self.headers)
1439 .field("trailers", &self.trailers)
1440 .field("end", &self.end)
1441 .field("done", &self.done)
1442 .finish()
1443 }
1444}
1445
1446#[doc(hidden)]
1448pub struct CallStreams {
1449 recv: Option<(Rc<ServerRecv>, HeaderMap, usize, Option<HeaderValue>)>,
1450 send: Option<Rc<SendShared>>,
1451}
1452
1453impl CallStreams {
1454 pub(crate) fn new(
1455 recv: Rc<ServerRecv>,
1456 headers: HeaderMap,
1457 max_size: usize,
1458 enc: Option<HeaderValue>,
1459 send: Rc<SendShared>,
1460 ) -> Self {
1461 CallStreams {
1462 recv: Some((recv, headers, max_size, enc)),
1463 send: Some(send),
1464 }
1465 }
1466
1467 pub fn recv<T>(&mut self) -> RecvStream<T> {
1473 let (recv, headers, max_size, enc) = self
1474 .recv
1475 .take()
1476 .expect("The request stream is taken already");
1477 RecvStream::server(recv, headers, max_size, enc)
1478 }
1479
1480 pub fn send<T>(&mut self) -> SendStream<T> {
1486 SendStream::from_shared(
1487 self.send
1488 .take()
1489 .expect("The response stream is taken already"),
1490 )
1491 }
1492}
1493
1494impl Drop for CallStreams {
1495 fn drop(&mut self) {
1496 if let Some((recv, ..)) = self.recv.take() {
1498 recv.close();
1499 }
1500 }
1501}
1502
1503impl fmt::Debug for CallStreams {
1504 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1505 f.debug_struct("CallStreams")
1506 .field("recv", &self.recv.is_some())
1507 .field("send", &self.send.is_some())
1508 .finish()
1509 }
1510}
1511
1512#[cfg(test)]
1513mod tests {
1514 use std::sync::Arc;
1515 use std::sync::atomic::{AtomicUsize, Ordering};
1516 use std::task::Wake;
1517
1518 use super::*;
1519
1520 #[derive(Default)]
1521 struct WakeCount(AtomicUsize);
1522
1523 impl WakeCount {
1524 fn waker() -> (Arc<Self>, Waker) {
1525 let count = Arc::new(Self::default());
1526 let waker = Waker::from(count.clone());
1527 (count, waker)
1528 }
1529
1530 fn take(&self) -> usize {
1531 self.0.swap(0, Ordering::Relaxed)
1532 }
1533 }
1534
1535 impl Wake for WakeCount {
1536 fn wake(self: Arc<Self>) {
1537 self.0.fetch_add(1, Ordering::Relaxed);
1538 }
1539 }
1540
1541 fn frame(flag: u8, payload: &[u8]) -> Vec<u8> {
1542 let mut buf = vec![flag];
1543 buf.extend_from_slice(&u32::try_from(payload.len()).unwrap().to_be_bytes());
1544 buf.extend_from_slice(payload);
1545 buf
1546 }
1547
1548 fn decoder(max_size: usize) -> Decoder {
1549 Decoder::new(max_size, None, GrpcStatus::Unimplemented)
1550 }
1551
1552 #[test]
1553 fn decoder_split_and_coalesced() {
1554 let mut dec = decoder(1024);
1555 let mut data = frame(0, b"hello");
1556 data.extend(frame(1, b""));
1557 data.extend(frame(0, b"world"));
1558
1559 for chunk in [&data[..3], &data[3..8]] {
1561 dec.data.push(Bytes::copy_from_slice(chunk), dec.max_size);
1562 assert!(!dec.has_message());
1563 assert_eq!(dec.next(), Ok(None));
1564 }
1565 dec.data
1566 .push(Bytes::copy_from_slice(&data[8..]), dec.max_size);
1567 assert!(dec.has_message());
1568 assert_eq!(dec.next(), Ok(Some((0, Bytes::from_static(b"hello")))));
1569 assert_eq!(dec.next(), Ok(Some((1, Bytes::new()))));
1570 assert_eq!(dec.next(), Ok(Some((0, Bytes::from_static(b"world")))));
1571 assert_eq!(dec.next(), Ok(None));
1572 assert!(!dec.has_data());
1573 }
1574
1575 #[test]
1576 fn decoder_partial_tail() {
1577 let mut dec = decoder(1024);
1578 let mut data = frame(0, b"one");
1579 data.extend(&frame(0, b"two")[..6]);
1580 dec.data.push(Bytes::from(data), dec.max_size);
1581 assert_eq!(dec.next(), Ok(Some((0, Bytes::from_static(b"one")))));
1582 assert_eq!(dec.next(), Ok(None));
1583 assert!(dec.has_data());
1584 dec.data.push(Bytes::from_static(b"wo"), dec.max_size);
1585 assert_eq!(dec.next(), Ok(Some((0, Bytes::from_static(b"two")))));
1586 }
1587
1588 #[test]
1589 fn decoder_max_size() {
1590 let mut dec = decoder(4);
1591 dec.data
1593 .push(Bytes::from(frame(0, b"hello")[..5].to_vec()), dec.max_size);
1594 assert_eq!(dec.next(), Err(5));
1595
1596 let mut dec = decoder(5);
1597 dec.data.push(Bytes::from(frame(0, b"hello")), dec.max_size);
1598 assert_eq!(dec.next(), Ok(Some((0, Bytes::from_static(b"hello")))));
1599 }
1600
1601 #[test]
1602 fn server_error_from_stream_error() {
1603 let cases = [
1604 (
1605 StreamError::Closed,
1606 GrpcStatus::Internal,
1607 "Stream is closed",
1608 ),
1609 (
1610 StreamError::Reset(Reason::ENHANCE_YOUR_CALM),
1611 GrpcStatus::ResourceExhausted,
1612 "Stream reset with ENHANCE_YOUR_CALM",
1613 ),
1614 (
1615 StreamError::Disconnected,
1616 GrpcStatus::Unavailable,
1617 "Connection is closed",
1618 ),
1619 (
1620 StreamError::Status(GrpcStatus::NotFound, HeaderValue::from_static("gone")),
1621 GrpcStatus::NotFound,
1622 "gone",
1623 ),
1624 ];
1625 for (err, st, msg) in cases {
1626 let err = ServerError::from(err);
1627 assert_eq!((err.status, err.message.to_str().unwrap()), (st, msg));
1628 }
1629 }
1630
1631 #[test]
1632 fn client_error_from_stream_error() {
1633 let err = Error::<ClientError>::from(StreamError::Status(
1634 GrpcStatus::ResourceExhausted,
1635 HeaderValue::from_static("too big"),
1636 ));
1637 assert!(matches!(*err, ClientError::InvalidRequest(_)));
1638 assert_eq!(err.status(), GrpcStatus::ResourceExhausted);
1639 assert_eq!(err.message(), "too big");
1640
1641 let err = Error::<ClientError>::from(StreamError::Reset(Reason::REFUSED_STREAM));
1643 assert!(matches!(*err, ClientError::NotProcessed(_)));
1644 assert_eq!(err.status(), GrpcStatus::Unavailable);
1645
1646 let err = Error::<ClientError>::from(StreamError::Reset(Reason::CANCEL));
1647 assert!(matches!(*err, ClientError::StreamFailed(_)));
1648 assert_eq!(err.status(), GrpcStatus::Cancelled);
1649
1650 let err = Error::<ClientError>::from(StreamError::Disconnected);
1651 assert!(matches!(
1652 *err,
1653 ClientError::Connection {
1654 request_sent: true,
1655 ..
1656 }
1657 ));
1658 assert_eq!(err.status(), GrpcStatus::Unavailable);
1659
1660 let err = Error::<ClientError>::from(StreamError::Closed);
1661 assert!(matches!(*err, ClientError::StreamFailed(_)));
1662 assert_eq!(err.status(), GrpcStatus::Cancelled);
1663
1664 for err in [
1666 StreamError::Closed,
1667 StreamError::Reset(Reason::REFUSED_STREAM),
1668 StreamError::Disconnected,
1669 ] {
1670 let st = err.status();
1671 assert_eq!(Error::<ClientError>::from(err).status(), st);
1672 }
1673 }
1674
1675 #[test]
1676 fn unattached_send_stream() {
1677 let tx = SendStream::<()>::new();
1678 let cx = Context::from_waker(Waker::noop());
1679 assert!(tx.poll_ready(&cx).is_pending());
1681 assert_eq!(tx.shared.ready_wakers.borrow().len(), 1);
1682 assert!(!tx.is_closed());
1683 assert!(!tx.header(consts::GRPC_STATUS, HeaderValue::from_static("0")));
1685
1686 let tx2 = tx.clone();
1687 drop(tx);
1688 assert!(!tx2.is_closed());
1689 drop(tx2.clone());
1690 assert!(!tx2.is_closed());
1691
1692 let shared = tx2.shared.clone();
1694 drop(tx2);
1695 assert!(shared.flags.get().contains(Flags::CLOSED));
1696 assert!(shared.ready_wakers.borrow().is_empty());
1697 assert!(matches!(
1698 shared.poll_ready(&cx),
1699 Poll::Ready(Err(StreamError::Closed))
1700 ));
1701 }
1702
1703 #[ntex::test]
1704 async fn cancelled_send_stream() {
1705 let tx = SendStream::<()>::default();
1706 tx.cancel();
1707 assert!(tx.is_closed());
1708 assert!(matches!(tx.send(&()).await, Err(StreamError::Closed)));
1709 assert!(matches!(tx.ready().await, Err(StreamError::Closed)));
1710
1711 let tx = SendStream::<()>::new();
1712 tx.close();
1713 tx.close();
1714 assert!(tx.is_closed());
1715 assert!(matches!(tx.ready().await, Err(StreamError::Closed)));
1716 }
1717
1718 #[test]
1719 fn send_waiter_handoff() {
1720 let shared = SendShared::new();
1721 shared.insert(Flags::ATTACHED | Flags::BUSY);
1722 let active = SendGuard(&shared);
1723 let (first_count, first_waker) = WakeCount::waker();
1724 let (second_count, second_waker) = WakeCount::waker();
1725 let (observer_count, observer_waker) = WakeCount::waker();
1726 let mut first_cx = Context::from_waker(&first_waker);
1727 let mut second_cx = Context::from_waker(&second_waker);
1728 let observer_cx = Context::from_waker(&observer_waker);
1729 let mut first = std::pin::pin!(SendWait::new(&shared));
1730 let mut second = std::pin::pin!(SendWait::new(&shared));
1731 assert!(first.as_mut().poll(&mut first_cx).is_pending());
1732 assert!(second.as_mut().poll(&mut second_cx).is_pending());
1733 assert!(shared.poll_ready(&observer_cx).is_pending());
1734
1735 let (moved_count, moved_waker) = WakeCount::waker();
1737 let mut moved_cx = Context::from_waker(&moved_waker);
1738 assert!(first.as_mut().poll(&mut moved_cx).is_pending());
1739 drop(active);
1740 assert_eq!(first_count.take(), 0);
1741 assert_eq!(moved_count.take(), 1);
1742 assert_eq!(second_count.take(), 0);
1743 assert_eq!(observer_count.take(), 1);
1744
1745 assert!(matches!(
1747 shared.poll_ready(&observer_cx),
1748 Poll::Ready(Ok(()))
1749 ));
1750 let (last_count, last_waker) = WakeCount::waker();
1751 let mut last_cx = Context::from_waker(&last_waker);
1752 let mut last = std::pin::pin!(SendWait::new(&shared));
1753 assert!(last.as_mut().poll(&mut last_cx).is_pending());
1754 let Poll::Ready(Ok(first)) = first.as_mut().poll(&mut moved_cx) else {
1755 panic!("selected sender did not acquire the slot");
1756 };
1757 assert!(second.as_mut().poll(&mut second_cx).is_pending());
1758 drop(first);
1759 assert_eq!(second_count.take(), 1);
1760 assert_eq!(last_count.take(), 0);
1761 let Poll::Ready(Ok(second)) = second.as_mut().poll(&mut second_cx) else {
1762 panic!("second sender did not acquire the slot");
1763 };
1764 drop(second);
1765 assert_eq!(last_count.take(), 1);
1766 let Poll::Ready(Ok(last)) = last.as_mut().poll(&mut last_cx) else {
1767 panic!("last sender did not acquire the slot");
1768 };
1769 drop(last);
1770 assert!(matches!(
1771 shared.poll_ready(&observer_cx),
1772 Poll::Ready(Ok(()))
1773 ));
1774 }
1775
1776 #[test]
1777 fn cancelled_send_waiters() {
1778 for selected in [false, true] {
1779 let shared = SendShared::new();
1780 shared.insert(Flags::ATTACHED | Flags::BUSY);
1781 let active = SendGuard(&shared);
1782 let (count, waker) = WakeCount::waker();
1783 let mut cx = Context::from_waker(&waker);
1784 let mut first = Box::pin(SendWait::new(&shared));
1785 let mut middle = Box::pin(SendWait::new(&shared));
1786 let mut last = Box::pin(SendWait::new(&shared));
1787 assert!(first.as_mut().poll(&mut cx).is_pending());
1789 assert!(middle.as_mut().poll(&mut cx).is_pending());
1790 assert!(last.as_mut().poll(&mut cx).is_pending());
1791 drop(middle);
1792 assert_eq!(count.take(), 0);
1793 if selected {
1794 drop(active);
1795 assert_eq!(count.take(), 1);
1796 drop(first);
1797 } else {
1798 drop(first);
1799 assert_eq!(count.take(), 0);
1800 drop(active);
1801 }
1802 assert_eq!(count.take(), 1);
1803 let Poll::Ready(Ok(last)) = last.as_mut().poll(&mut cx) else {
1804 panic!("cancelled sender stranded its successor");
1805 };
1806 drop(last);
1807 assert!(shared.send_waiters.borrow().is_empty());
1808 assert!(matches!(shared.poll_ready(&cx), Poll::Ready(Ok(()))));
1809 }
1810 }
1811
1812 #[test]
1813 fn send_waiters_terminal_wakeups() {
1814 let transitions: [fn(&SendShared); 5] = [
1815 SendShared::close,
1816 SendShared::cancel,
1817 |shared| {
1818 shared.finish();
1819 },
1820 |shared| {
1821 shared.set_error(StreamError::Disconnected);
1822 drop(SendGuard(shared));
1823 },
1824 |shared| {
1825 shared.insert(Flags::SENDING);
1826 drop(SendGuard(shared));
1827 },
1828 ];
1829 for attached in [false, true] {
1830 for transition in transitions {
1831 let shared = SendShared::new();
1832 if attached {
1833 shared.insert(Flags::ATTACHED | Flags::BUSY);
1834 }
1835 let (first_count, first_waker) = WakeCount::waker();
1836 let (second_count, second_waker) = WakeCount::waker();
1837 let (observer_count, observer_waker) = WakeCount::waker();
1838 let mut first_cx = Context::from_waker(&first_waker);
1839 let mut second_cx = Context::from_waker(&second_waker);
1840 let observer_cx = Context::from_waker(&observer_waker);
1841 let mut first = std::pin::pin!(shared.send(&()));
1842 let mut second = std::pin::pin!(shared.send(&()));
1843 assert!(first.as_mut().poll(&mut first_cx).is_pending());
1844 assert!(second.as_mut().poll(&mut second_cx).is_pending());
1845 assert!(shared.poll_ready(&observer_cx).is_pending());
1846
1847 transition(&shared);
1848 assert_eq!(first_count.take(), 1);
1849 assert_eq!(second_count.take(), 1);
1850 assert_eq!(observer_count.take(), 1);
1851 assert!(matches!(
1852 first.as_mut().poll(&mut first_cx),
1853 Poll::Ready(Err(_))
1854 ));
1855 assert!(matches!(
1856 second.as_mut().poll(&mut second_cx),
1857 Poll::Ready(Err(_))
1858 ));
1859 assert!(matches!(
1860 shared.poll_ready(&observer_cx),
1861 Poll::Ready(Err(_))
1862 ));
1863 assert!(shared.send_waiters.borrow().is_empty());
1864 }
1865 }
1866 }
1867
1868 #[ntex::test]
1869 async fn server_recv_end() {
1870 let mut rx = RecvStream::<()>::server(ServerRecv::new(true), HeaderMap::new(), 64, None);
1872 assert_eq!(rx.status(), None);
1873 assert!(rx.recv().await.is_none());
1874 assert_eq!(rx.status(), Some(GrpcStatus::Ok));
1875 assert!(rx.recv().await.is_none());
1876
1877 let recv = ServerRecv::new(false);
1879 let mut rx = RecvStream::<()>::server(recv.clone(), HeaderMap::new(), 64, None);
1880 recv.push_end(Err(StreamError::Reset(Reason::CANCEL)));
1881 recv.push_end(Ok(HeaderMap::new()));
1882 assert!(recv.is_ended());
1883 assert!(matches!(
1884 rx.recv().await,
1885 Some(Err(StreamError::Reset(Reason::CANCEL)))
1886 ));
1887 assert_eq!(rx.status(), Some(GrpcStatus::Cancelled));
1888 assert!(rx.recv().await.is_none());
1889
1890 let mut trailers = HeaderMap::new();
1891 trailers.insert(consts::GRPC_STATUS, HeaderValue::from_static("0"));
1892 let recv = ServerRecv::new(false);
1893 let mut rx = RecvStream::<()>::server(recv.clone(), HeaderMap::new(), 64, None);
1894 recv.push_end(Ok(trailers.clone()));
1895 assert!(rx.recv().await.is_none());
1896 assert_eq!(rx.trailers(), &trailers);
1897 }
1898
1899 #[ntex::test]
1900 async fn server_recv_failure_clears_buffered_data() {
1901 use crate::google_types::BytesValue;
1902
1903 for (invalid, status) in [
1904 (frame(2, b""), GrpcStatus::Internal),
1905 (frame(0, b"\xff"), GrpcStatus::InvalidArgument),
1906 (frame(0, &[0; 65]), GrpcStatus::ResourceExhausted),
1907 ] {
1908 let mut data = frame(0, b"");
1909 data.extend(invalid);
1910 data.extend(frame(0, b""));
1911 let mut rx = RecvStream::<BytesValue>::server(
1912 ServerRecv::new(true),
1913 HeaderMap::new(),
1914 64,
1915 None,
1916 );
1917 rx.decoder.data.push(Bytes::from(data), 64);
1918
1919 assert_eq!(rx.recv().await.unwrap().unwrap(), BytesValue::default());
1920 let err = rx.recv().await.unwrap().unwrap_err();
1921 assert_eq!(err.status(), status);
1922 assert_eq!(rx.status(), Some(status));
1923 assert!(!rx.decoder.has_data());
1924 assert!(rx.recv().await.is_none());
1925 }
1926 }
1927
1928 #[ntex::test]
1929 async fn server_recv_single() {
1930 let rx = RecvStream::<()>::server(ServerRecv::new(true), HeaderMap::new(), 64, None);
1931 let err = rx.single().await.unwrap_err();
1932 assert_eq!(err.status, GrpcStatus::InvalidArgument);
1933
1934 let recv = ServerRecv::new(false);
1935 let rx = RecvStream::<()>::server(recv.clone(), HeaderMap::new(), 64, None);
1936 recv.push_end(Err(StreamError::Disconnected));
1937 let err = rx.single().await.unwrap_err();
1938 assert_eq!(err.status, GrpcStatus::Unavailable);
1939 }
1940
1941 #[test]
1942 fn call_streams_drop_closes_recv() {
1943 let recv = ServerRecv::new(false);
1944 let streams = CallStreams::new(
1945 recv.clone(),
1946 HeaderMap::new(),
1947 64,
1948 None,
1949 Rc::new(SendShared::new()),
1950 );
1951 drop(streams);
1952
1953 recv.push_end(Ok(HeaderMap::new()));
1955 assert!(recv.is_ended());
1956 let mut trailers = HeaderMap::new();
1957 let cx = Context::from_waker(Waker::noop());
1958 assert!(matches!(
1959 recv.poll(&mut trailers, &cx),
1960 Poll::Ready(Item::End(Err(StreamError::Closed)))
1961 ));
1962 }
1963
1964 #[test]
1965 fn call_streams_take() {
1966 let recv = ServerRecv::new(false);
1967 let mut streams = CallStreams::new(
1968 recv.clone(),
1969 HeaderMap::new(),
1970 64,
1971 None,
1972 Rc::new(SendShared::new()),
1973 );
1974 let rx = streams.recv::<()>();
1975 let _tx = streams.send::<()>();
1976 drop(streams);
1977 recv.push_end(Ok(HeaderMap::new()));
1979 let cx = Context::from_waker(Waker::noop());
1980 let mut trailers = HeaderMap::new();
1981 assert!(matches!(
1982 recv.poll(&mut trailers, &cx),
1983 Poll::Ready(Item::End(Ok(())))
1984 ));
1985 drop(rx);
1986 }
1987}