util: add Wrap

tower/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Unreleased
 
+### Added
+
+- **util**: Add `Wrap` for wrapping a `Service` with pre/post functions
+
 ### Changed
 
 - **util**: Removed deprecated `ServiceExt::ready_and` method and `ReadyAnd`

tower/src/builder/mod.rs

@@ -808,6 +808,76 @@ impl<L> ServiceBuilder<L> {
     {
         self.layer(crate::util::BoxCloneService::layer())
     }
+
+    /// Wraps the inner service with async pre/post functions.
+    ///
+    /// The `pre` function is any function that looks like
+    /// <code>async [FnMut]\(Request) -> [Result]<(SRequest, T), [Result]<Response, Err>></code>
+    /// where `Request` is the request given to the `Wrap` service.
+    ///
+    /// * If this function outputs <code>[Ok]\((SRequest, T))</code> the `SRequest` is passed to
+    ///   the inner service and the `T` is retained as shared state. When the inner service outputs
+    ///   a result, this result is passed to `post` along with the shared `T`.
+    /// * If this function returns <code>[Err]\(result)</code> the result is output by the `Wrap`
+    ///   service without calling the inner service or the `post` function.
+    ///
+    /// The `post` function is any function that looks like
+    /// <code>async [FnOnce]\(Res, T) -> [Result]<Response, Err></code> where `Res` is the result
+    /// returned by the inner service, `T` is the shared state provided by `pre`, and `Response`
+    /// and `Err` match the types used in `pre`. The returned [`Result`] is the overall result from
+    /// the wrapped service.
+    /// 
+    /// See the documentation for the [`decorate` combinator][] for examples.
+    ///
+    /// See also [`wrap`](Self::wrap) for when you just need to share state across the inner
+    /// service and don't need asynchronous functions or short-circuiting.
+    /// 
+    /// [`decorate` combinator]: crate::util::ServiceExt::decorate
+    #[cfg(feature = "util")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "util")))]
+    pub fn decorate<Pre, Post>(
+        self,
+        pre: Pre,
+        post: Post,
+    ) -> ServiceBuilder<Stack<crate::util::WrapLayer<Pre, Post>, L>>
+    where
+        Self: Sized,
+    {
+        self.layer(crate::util::WrapLayer::decorate(pre, post))
+    }
+
+    /// Wraps the inner service with a synchronous pre function that returns a post function.
+    ///
+    /// The given function is any function that looks like
+    /// <code>[FnMut]\(&mut Request) -> [FnOnce]\(Response) -> T</code> where `Request` is the
+    /// request given to the `Wrap` service, `Response` is the response returned from the inner
+    /// service, and `T` is the response returned from the `Wrap` service. If the inner service
+    /// returns an error the error is output directly without being given to the post function.
+    /// 
+    /// See the documentation for the [`wrap` combinator][] for examples.
+    /// 
+    /// [`wrap` combinator]: crate::util::ServiceExt::wrap
+    #[cfg(feature = "util")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "util")))]
+    pub fn wrap<F, F2, Request, Response, T, E>(
+        self,
+        f: F,
+    ) -> ServiceBuilder<
+        Stack<
+            crate::util::WrapLayer<
+                crate::util::helper::Pre<Request, T, E, F>,
+                crate::util::helper::Post<F2, E>,
+            >,
+            L,
+        >,
+    >
+    where
+        Self: Sized,
+        F: FnMut(&mut Request) -> F2,
+        F2: FnOnce(Response) -> T,
+    {
+        self.layer(crate::util::WrapLayer::new(f))
+    }
 }
 
 impl<L: fmt::Debug> fmt::Debug for ServiceBuilder<L> {

tower/src/util/mod.rs

@@ -18,6 +18,7 @@ mod optional;
 mod ready;
 mod service_fn;
 mod then;
+mod wrap;
 
 pub use self::{
     and_then::{AndThen, AndThenLayer},
@@ -35,6 +36,7 @@ pub use self::{
     ready::{Ready, ReadyOneshot},
     service_fn::{service_fn, ServiceFn},
     then::{Then, ThenLayer},
+    wrap::{Wrap, WrapLayer},
 };
 
 pub use self::call_all::{CallAll, CallAllUnordered};
@@ -58,6 +60,13 @@ pub mod future {
     pub use super::map_result::MapResultFuture;
     pub use super::optional::future as optional;
     pub use super::then::ThenFuture;
+    pub use super::wrap::{WrapFuture, WrapPostFuture, WrapPreFuture};
+}
+
+pub mod helper {
+    //! Helper types
+
+    pub use super::wrap::{Post, PostFn, Pre, PreFn};
 }
 
 /// An extension trait for `Service`s that provides a variety of convenient
@@ -1036,6 +1045,145 @@ pub trait ServiceExt<Request>: tower_service::Service<Request> {
     {
         BoxCloneService::new(self)
     }
+
+    /// Wraps this service with async pre/post functions.
+    ///
+    /// The `pre` function is any function that looks like
+    /// <code>async [FnMut]\(Request) -> [Result]<(SRequest, T), [Result]<Response, Err>></code>
+    /// where `Request` is the request given to the `Wrap` service.
+    ///
+    /// * If this function outputs <code>[Ok]\((SRequest, T))</code> the `SRequest` is passed to
+    ///   the inner service and the `T` is retained as shared state. When the inner service outputs
+    ///   a result, this result is passed to `post` along with the shared `T`.
+    /// * If this function returns <code>[Err]\(result)</code> the result is output by the `Wrap`
+    ///   service without calling the inner service or the `post` function.
+    ///
+    /// The `post` function is any function that looks like
+    /// <code>async [FnOnce]\(Res, T) -> [Result]<Response, Err></code> where `Res` is the result
+    /// returned by the inner service, `T` is the shared state provided by `pre`, and `Response`
+    /// and `Err` match the types used in `pre`. The returned [`Result`] is the overall result from
+    /// the wrapped service.
+    ///
+    /// See also [`wrap`](Self::wrap) for when you just need to share state across the inner
+    /// service and don't need asynchronous functions or short-circuiting.
+    ///
+    /// # Examples
+    ///
+    /// Adding a cache in front of a service:
+    ///
+    /// ```rust
+    /// use tower::{Service, ServiceExt};
+    /// use std::collections::HashMap;
+    /// use std::sync::{Arc, Mutex};
+    /// # #[derive(Clone, PartialEq, Eq, Hash)] struct Request;
+    /// # impl Request { fn new(_: &str) -> Self { Request }}
+    /// # #[derive(Clone)] struct Response;
+    ///
+    /// # fn main() {
+    /// #     async {
+    /// // A service returning Result<Response, Error>
+    /// let service = /* ... */
+    /// # tower::service_fn(|_: Request| async { Ok::<_, ()>(Response) });
+    ///
+    /// // Wrap the service in a cache for successful responses
+    /// let cache = Arc::new(Mutex::new(HashMap::new()));
+    /// let mut new_service = service.decorate(
+    ///     {
+    ///         let cache = cache.clone();
+    ///         move |req: Request| std::future::ready(
+    ///             cache.lock().unwrap().get(&req).cloned().map_or_else(
+    ///                 || Ok((req.clone(), req)),
+    ///                 |response| Err(Ok(response))),
+    ///         )
+    ///     },
+    ///     move |res: Result<Response, _>, req| async move {
+    ///         if let Ok(ref val) = res {
+    ///             cache.lock().unwrap().insert(req, val.clone());
+    ///         }
+    ///         res
+    ///     },
+    /// );
+    ///
+    /// // Call the new service
+    /// let request = Request::new("cats");
+    /// let response = new_service
+    ///     .ready()
+    ///     .await?
+    ///     .call(request)
+    ///     .await?;
+    /// #     Ok::<(), ()>(())
+    /// #     };
+    /// # }
+    /// ```
+    fn decorate<Pre, Post>(self, pre: Pre, post: Post) -> Wrap<Self, Pre, Post>
+    where
+        Self: Sized,
+    {
+        Wrap::decorate(self, pre, post)
+    }
+
+    /// Wraps this service with a synchronous pre function that returns a post function.
+    ///
+    /// The given function is any function that looks like
+    /// <code>[FnMut]\(&mut Request) -> [FnOnce]\(Response) -> T</code> where `Request` is the
+    /// request given to the `Wrap` service, `Response` is the response returned from the inner
+    /// service, and `T` is the response returned from the `Wrap` service. If the inner service
+    /// returns an error the error is output directly without being given to the post function.
+    ///
+    /// # Examples
+    ///
+    /// Attaching a unique identifier to the request and response:
+    ///
+    /// ```rust
+    /// use tower::{Service, ServiceExt};
+    /// # #[derive(Clone)] struct UUID;
+    /// # impl UUID { fn new() -> Self { UUID } }
+    /// # struct Request;
+    /// # impl Request {
+    /// #     fn new(_: &str) -> Self { Request }
+    /// #     fn set_identifier(&mut self, _: UUID) {}
+    /// # }
+    /// # struct Response;
+    /// # impl Response { fn set_identifier(&mut self, _: UUID) {} }
+    ///
+    /// # fn main() {
+    /// #     async {
+    /// // A service returning Result<Response, Error>
+    /// let service = /* ... */
+    /// # tower::service_fn(|_: Request| async { Ok::<_, ()>(Response) });
+    ///
+    /// // Attach a unique identifier to each request and response
+    /// let mut new_service = service.wrap(|req| {
+    ///     let uuid = UUID::new();
+    ///     req.set_identifier(uuid.clone());
+    ///     move |mut res| {
+    ///         res.set_identifier(uuid);
+    ///         res
+    ///     }
+    /// });
+    ///
+    /// // Call the new service
+    /// let request = Request::new("cats");
+    /// let response = new_service
+    ///     .ready()
+    ///     .await?
+    ///     .call(request)
+    ///     .await?;
+    /// #     Ok::<(), ()>(())
+    /// #     };
+    /// # }
+    /// ```
+    fn wrap<F, F2, T>(
+        self,
+        f: F,
+    ) -> Wrap<Self, wrap::Pre<Request, T, Self::Error, F>, wrap::Post<F2, Self::Error>>
+    where
+        Self: Sized,
+        F: FnMut(&mut Request) -> F2,
+        F2: FnOnce(Self::Response) -> T,
+    {
+        Wrap::new(self, f)
+    }
 }
 
 impl<T: ?Sized, Request> ServiceExt<Request> for T where T: tower_service::Service<Request> {}