ntex_grpc/client/
stream.rs

1//! Streaming calls.
2use std::{fmt, marker::PhantomData, rc::Rc, time::Instant};
3
4use ntex_error::Error;
5use ntex_h2::{self as h2, OperationError, frame::Reason};
6use ntex_http::{HeaderMap, HeaderName, HeaderValue, Method, error::Error as HttpError};
7use ntex_util::time;
8
9use super::request::{RequestContext, RequestContextBuilder, Response};
10use super::transport::{self as tr, with_deadline};
11use super::{CallFailure, Client, ClientError};
12#[cfg(feature = "compression")]
13use crate::Compression;
14use crate::stream::{ClientRecv, SendConfig, SendShared};
15use crate::{GrpcStatus, RecvStream, SendStream, service::MethodDef, utils::Enc};
16
17/// Opens the streams of streaming calls of method `M`.
18///
19/// Generated clients go through this trait for client-streaming,
20/// server-streaming and bidi calls. [`Client`], [`ntex_h2::client::Client`]
21/// and [`ntex_h2::client::SimpleClient`] implement it. A custom transport
22/// wraps one of them, a [`StreamCall`] is only created by them.
23pub trait StreamTransport<M: MethodDef, St = ()> {
24    /// Opens the h2 stream of a call.
25    ///
26    /// `ctx` holds the headers and timeout set on the [`StreamRequest`]. The
27    /// timeout covers the whole call, the time spent opening the stream too.
28    /// A wrapper must adjust the context before delegating to a built-in
29    /// transport, which sends the request headers.
30    async fn open(
31        &self,
32        ctx: &mut RequestContext<'_, St>,
33    ) -> Result<StreamCall, Error<ClientError>>;
34}
35
36/// The h2 stream of a streaming call.
37pub struct StreamCall {
38    tx: h2::client::SendStream,
39    rx: h2::client::RecvStream,
40    deadline: Option<Instant>,
41    service: &'static str,
42}
43
44impl StreamCall {
45    /// The `SendStream` and `RecvStream` of the call.
46    fn attach<Out, St>(
47        self,
48        shared: &Rc<SendShared>,
49        ctx: &RequestContext<'_, St>,
50    ) -> Result<RecvStream<Out>, Error<ClientError>> {
51        let StreamCall {
52            tx,
53            rx,
54            deadline,
55            service,
56        } = self;
57        let cfg = SendConfig {
58            max_size: ctx.get_max_send_message_size(),
59            enc: send_enc(ctx),
60            deadline,
61            service,
62        };
63        shared
64            .attach_client(tx, cfg)
65            .map_err(|e| Error::<ClientError>::from(e).with_service(service))?;
66        Ok(RecvStream::client(
67            ClientRecv::new(rx, deadline, service, Some(shared.clone())),
68            ctx.get_max_message_size(),
69        ))
70    }
71}
72
73impl fmt::Debug for StreamCall {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        f.debug_struct("StreamCall")
76            .field("deadline", &self.deadline)
77            .field("service", &self.service)
78            .finish_non_exhaustive()
79    }
80}
81
82#[cfg(feature = "compression")]
83fn send_enc<St>(ctx: &RequestContext<'_, St>) -> Option<Enc> {
84    ctx.get_compression()
85}
86
87#[cfg(not(feature = "compression"))]
88fn send_enc<St>(_: &RequestContext<'_, St>) -> Option<Enc> {
89    None
90}
91
92impl<M: MethodDef, St> StreamTransport<M, St> for Client {
93    #[inline]
94    async fn open(
95        &self,
96        ctx: &mut RequestContext<'_, St>,
97    ) -> Result<StreamCall, Error<ClientError>> {
98        ctx.message_size_defaults(self.max_message_size, self.max_send_message_size);
99        StreamTransport::<M, St>::open(&self.pool, ctx).await
100    }
101}
102
103impl<M: MethodDef, St> StreamTransport<M, St> for h2::client::Client {
104    /// The timeout covers waiting for a connection too.
105    async fn open(
106        &self,
107        ctx: &mut RequestContext<'_, St>,
108    ) -> Result<StreamCall, Error<ClientError>> {
109        with_deadline(ctx.get_timeout(), async |deadline| {
110            let client = self.client().await.map_err(|e| e.map(tr::connect_error))?;
111            open_stream::<M, St>(&client, ctx, deadline).await
112        })
113        .await
114    }
115}
116
117impl<M: MethodDef, St> StreamTransport<M, St> for h2::client::SimpleClient {
118    async fn open(
119        &self,
120        ctx: &mut RequestContext<'_, St>,
121    ) -> Result<StreamCall, Error<ClientError>> {
122        with_deadline(ctx.get_timeout(), async |deadline| {
123            open_stream::<M, St>(self, ctx, deadline).await
124        })
125        .await
126        .map_err(|e| e.with_service(self.service()))
127    }
128}
129
130async fn open_stream<M: MethodDef, St>(
131    client: &h2::client::SimpleClient,
132    ctx: &mut RequestContext<'_, St>,
133    deadline: Option<Instant>,
134) -> Result<StreamCall, Error<ClientError>> {
135    if let Some(err) = ctx.take_error() {
136        return Err(Error::from(tr::request_error(err)).with_service(client.service()));
137    }
138    let hdrs = tr::request_headers(ctx);
139    let (tx, rx) = client
140        .send(Method::POST, M::PATH, hdrs, false)
141        .await
142        .map_err(|e| {
143            tr::operation_error(e, false, HeaderMap::new(), deadline)
144                .with_service(client.service())
145        })?;
146    if ctx.get_disconnect_on_drop() {
147        tx.disconnect_on_drop();
148    }
149    Ok(StreamCall {
150        tx,
151        rx,
152        deadline,
153        service: client.service(),
154    })
155}
156
157const STREAM_USED: &str = "SendStream is used by another call already";
158
159/// The kind of a bidi call, see [`StreamRequest::open()`].
160#[derive(Debug)]
161pub struct Bidi;
162
163/// The kind of a server-streaming call, holds the request message.
164pub struct ServerStreaming<'a, In>(&'a In);
165
166impl<In> fmt::Debug for ServerStreaming<'_, In> {
167    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
168        f.debug_struct("ServerStreaming").finish_non_exhaustive()
169    }
170}
171
172/// The kind of a client-streaming call, holds the request stream.
173pub struct ClientStreaming(Rc<SendShared>);
174
175impl fmt::Debug for ClientStreaming {
176    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
177        f.debug_struct("ClientStreaming").finish_non_exhaustive()
178    }
179}
180
181/// A streaming call that is ready to be opened.
182///
183/// Generated client methods return it for client-streaming,
184/// server-streaming and bidi methods. Set headers or a timeout if you need
185/// them, then open the call:
186///
187/// * bidi: `open()` returns the request and the response streams.
188/// * server-streaming: `open()` sends the request message and returns the
189///   response stream.
190/// * client-streaming: the request stream is passed to the method, `send()`
191///   waits for the reply while the caller sends messages on the stream.
192///
193/// The timeout covers the whole call. Once it runs out, the streams fail
194/// with [`ClientError::DeadlineExceeded`] and the call is reset.
195pub struct StreamRequest<'a, T, M, K> {
196    transport: &'a T,
197    kind: K,
198    ctx: RequestContextBuilder,
199    _m: PhantomData<fn(M)>,
200}
201
202impl<'a, T, M, K> StreamRequest<'a, T, M, K> {
203    fn with_kind(transport: &'a T, kind: K) -> Self {
204        StreamRequest {
205            transport,
206            kind,
207            ctx: RequestContextBuilder::new(),
208            _m: PhantomData,
209        }
210    }
211
212    /// Set a request header, replacing any existing value for the same name.
213    ///
214    /// See [`Request::header()`](super::Request::header).
215    pub fn header<K1, V>(&mut self, key: K1, value: V) -> &mut Self
216    where
217        HeaderName: TryFrom<K1>,
218        HeaderValue: TryFrom<V>,
219        <HeaderName as TryFrom<K1>>::Error: Into<HttpError>,
220        <HeaderValue as TryFrom<V>>::Error: Into<HttpError>,
221    {
222        self.ctx.header(key, value);
223        self
224    }
225
226    /// Add a request header, keeping existing values for the same name.
227    ///
228    /// See [`Request::append_header()`](super::Request::append_header).
229    pub fn append_header<K1, V>(&mut self, key: K1, value: V) -> &mut Self
230    where
231        HeaderName: TryFrom<K1>,
232        HeaderValue: TryFrom<V>,
233        <HeaderName as TryFrom<K1>>::Error: Into<HttpError>,
234        <HeaderValue as TryFrom<V>>::Error: Into<HttpError>,
235    {
236        self.ctx.append_header(key, value);
237        self
238    }
239
240    /// Set the max duration of the whole call.
241    ///
242    /// The duration is sent in the `grpc-timeout` header. Built-in
243    /// transports reset the call once it runs out, the streams fail with
244    /// [`ClientError::DeadlineExceeded`]. A zero timeout fails the call
245    /// without sending it.
246    pub fn timeout<U>(&mut self, timeout: U) -> &mut Self
247    where
248        std::time::Duration: From<U>,
249    {
250        self.ctx.timeout(timeout);
251        self
252    }
253
254    /// Set the size limit of each response message.
255    ///
256    /// Overrides the [`Client`] default for this call. Without a client
257    /// default or an override, the limit is 4 MiB.
258    ///
259    /// A larger message fails the call with
260    /// [`GrpcStatus::ResourceExhausted`].
261    pub fn max_message_size(&mut self, size: usize) -> &mut Self {
262        self.ctx.max_message_size(size);
263        self
264    }
265
266    /// Set the size limit of each request message.
267    ///
268    /// Overrides the [`Client`] default for this call. Without a client
269    /// default or an override, the limit is 2 GiB - 1.
270    ///
271    /// A larger message is not sent, sending it fails with
272    /// [`GrpcStatus::ResourceExhausted`].
273    pub fn max_send_message_size(&mut self, size: usize) -> &mut Self {
274        self.ctx.max_send_message_size(size);
275        self
276    }
277
278    /// Compress the request messages with gzip or zstd.
279    ///
280    /// Requires the `compression` feature.
281    #[cfg(feature = "compression")]
282    pub fn compression(&mut self, compression: Compression) -> &mut Self {
283        self.ctx.compression(compression);
284        self
285    }
286
287    /// Disconnect the connection if the call is dropped before its end.
288    pub fn disconnect_on_drop(&mut self) -> &mut Self {
289        self.ctx.disconnect_on_drop();
290        self
291    }
292}
293
294impl<'a, T, M: MethodDef> StreamRequest<'a, T, M, Bidi> {
295    /// Create a bidi call through `transport`.
296    ///
297    /// Generated clients do this for you.
298    pub fn bidi(transport: &'a T) -> Self {
299        Self::with_kind(transport, Bidi)
300    }
301
302    /// Open the call, returns the request and the response streams.
303    ///
304    /// Close the request stream, or drop it, once all messages are sent.
305    /// Dropping the response stream before its end cancels the call.
306    pub async fn open(
307        self,
308    ) -> Result<(SendStream<M::Input>, RecvStream<M::Output>), Error<ClientError>>
309    where
310        T: StreamTransport<M, ()>,
311    {
312        self.open_with(&()).await
313    }
314
315    /// Open the call, carrying `request_state` through to the transport.
316    pub async fn open_with<St>(
317        self,
318        request_state: &St,
319    ) -> Result<(SendStream<M::Input>, RecvStream<M::Output>), Error<ClientError>>
320    where
321        T: StreamTransport<M, St>,
322    {
323        let mut ctx = self.ctx.build(request_state);
324        let call = self.transport.open(&mut ctx).await?;
325        let tx = SendStream::new();
326        let rx = call.attach::<M::Output, St>(tx.shared(), &ctx)?;
327        Ok((tx, rx))
328    }
329}
330
331impl<'a, T, M: MethodDef> StreamRequest<'a, T, M, ServerStreaming<'a, M::Input>> {
332    /// Create a server-streaming call that sends `input` through `transport`.
333    ///
334    /// Generated clients do this for you.
335    pub fn server_streaming(transport: &'a T, input: &'a M::Input) -> Self {
336        Self::with_kind(transport, ServerStreaming(input))
337    }
338
339    /// Send the request message, returns the response stream.
340    ///
341    /// The transport's final size limit and compression settings are used to
342    /// encode the request. If encoding fails, the opened call is cancelled
343    /// without sending its body.
344    ///
345    /// Dropping the response stream before its end cancels the call.
346    pub async fn open(self) -> Result<RecvStream<M::Output>, Error<ClientError>>
347    where
348        T: StreamTransport<M, ()>,
349    {
350        self.open_with(&()).await
351    }
352
353    /// Send the request message, carrying `request_state` through to the
354    /// transport.
355    pub async fn open_with<St>(
356        self,
357        request_state: &St,
358    ) -> Result<RecvStream<M::Output>, Error<ClientError>>
359    where
360        T: StreamTransport<M, St>,
361    {
362        let mut ctx = self.ctx.build(request_state);
363        let StreamCall {
364            tx,
365            rx,
366            deadline,
367            service,
368        } = self.transport.open(&mut ctx).await?;
369
370        let deadline_exceeded = || {
371            tx.reset(Reason::CANCEL);
372            tr::deadline_exceeded(tr::TIMEOUT_RAN_OUT, HeaderMap::new()).with_service(service)
373        };
374        let send = async {
375            let buf = tr::encode_message(self.kind.0, &ctx).await.map_err(|err| {
376                tx.reset(Reason::CANCEL);
377                Error::from(err).with_service(service)
378            })?;
379            // Encoding can finish without yielding to the timeout.
380            if deadline.is_some_and(|d| Instant::now() >= d) {
381                return Err(deadline_exceeded());
382            }
383            if let Err(err) = tx.send_pages(buf, true).await {
384                // the server can reply and reset the stream before it reads the
385                // whole request, the reply or the reset is in the receive queue
386                if !matches!(
387                    *err,
388                    OperationError::RemoteReset(_) | OperationError::Closed(Some(_))
389                ) {
390                    return Err(tr::operation_error(err, true, HeaderMap::new(), deadline)
391                        .with_service(service));
392                }
393            }
394            Ok(())
395        };
396        if let Some(deadline) = deadline {
397            let remaining = deadline.saturating_duration_since(Instant::now());
398            if remaining.is_zero() {
399                return Err(deadline_exceeded());
400            }
401            match time::timeout(remaining, send).await {
402                Ok(res) => res?,
403                Err(()) => return Err(deadline_exceeded()),
404            }
405        } else {
406            send.await?;
407        }
408        Ok(RecvStream::client(
409            ClientRecv::new(rx, deadline, service, None),
410            ctx.get_max_message_size(),
411        ))
412    }
413}
414
415impl<'a, T, M: MethodDef> StreamRequest<'a, T, M, ClientStreaming> {
416    /// Create a client-streaming call through `transport`, the request
417    /// messages are sent on `tx`.
418    ///
419    /// Generated clients do this for you.
420    pub fn client_streaming(transport: &'a T, tx: &SendStream<M::Input>) -> Self {
421        Self::with_kind(transport, ClientStreaming(tx.shared().clone()))
422    }
423
424    /// Open the call and wait for the reply.
425    ///
426    /// Send the request messages on the stream passed to the method while
427    /// this future runs, then close the stream or drop it. Messages sent
428    /// before the call is opened wait for it. Once this future is polled,
429    /// the stream belongs to this call, even while opening. Reusing it for
430    /// another call fails with [`ClientError::InvalidRequest`].
431    ///
432    /// An opening error also fails the sender with the same client error.
433    /// Dropping this future while opening closes the sender and wakes
434    /// waiting sends.
435    ///
436    /// ```rust,ignore
437    /// let tx = ntex_grpc::SendStream::new();
438    /// let (res, _) = futures::join!(client.record_route(&tx).send(), async {
439    ///     for point in points {
440    ///         tx.send(&point).await?;
441    ///     }
442    ///     tx.close();
443    ///     Ok::<_, ntex_grpc::StreamError>(())
444    /// });
445    /// ```
446    pub async fn send(self) -> Result<Response<M>, Error<ClientError>>
447    where
448        T: StreamTransport<M, ()>,
449    {
450        self.send_with(&()).await
451    }
452
453    /// Open the call and wait for the reply, carrying `request_state`
454    /// through to the transport.
455    pub async fn send_with<St>(self, request_state: &St) -> Result<Response<M>, Error<ClientError>>
456    where
457        T: StreamTransport<M, St>,
458    {
459        let shared = self.kind.0;
460        let opening = shared.start_open().ok_or_else(|| {
461            Error::from(ClientError::InvalidRequest(CallFailure::new(
462                GrpcStatus::Internal,
463                STREAM_USED,
464            )))
465        })?;
466        let mut ctx = self.ctx.build(request_state);
467        let call = match self.transport.open(&mut ctx).await {
468            Ok(call) => call,
469            Err(err) => {
470                opening.fail(err.clone());
471                return Err(err);
472            }
473        };
474        let service = call.service;
475        let mut rx = call.attach::<M::Output, St>(&shared, &ctx)?;
476        drop(opening);
477
478        let output = match rx.recv().await {
479            Some(Ok(output)) => output,
480            Some(Err(err)) => return Err(err.into()),
481            None => {
482                let failure = CallFailure::new(GrpcStatus::Internal, tr::NO_MESSAGE)
483                    .with_headers(rx.trailers().clone());
484                let err = Error::from(ClientError::InvalidResponse(failure));
485                return Err(err.with_service(service));
486            }
487        };
488        match rx.recv().await {
489            None => {}
490            Some(Err(err)) => return Err(err.into()),
491            Some(Ok(_)) => {
492                let hdrs = rx.headers().clone();
493                let err = tr::invalid_response(GrpcStatus::Internal, tr::EXTRA_DATA, hdrs);
494                return Err(err.with_service(service));
495            }
496        }
497        Ok(Response {
498            output,
499            headers: rx.headers().clone(),
500            trailers: rx.trailers().clone(),
501            req_size: shared.sent(),
502            res_size: rx.res_size(),
503        })
504    }
505}
506
507impl<T, M, K: fmt::Debug> fmt::Debug for StreamRequest<'_, T, M, K> {
508    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
509        f.debug_struct("StreamRequest")
510            .field("kind", &self.kind)
511            .finish_non_exhaustive()
512    }
513}
ntex_grpc/client/stream.rs

ntex_grpc/client/
stream.rs
