Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Choose Listener type at runtime #3172

Closed
1 task done
jasonaowen opened this issue Jan 13, 2025 · 6 comments
Closed
1 task done

Choose Listener type at runtime #3172

jasonaowen opened this issue Jan 13, 2025 · 6 comments

Comments

@jasonaowen
Copy link

  • I have looked for existing issues (including closed) about this

Feature Request

Motivation

I want to choose between a TcpListener and a UnixListener at runtime, and use the same simple serve function.

I was trying to make use of the new feature introduced in #2941 to switch between a TcpListener and a UnixListener depending on a command-line argument. It seems the genericity is for allowing Unix sockets at compile-time - which is great, thank you!

Proposal

First I tried to bind one listener or the other outside of the function, so I could then pass it as an argument:

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    assert!(args.len() == 2);

    let app = axum::Router::new().route("/", axum::routing::get(|| async { "Hello, world!" }));

    match args[1].as_str() {
        "t" => run(tokio::net::TcpListener::bind(("localhost", 3000)).await?, app).await,
        "u" => run(tokio::net::UnixListener::bind("/tmp/axum-example-socket")?, app).await,
        _ => panic!(),
    }
}

async fn run(listener: impl axum::serve::Listener, app: axum::Router) -> std::io::Result<()> {
    axum::serve::serve(listener, app).await?
}

But this fails with

  --> src/main.rs:17:39
   |
17 |     axum::serve::serve(listener, app).await?;
   |     ----------------------------------^^^^^
   |     |                                ||
   |     |                                |`Serve<impl axum::serve::Listener, Router, Router>` is not a future
   |     |                                help: remove the `.await`
   |     this call returns `Serve<impl axum::serve::Listener, Router, Router>`
   |
   = help: the trait `IntoFuture` is not implemented for `Serve<impl axum::serve::Listener, Router, Router>`
   = note: Serve<impl axum::serve::Listener, Router, Router> must be a future or must implement `IntoFuture` to be awaited

Then, while trying to simplify, I tried this:

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    assert!(args.len() == 2);

    let listener = match args[1].as_str() {
        "t" => tokio::net::TcpListener::bind(("localhost", 3000)).await?,
        "u" => tokio::net::UnixListener::bind("/tmp/axum-example-socket")?,
        _ => panic!(),
    };

    let app = axum::Router::new().route("/", axum::routing::get(|| async { "Hello, world!" }));

    axum::serve::serve(listener, app).await?
}

This fails with

error[E0308]: `match` arms have incompatible types
  --> src/main.rs:8:16
   |
6  |       let listener = match args[1].as_str() {
   |  ____________________-
7  | |         "t" => tokio::net::TcpListener::bind(("localhost", 3000)).await?,
   | |                --------------------------------------------------------- this is found to be of type `tokio::net::TcpListener`
8  | |         "u" => tokio::net::UnixListener::bind("/tmp/axum-example-socket")?,
   | |                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `TcpListener`, found `UnixListener`
9  | |         _ => panic!()
10 | |     };
   | |_____- `match` arms have incompatible types

Alternatives

I can call serve in two different branches - which, again, is much more than I could before v0.8.0!

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    assert!(args.len() == 2);

    let app = axum::Router::new().route("/", axum::routing::get(|| async { "Hello, world!" }));

    match args[1].as_str() {
        "t" => axum::serve::serve(tokio::net::TcpListener::bind(("localhost", 3000)).await?, app).await,
        "u" => axum::serve::serve(tokio::net::UnixListener::bind("/tmp/axum-example-socket")?, app).await,
        _ => panic!()
    }?;

    Ok(())
}

This looks reasonable in this small example, but with_graceful_shutdown makes it slightly more cumbersome.


Thanks again for axum and v0.8!

@mladedav
Copy link
Collaborator

You can create something like Either<T1, T2> or just this for your specific use:

enum TcpOrUnix {
   Tcp(TcpListener),
   Unix(UnixListener),
}

impl Listener for TcpOrUnix {
   // match self and then just call the inner one...
}

Is that good enough for you?

@jasonaowen
Copy link
Author

jasonaowen commented Jan 15, 2025

Ah, thank you, that approach hadn't occurred to me!

I'm trying to implement it, and having some trouble working out what the types of Listener::Io and Listener::Addr should be. I haven't found an existing enum that wraps both tokio::net::TcpStream and tokio::net::UnixStream, or one that wraps both std::net::SocketAddr and std::os::unix::net::SocketAddr, and applying the same approach got me pretty deep in the weeds trying to impl tokio::io::AsyncRead for TcpOrUnixIo.

Am I way off track? or what am I missing?

@mladedav
Copy link
Collaborator

It's a bit more boilerplate than I expected, but you can use the same idea for the Io and Addr. You can create your own type that wraps both. I think you tried that but fought with the pins in the implementation?

enum TcpOrUnix {
    Tcp(TcpListener),
    Unix(tokio::net::UnixListener),
}

enum TcpOrUnixIo {
    Tcp(TcpStream),
    Unix(tokio::net::UnixStream),
}

impl AsyncRead for TcpOrUnixIo {
    fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            TcpOrUnixIo::Tcp(tcp_stream) => std::pin::pin!(tcp_stream).poll_read(cx, buf),
            TcpOrUnixIo::Unix(unix_stream) => std::pin::pin!(unix_stream).poll_read(cx, buf),
        }
    }
}

impl AsyncWrite for TcpOrUnixIo {
    // Same idea as `AsyncRead`
}

enum TcpOrUnixAddr {
    Tcp(std::net::SocketAddr),
    Unix(tokio::net::unix::SocketAddr),
}

impl Listener for TcpOrUnix {
    type Io = TcpOrUnixIo;

    type Addr = TcpOrUnixAddr;

    fn accept(&mut self) -> impl Future<Output = (Self::Io, Self::Addr)> + Send {
        match self {
            TcpOrUnix::Tcp(tcp_listener) => futures_util::future::Either::Left(
                tcp_listener
                    .accept()
                    .map(|(io, addr)| (TcpOrUnixIo::Tcp(io), TcpOrUnixAddr::Tcp(addr))),
            ),
            TcpOrUnix::Unix(unix_listener) => futures_util::future::Either::Right(
                unix_listener
                    .accept()
                    .map(|(io, addr)| (TcpOrUnixIo::Unix(io), TcpOrUnixAddr::Unix(addr))),
            ),
        }
    }

    fn local_addr(&self) -> io::Result<Self::Addr> {
        match self {
            TcpOrUnix::Tcp(tcp_listener) => tcp_listener.local_addr().map(TcpOrUnixAddr::Tcp),
            TcpOrUnix::Unix(unix_listener) => unix_listener.local_addr().map(TcpOrUnixAddr::Unix),
        }
    }
}

@jasonaowen
Copy link
Author

It's a bit more boilerplate than I expected, but you can use the same idea for the Io and Addr. You can create your own type that wraps both. I think you tried that but fought with the pins in the implementation?

            TcpOrUnixIo::Tcp(tcp_stream) => std::pin::pin!(tcp_stream).poll_read(cx, buf),
            TcpOrUnixIo::Unix(unix_stream) => std::pin::pin!(unix_stream).poll_read(cx, buf),

Yes, that's exactly right! I did try that, and did fight with the pins. Thanks for working out what it needed to be, @mladedav!

So, do you think there is room in axum for code like this? Or is this too specific a use case?

@mladedav
Copy link
Collaborator

I think this is too specific. If there are more similar requests we might consider it (probably with generics though), but until then you can either use the code as is or publish it on crates in another crate.

@jasonaowen
Copy link
Author

Okay, thanks for taking the time to work on this with me, @mladedav!

@jasonaowen jasonaowen closed this as not planned Won't fix, can't repro, duplicate, stale Jan 17, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants