mysten_network/callback/
service.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use super::{CallbackLayer, MakeCallbackHandler, ResponseBody, ResponseFuture};
5use http::{Request, Response};
6use std::task::{Context, Poll};
7use tower::Service;
8
9/// Middleware that adds callbacks to a [`Service`].
10///
11/// See the [module docs](crate::callback) for an example.
12///
13/// [`Service`]: tower::Service
14#[derive(Debug, Clone, Copy)]
15pub struct Callback<S, M> {
16    pub(crate) inner: S,
17    pub(crate) make_callback_handler: M,
18}
19
20impl<S, M> Callback<S, M> {
21    /// Create a new [`Callback`].
22    pub fn new(inner: S, make_callback_handler: M) -> Self {
23        Self {
24            inner,
25            make_callback_handler,
26        }
27    }
28
29    /// Returns a new [`Layer`] that wraps services with a [`CallbackLayer`] middleware.
30    ///
31    /// [`Layer`]: tower::layer::Layer
32    pub fn layer(make_handler: M) -> CallbackLayer<M>
33    where
34        M: MakeCallbackHandler,
35    {
36        CallbackLayer::new(make_handler)
37    }
38
39    /// Gets a reference to the underlying service.
40    pub fn inner(&self) -> &S {
41        &self.inner
42    }
43
44    /// Gets a mutable reference to the underlying service.
45    pub fn inner_mut(&mut self) -> &mut S {
46        &mut self.inner
47    }
48
49    /// Consumes `self`, returning the underlying service.
50    pub fn into_inner(self) -> S {
51        self.inner
52    }
53}
54
55impl<S, M, RequestBody, ResponseBodyT> Service<Request<RequestBody>> for Callback<S, M>
56where
57    S: Service<
58            Request<RequestBody>,
59            Response = Response<ResponseBodyT>,
60            Error: std::fmt::Display + 'static,
61        >,
62    M: MakeCallbackHandler,
63    RequestBody: http_body::Body<Error: std::fmt::Display + 'static>,
64    ResponseBodyT: http_body::Body<Error: std::fmt::Display + 'static>,
65{
66    type Response = Response<ResponseBody<ResponseBodyT, M::Handler>>;
67    type Error = S::Error;
68    type Future = ResponseFuture<S::Future, M::Handler>;
69
70    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
71        self.inner.poll_ready(cx)
72    }
73
74    fn call(&mut self, request: Request<RequestBody>) -> Self::Future {
75        let (head, body) = request.into_parts();
76        let handler = self.make_callback_handler.make_handler(&head);
77        let request = Request::from_parts(head, body);
78
79        ResponseFuture {
80            inner: self.inner.call(request),
81            handler: Some(handler),
82        }
83    }
84}