Require less/no boiler plate to allow custom middleware to extract `State`

[jsoneaday]

Feature Request

It would be nice if it required much less boiler plate to get AppState passed as a parameter to a custom middleware function. The best would be if no additional code was needed at all and the AppState was simply included into the function's parameter signature, like how handlers are allowed to do.

Motivation

In larger projects there could be multiple layers of middleware and it can be tedious to impl all the Traits just to get access to AppState

Proposal

Ideally, simply allow all custom middleware to simply add AppState as one of their parameters during definition of the function

Alternatives

Not as far as I know.

Additional Comment

I am new to Axum but really like it. If someone more experienced could comment on feasibility of this feature, I'd be happy to take a stab at implementing it myself because I really want it. Note I'm referring to this https://docs.rs/axum/latest/axum/middleware/index.html#accessing-state-in-middleware

[davidpdrsn]

Yes I agree that would be nice but we don't currently know how to do it.

The issue is that Router::layer ideally should accept any L: tower::Layer. That trait has no concept of state (that's an axum thing and not a tower thing) so there isn't really a way to check which state the layer requires and pass it in.

If we changed the signature of Router::layer then it wouldn't be consistent with ServiceBuilder::layer. So the axum version would support state but ServiceBuilder wouldn't. I don't think that's good since ServiceBuilder is the recommended way to add multiple middleware at once.

[jplatte]

I could see tower adopting state as a concept, but I wonder whether passing it by shared reference and cloning where needed really makes sense for tower as a whole, and what the consequences would be if tower adopted a slightly different state concept.

[jplatte]

What we could do I guess is have our own layer trait that is automatically implemented for every tower layer. Right?

[jsoneaday]

What we could do I guess is have our own layer trait that is automatically implemented for every tower layer. Right?

Yeah that makes sense. So potentially, just thinking out loud, we could have that layer update the Request object when it is passed to the call next.run(request).await and inject the AppState.

But before I say anything further I need to go through the code. Right now I'm just learning the Axum api.

But if this is possible without disturbing the existing Router::layer and ServiceBuilder::layer api, perhaps we could even extend them with wrappers?, it would be a fun challenge to try and implement it.

[davidpdrsn]

What we could do I guess is have our own layer trait that is automatically implemented for every tower layer. Right?

Yeah we could try that and see how it works but it would also require our own ServiceBuilder I suppose, or some other way to compose multiple layers.

But if this is possible without disturbing the existing Router::layer and ServiceBuilder::layer api, perhaps we could even extend them with wrappers?

Its not possible without at least changing Router::layer. ServiceBuilder is defined in tower so we can't easily change that. I'm also not sure that state as concept makes sense in tower but haven't thought much about it.

Not sure what you mean by "wrappers".

it would be a fun challenge to try and implement it.

You're very welcome to give it a shot! However state propagation is probably the most complex part of axum so might be a little challenge depending on your level of experience. Digging into Router::with_state and boxed.rs is probably a good place to start.

[commonsensesoftware]

Fan and enthusiast here. I too am new axum and tower, but having worked on many web frameworks, I've got the gist. (sorry for the self-promotion) I've just put out more-di-axum, which brings full DI support to axum. This is not just simple AppState, but full support for singleton, transient, and scoped (e.g. per-request) lifetimes to services. In the spirit of axum ergonomics, I was able to integrate more-di (disclosure, I am the owner) in just a couple of days. Kudos on a well-thought out design.

For the sake of discussion, an abridged integration looks like:

use di::*;
use di_axum::*;

#[injectable]
struct Person;

impl Person {
    fn greet(&self) -> &str {
        "Hello world!"
    }        
}

async fn say_hello(Inject(person): Inject<Person>) -> String {
    person.greet().into()
}

#[tokio::main]
async fn main() {
    let provider = ServiceCollection::new()
        .add(Person::scoped())
        .build_provider()
        .unwrap();

    let app = Router::new()
        .route("/", get(say_hello))
        .with_provider(provider);

    let listener = TcpListener::bind("127.0.0.1:5000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

This does not use with_state or State in the handler at all. Under the hood a single middleware attaches the ServiceProvider as (middleware) state. When the middleware executes, it creates a new scope for the incoming request and adds it as an extension. Singletons are instantiated just once for the lifetime of the application, whereas scoped will be created just once for the lifetime of the request. The Clone requirement will simply clone Arc (which was already by design for DI 😉). Each of the various ways a service (struct or trait) can be resolved has a complementary extractor. If a service must be injected, versus say an optional service with TryInject, then HTTP status code 500 (Internal Server Error) is returned with a message that states the service is not registered (e.g. it's a bug).

As I'm still new to the ecosystem, I'm wondering if this concept can or should be pushed down further in to tower. I also can't help, but wonder if the same concept can't also be used for composing middleware with state. What's possible with State in axum today is pretty clunky. It's a concept needed - for sure, but it's probably not one that axum should solve beyond the barebones support it already has. I see a lot of synergy in marrying DI + tower + axum together . I'm interested in doing my part to fill that gap.

Should you be curious as to what a runnable example looks like, you can find one here. It is a reimagining of this axum example.