Tower

When a middleware needs to modify the behavior of a response (e.g., adding a timeout), you cannot simply return the inner service's future. Instead, you must implement a custom Future that wraps the inner future and any other auxiliary futures (like a timer).

To avoid the overhead of Box<dyn Future>, implement the Future trait manually on a custom struct. Because polling a field of a pinned struct requires 'pin projection', it is highly recommended to use the pin-project crate. This allows you to safely obtain Pin<&mut Field> from a Pin<&mut Struct>.

use pin_project::pin_project;
use std::{pin::Pin, future::Future, task::{Context, Poll}};
use tokio::time::Sleep;

#[pin_project]
pub struct ResponseFuture<F> {
    #[pin]
    response_future: F,
    #[pin]
    sleep: Sleep,
}

impl<F, Response, Error> Future for ResponseFuture<F>
where
    F: Future<Output = Result<Response, Error>>,
{
    type Output = Result<Response, Error>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();

        // 1. Poll the inner response future
        match this.response_future.poll(cx) {
            Poll::Ready(result) => return Poll::Ready(result),
            Poll::Pending => {}
        }

        // 2. Poll the auxiliary future (e.g., a timer)
        match this.sleep.poll(cx) {
            Poll::Ready(()) => {
                // Handle timeout/auxiliary completion
                todo!("Return error")
            }
            Poll::Pending => {}
        }

        // 3. Neither is ready
        Poll::Pending
    }
}

When building middleware, you often need to return an error type that can accommodate both the inner service's error and your own middleware-specific errors (e.g., a TimeoutError).

Tower uses a boxed trait object pattern to solve this. Instead of complex nested enums, define your middleware's error type to be BoxError, which is a type alias for Box<dyn std::error::Error + Send + Sync>.

Advantages:

• Stability: Changing the order of middleware doesn't change the final error type.
• Constant Size: The error type has a fixed size on the stack regardless of how many middleware layers are applied.
• Ease of use: You can extract specific errors using error.downcast_ref::<YourErrorType>() instead of massive match statements.

Requirements: To use this pattern, your middleware's Service implementation must require that the inner service's error type implements Into<BoxError>.

// Use this type alias for middleware errors
type BoxError = Box<dyn std::error::Error + Send + Sync>;

// Ensure your Service implementation has this bound
impl<S, Request> Service<Request> for MyMiddleware<S>
where
    S: Service<Request>,
    S::Error: Into<BoxError>,
{
    type Response = S::Response;
    type Error = BoxError;
    type Future = MyMiddlewareFuture<S::Future>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx).map_err(Into::into)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        // ...
    }
}

To build a middleware in Tower, you typically wrap an inner Service in a new struct. This struct should implement the Service trait, allowing it to intercept requests and responses.

Key requirements for a robust middleware struct:

• Implement Clone: Services should be clonable so that &mut self from Service::call can be converted into an owned self if needed for the response future.
• Implement Debug: For better observability.
• Constructor: Provide a new(inner: S, ...) method. It is recommended to omit trait bounds on the generic parameter S in the constructor, even if you expect it to implement Service later.
• poll_ready: Forward the readiness check to the inner service to respect backpressure.
• call: Intercept the request and return a custom Future that wraps the inner service's future.

use std::time::Duration;
use tower::Service;
use std::task::{Context, Poll};

#[derive(Debug, Clone)]
struct MyMiddleware<S> {
    inner: S,
    config: Duration,
}

impl<S> MyMiddleware<S> {
    pub fn new(inner: S, config: Duration) -> Self {
        MyMiddleware { inner, config }
    }
}

impl<S, Request> Service<Request> for MyMiddleware<S>
where
    S: Service<Request>,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future; // Or a custom Future type

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        self.inner.call(request)
    }
}

The standard pattern for implementing Tower middleware is to create a struct that wraps an inner Service and returns a custom Future that wraps the inner service's Future.

Steps to implement:

1. Define the Middleware struct: Holds the inner service and any configuration (like a Duration for timeouts).
2. Implement Service for the struct:
   • poll_ready should call the inner service's poll_ready and map the error using .map_err(Into::into) if using BoxError.
   • call should call the inner service's call and return your custom ResponseFuture.
3. Define the ResponseFuture struct: Use pin_project to manage the pinning of the inner future and any other state (like a timer).
4. Implement Future for ResponseFuture:
   • In poll, poll the inner future first. If it returns Ready(Ok(res)), return the result.
   • If the inner future returns an error, map it using .map_err(Into::into).
   • If the inner future is Pending, poll your middleware's internal state (e.g., a Sleep timer). If the timer expires, return your middleware's specific error.

use pin_project::pin_project;
use std::task::{Context, Poll};
use std::pin::Pin;
use tower::Service;

#[derive(Debug, Clone)]
struct MyMiddleware<S> {
    inner: S,
    // configuration...
}

impl<S, Request> Service<Request> for MyMiddleware<S>
where
    S: Service<Request>,
    S::Error: Into<BoxError>,
{
    type Response = S::Response;
    type Error = BoxError;
    type Future = MyMiddlewareFuture<S::Future>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx).map_err(Into::into)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        let response_future = self.inner.call(request);
        MyMiddlewareFuture { response_future /* ... */ }
    }
}

#[pin_project]
struct MyMiddlewareFuture<F> {
    #[pin]
    response_future: F,
    // other state like sleep...
}

impl<F, Response, Error> Future for MyMiddlewareFuture<F>
where
    F: Future<Output = Result<Response, Error>>,
    Error: Into<BoxError>,
{
    type Output = Result<Response, BoxError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();

        match this.response_future.poll(cx) {
            Poll::Ready(result) => {
                return Poll::Ready(result.map_err(Into::into));
            }
            Poll::Pending => {}
        }

        // Check middleware logic (e.g. timeout)
        // ...

        Poll::Pending
    }
}