tower-rs / tower

async fn(Request) -> Result<Response, Error>
https://docs.rs/tower
MIT License
3.56k stars 281 forks source link

Feature Request: Add a way to branch on the the service pipeline #797

Open Sagebati opened 1 month ago

Sagebati commented 1 month ago

So the I have a case where I want to have two behaviors if a token is set or not. I didn't want to write my service, so I tried to use the service contaminators (optional, filter, either, etc..)

But what I wanted to do I not handled (I think) example.

I have a Predicate on the request and I want to Run either A or B depending of the predicate

let predicate_a = |req| Ok(req);
let predicate_b = |req| Err(req);
let layer_b = ...;
let layer_a = ...;

let global = ServiceBuilder::new()
   .layer(x)
   .layer(bar).
   .branch(predicate_a, service_a) // service_a is always called because it's always Ok(_)
   .branch(predicate_a, service_b) // service_b is never called because it's always Err(_)
   .layer(etc);

I'm working on an implementation that could live in the mod util

Draft:

// I took took the predicate in tower and added a Error type for unification
pub trait Predicate<Request> {
    type Request;

    type Error;

    fn check(&mut self, request: Request) -> Result<Self::Request, Self::Error>;
}

pub struct BrachService<P, L, L2, R, E> {
    predicate: P,
    predicate_result: bool, // This i'm not sure yet
    layer: L,
    layer2: L2,
    _response_error: PhantomData<(R, E)>,
}

impl<T, R, E, P, L, L2> Service<T> for BrachService<P, L, L2, R, E>
where
    L: Service<T, Response = R, Error = E>,
    L2: Service<T, Response = R, Error = E>,
    L::Future: 'static,
    L2::Future: 'static,
    P: Predicate<T, Request = T, Error = T>,
{
    type Response = R;
    type Error = E;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        if self.predicate_result {
            self.layer.poll_ready(cx)
        } else {
            self.layer2.poll_ready(cx)
        }
    }

    fn call(&mut self, req: T) -> Self::Future {
        match self.predicate.check(req) {
            Ok(req) => Box::pin(self.layer.call(req)),
            Err(req) => Box::pin(self.layer2.call(req)),
        }
    }
}
seanmonstar commented 1 month ago

Would Steer work? https://docs.rs/tower/latest/tower/steer/index.html

Sagebati commented 1 month ago

Steer seems to have everything I wanted, but I struggle passing it as a layer

let mut rate_limiter = Steer::new(
        vec![BoxLayer::new(user_rate_limiter), BoxLayer::new(global_rate_limiter)],
        |req: Request, services| {
            let header = req.headers().typed_get::<Authorization<Bearer>>();

            match header
                .as_ref()
                .and_then(|header| BearerToken::from_str(header.token()).ok())
            {
                Some(_) => 0, // If some it means that we have a valid token so we authorize more request
                None => 1,    // If not we limit the number of request per second
            }
        },
    );
    base.layer(middleware!())
        .layer(rate_limiter.as_service()) // 

Got

  .layer(rate_limiter)
    |          ----- ^^^^^^^^^^^^ the trait `tower::Layer<Route>` is not implemented for `Steer<BoxLayer<_, http::Request<_>, _,
seanmonstar commented 1 month ago

Ah, because steer doesn't have a layer impl.

Sagebati commented 1 month ago

So there is no way to do it, with Steer ? I misinformed, from the beginning I was trying to combine layers and not services. I still struggle to differentiate them

I saw that there was a draft of the "same" idea https://github.com/tower-rs/tower/pull/562. Is this something that has some value for a feature request ?