Common data type for functions with different parameters (e.g. Axum route handlers)

Hi,
I'm trying to implement a router that calls different handler functions based on a context object. I took Axum's router as inspiration. I already found a very good explanation of how Axum extracts different parts of a request with so called "extractors" (GitHub - alexpusch/rust-magic-function-params: Example for Axum style magic function parameter passing).

I adapted this example a little and now I'm trying to figure out, how to store handler functions in a HashMap but struggle with finding a correct data type that allows me to store functions with a different amount of parameters.

This is my code at the moment:

Code
// Context and FromContext trait ///////////////////////////////////////////////
#[derive(Clone)]
pub struct Context {
    path: String,
    id: String,
}

impl Context {
    pub fn default() -> Self {
        Context {
            path: "/path/to/hello_world".to_string(),
            id  : "123456789".to_string()
        }
    }
}

pub trait FromContext {
    fn from_context(ctx: &Context) -> Self;
}


// Extractors //////////////////////////////////////////////////////////////////

pub struct Path(String);
impl FromContext for Path {
    fn from_context(ctx: &Context) -> Self {
        Path(ctx.path.to_string())
    }
}

pub struct Id(String);
impl FromContext for Id {
    fn from_context(ctx: &Context) -> Self {
        Id(ctx.id.to_string())
    }
}

// Handler Trait and Implementation ////////////////////////////////////////////

pub trait Handler<T> {
    fn call(self, ctx: Context);
}

impl<F, T> Handler<T> for F
where
    F: Fn(T),
    T: FromContext,
{
    fn call(self, ctx: Context) {
        (self)(T::from_context(&ctx));
    }
}

impl<F, T1, T2> Handler<(T1, T2)> for F
where
    F: Fn(T1, T2),
    T1: FromContext,
    T2: FromContext,
{
    fn call(self, ctx: Context) {
        (self)(T1::from_context(&ctx), T2::from_context(&ctx));
    }
}

// Trigger Function ////////////////////////////////////////////////////////////

pub fn trigger<T, H>(ctx: Context, handler: H)
where
    H: Handler<T>,
{
    handler.call(ctx);
}

// Handler /////////////////////////////////////////////////////////////////////

fn handler_1(Path(p): Path, Id(id): Id) {
    println!("Handler 1: path={}; id={}", p, id);
}

fn handler_2(Path(p): Path) {
    println!("Handler 2: path={}", p);
}

// Main ////////////////////////////////////////////////////////////////////////

fn main() {
    let ctx = Context::default();
    trigger(ctx.clone(), handler_1);
    trigger(ctx.clone(), handler_2);

    let h1: Box<dyn Handler<(Path, Id)>>    = Box::new(handler_1);
    let h2: Box<dyn Handler<Path>>          = Box::new(handler_2);

    // How to store h1 and h2 in the same hash map?
}

And if you want to run it here's a lilnk to Rust playground.

As you can see I tried boxing the Handler<T> but the generic type parameter T remains a problem.

I hope someone of you has an idea of how to solve this. Thanks a lot in advance :smiley:

type FnPathId = fn(Path, Id);
type FnPath = fn(Path);

#[derive(Hash, PartialEq, Eq)]
enum Handlers {
    PathId(FnPathId),
    Path(FnPath),
}

let mut map: HashMap<Handlers, ()> = HashMap::new();
map.extend([
    (Handlers::PathId(handler_1), ()),
    (Handlers::Path(handler_2), ()),
]);

playground

Thank you @vague for you quick reply.
Your solution works, but with that I'd need to define a type for every function signature. Do I understand that correctly? Is there a way of getting around that and just accept any function signature?

Because as you might imagine I want to create routes with a syntax like this:
router.add_route("/hello_world", handler_1);
And for that I need the add_route function as well as the underlying HashMap to accept handler functions of any type without defining a type corresponding to its signature first.

If you try to imitate Axum, why not look at how Axum does it?

If you read the documentation for routing::get(), for example, you'll see that it takes a Handler<T, S, B> and returns a MethodRouter<S, B>. If you follow its implementation through macros, you'll see it merely calls method_routing::on(), which in turn calls MethodRouter::on().

MethodRouter::on() then creates a BoxedIntoRoute::from_handler(handler). You can find this method here. This already hints at type erasure. And indeed:

pub(crate) struct BoxedIntoRoute<S, B, E>(Box<dyn ErasedIntoRoute<S, B, E>>);

and the concrete type underlying the boxed trait object is:

pub(crate) struct MakeErasedHandler<H, S, B> {
    pub(crate) handler: H,
    pub(crate) into_route: fn(H, S) -> Route<B>,
}

Finally, it implements ErasedIntoRoute<S, B> regardless of the concrete type parameter H:

impl<H, S, B> ErasedIntoRoute<S, B, Infallible> for MakeErasedHandler<H, S, B>
where
    H: Clone + Send + 'static,
    S: 'static,
    B: HttpBody + 'static,
{
    ...
}

And this is how the handler type is erased from the stored handler's type. Actually, I'd say it's a pretty common technique for type erasure (I did discover it myself for a personal side project, where I needed to erase a concrete private type from a public API).

4 Likes

That's not how it is actually implemented in Axum, though. It also doesn't scale to arbitrary combinations of parameters.

Thank you very much @H2CO3. Your explanation brought me on the right tracks. I had a look at Axum's code before but couldn't figure out. But even with your explanation I had to implement it myself and figure out what the code actually does an why it works.

For me the implementation held two problems:

  1. Storing the functions in a HashMap
  2. Calling the stored functions

The first problem was not too hard to solve, because of your explanation, the links to the documentation and looking at Axum's code. I'd explain it like that (please correct me if I'm wrong):
The type we use to store functions in the HashMap is BoxedIntoRoute. BoxedIntoRoute has a field of type ErasedIntoRoute which is a trait. This trait is used to "hide" the handlers type. When creating a BoxedIntoRoute object we pass it an object of MakeEraseHandler<H> which implements the handler-type-agnostic trait ErasedIntoRoute but MakeEraseHandler<H> itself depends on the handler type. So the trait ErasedIntoRoute is used to "hide" the handler type. But it still needs to expose "something", which lets us use the underlying handler (stored in MakeEraseHandler<H>).

And this is where we get to the second problem - calling a stored function:
The "something" mentioned above is the Route struct. But as the route struct cannot depend on the handler type (because then the ErasedIntoRoute trait would also depend on the handler type) we need a trait again to hide the type.
The ErasedIntoRoute trait defines a function into_route which returns an object of type Route. Route is a struct with a field of type Box<dyn Service>. Service is a trait which again hides the type of the handler behind its interface. HandlerService<H, T> is a struct which implements the Service trait. With this construct the Route struct does not depend on the handler type but it can use the handler through the interface defined by the Service trait.

As far as I understand I used the same pattern twice. First for storing the function and then again for making it callable. I created a draft of a diagram of the relation between the structs and traits.

With all of this and some little changes here and there I finally got to write the following code:

pub struct Router {
    routes: HashMap<i32, BoxedIntoRoute>
}

impl Router {
    fn new() -> Self {
        Router { routes: HashMap::new() }
    }

    fn add<HandlerParamType, HandlerType>(&mut self, i: i32, h: HandlerType)
    where
        HandlerType: Handler<HandlerParamType>,
        HandlerParamType: 'static
    {
        self.routes.insert(i, BoxedIntoRoute::from_handler(h));
    }

    fn call(&self, i: i32, ctx: Context) {
        self.routes.get(&i).unwrap().clone().call_with_ctx(ctx);
    }
}

fn main() {
    let ctx = Context::default();

    let mut router = Router::new();
    router.add(1, handler_1);
    router.add(2, handler_2);

    router.call(1, ctx.clone());
    router.call(2, ctx.clone());
}

I'm still refactoring and trying to make it a little more understandable. But when I'm done doing that I'll post the final version of it for everyone who's interested in this topic :smiley:

Hey, I finally had time to clean up the code and make some changes too. The changes I did are:

  1. Make the router accept async functions because that's what I need in the project I'm working on.
  2. Remove the BoxedIntoRoute and Route struct because all they do is putting a trait into a Box<...> and giving it a new name (see the diagram in my post above). As this made the code more difficult to read (at least that's what I think) I removed them.

On this Rust Playground you can see how I implemented it in the end. What do you think about it?

To keep it more readable I chose to not include the macro I use to generate implementations of the Handler trait for different numbers of parameters.
If someone is interested in the macro, here it is (actually it's two):

use async_trait::async_trait;
use futures::join;

macro_rules! impl_handler {
    (
        $($tn:ident),*
        ;
        $($Tn:ident),*
    ) => {
        #[async_trait]
        impl<H, $( $Tn ),* , Fut> Handler<( $( $Tn ),* ,)> for H
        where
            H: Fn( $( $Tn ), * ) -> Fut + Clone + Send + Sync + 'static,
            Fut: Future<Output = ()> + Send,
            $(
                $Tn: FromContext + Send
            ),* ,
        {
            async fn call(self, ctx: Context) {
                // Execute from_context on all parameters concurrently
                let (
                    $( $tn ),* ,
                ) = join!(
                    $(
                        $Tn::from_context(&ctx)
                    ),*
                );

                // Call the function
                (self)( $( $tn ),* ).await
            }
        }
    };
}

#[rustfmt::skip]
macro_rules! all_the_parameters {
    ($name:ident) => {
        $name!(t1; T1);
        $name!(t1, t2; T1, T2);
        $name!(t1, t2, t3; T1, T2, T3);
        $name!(t1, t2, t3, t4; T1, T2, T3, T4);
        $name!(t1, t2, t3, t4, t5; T1, T2, T3, T4, T5);
    };
}

// Implement the handler trait for functions from 1 to 5 parameters
all_the_parameters!(impl_handler);

This is quite similar to how Axum does it. The impl_handler macro generates an implementation of the Handler trait and the all_the_parameters macro provides the number of parameters I want to generate Handler implementations for.

Edit: I forgot to include that I also made the extractors async. I updated the link to the Rust Playground.

1 Like

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.