FP Complete


This is the fourth and final post in a series on combining web and gRPC services into a single service using Tower, Hyper, Axum, and Tonic. The full four parts are:

  1. Overview of Tower
  2. Understanding Hyper, and first experiences with Axum
  3. Demonstration of Tonic for a gRPC client/server
  4. Today’s post: How to combine Axum and Tonic services into a single service

Single port, two protocols

That heading is a lie. Both an Axum web application and a gRPC server speak the same protocol: HTTP/2. It may be more fair to say they speak different dialects of it. But importantly, it’s trivially easy to look at a request and determine whether it wants to talk to the gRPC server or not. gRPC requests will all include the header Content-Type: application/grpc. So our final step today is to write something that can accept both a gRPC Service and a normal Service, and return one unified service. Let’s do it! For reference, complete code is in src/bin/server-hybrid.rs.

Let’s start off with our main function, and demonstrate what we want this thing to look like:

#[tokio::main]
async fn main() {
    let addr = SocketAddr::from(([0, 0, 0, 0], 3000));

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

    let grpc_service = tonic::transport::Server::builder()
        .add_service(EchoServer::new(MyEcho))
        .into_service();

    let hybrid_make_service = hybrid(axum_make_service, grpc_service);

    let server = hyper::Server::bind(&addr).serve(hybrid_make_service);

    if let Err(e) = server.await {
        eprintln!("server error: {}", e);
    }
}

We set up simplistic axum_make_service and grpc_service values, and then use the hybrid function to combine them into a single service. Notice the difference in those names, and the fact that we called into_make_service for the former and into_service for the latter. Believe it or not, that’s going to cause us a lot of pain very shortly.

Anyway, with that yet-to-be-explained hybrid function, spinning up a hybrid server is a piece of cake. But the devil’s in the details!

Also: there are simpler ways of going about the code below using trait objects. I avoided any type erasure techniques, since (1) I thought the code was a bit clearer this way, and (2) it turns into a nicer tutorial in my opinion. The one exception is that I am using a trait object for errors, since Hyper itself does so, and it simplifies the code significantly to use the same error representation across services.

Defining hybrid

Our hybrid function is going to return a HybridMakeService value:

fn hybrid<MakeWeb, Grpc>(make_web: MakeWeb, grpc: Grpc) -> HybridMakeService<MakeWeb, Grpc> {
    HybridMakeService { make_web, grpc }
}

struct HybridMakeService<MakeWeb, Grpc> {
    make_web: MakeWeb,
    grpc: Grpc,
}

I’m going to be consistent and verbose with the type variable names throughout. Here, we have the type variables MakeWeb and Grpc. This reflects the difference between what Axum and Tonic provide from an API perspective. We’ll need to provide Axum’s MakeWeb with connection information in order to get the request-handling Service. With Grpc, we won’t have to do that.

In any event, we’re ready to implement our Service for HybridMakeService:

impl<ConnInfo, MakeWeb, Grpc> Service<ConnInfo> for HybridMakeService<MakeWeb, Grpc>
where
    MakeWeb: Service<ConnInfo>,
    Grpc: Clone,
{
    // ...
}

We have the two expected type variables MakeWeb and Grpc, as well as ConnInfo, to represent whatever connection information we’re given. Grpc won’t care about that at all, but the ConnInfo must match up with what MakeWeb is receiving. Therefore, we have the bound MakeWeb: Service<ConnInfo>. The Grpc: Clone bound will make sense shortly.

When we receive an incoming connection, we’ll need to do two things:

Once we have the new Web Service and the cloned Grpc, we’ll wrap these up into a new struct, HybridService. We’re also going to need some help to perform the necessary async actions, so we’ll create a new helper Future type. This all looks like:

type Response = HybridService<MakeWeb::Response, Grpc>;
type Error = MakeWeb::Error;
type Future = HybridMakeServiceFuture<MakeWeb::Future, Grpc>;

fn poll_ready(
    &mut self,
    cx: &mut std::task::Context,
) -> std::task::Poll<Result<(), Self::Error>> {
    self.make_web.poll_ready(cx)
}

fn call(&mut self, conn_info: ConnInfo) -> Self::Future {
    HybridMakeServiceFuture {
        web_future: self.make_web.call(conn_info),
        grpc: Some(self.grpc.clone()),
    }
}

Note that we’re deferring to self.make_web to say it’s ready and passing along its errors. Let’s tie this piece off by looking at HybridMakeServiceFuture:

#[pin_project]
struct HybridMakeServiceFuture<WebFuture, Grpc> {
    #[pin]
    web_future: WebFuture,
    grpc: Option<Grpc>,
}

impl<WebFuture, Web, WebError, Grpc> Future for HybridMakeServiceFuture<WebFuture, Grpc>
where
    WebFuture: Future<Output = Result<Web, WebError>>,
{
    type Output = Result<HybridService<Web, Grpc>, WebError>;

    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context) -> Poll<Self::Output> {
        let this = self.project();
        match this.web_future.poll(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Ready(Ok(web)) => Poll::Ready(Ok(HybridService {
                web,
                grpc: this.grpc.take().expect("Cannot poll twice!"),
            })),
        }
    }
}

We need to pull in pin_project to allow us to project the pinned web future inside our poll implementation. (If you’re not familiar with pin_project, don’t worry, we’ll describe things later on with HybridFuture.) When we poll web_future, we could end up in one of three states:

There’s some funny business with that this.grpc.take() to get the cloned Grpc value out of the Option. Futures have an invariant that, once they return Ready, they cannot be polled again. Therefore, it’s safe to assume that take will only ever be called once. But all of this pain could be avoided if Axum exposed an into_service method instead.

HybridService

The previous types will ultimately produce a HybridService. Let’s look at what that is:

struct HybridService<Web, Grpc> {
    web: Web,
    grpc: Grpc,
}

impl<Web, Grpc, WebBody, GrpcBody> Service<Request<Body>> for HybridService<Web, Grpc>
where
    Web: Service<Request<Body>, Response = Response<WebBody>>,
    Grpc: Service<Request<Body>, Response = Response<GrpcBody>>,
    Web::Error: Into<Box<dyn std::error::Error + Send + Sync + 'static>>,
    Grpc::Error: Into<Box<dyn std::error::Error + Send + Sync + 'static>>,
{
    // ...
}

This HybridService will take Request<Body> as input. The underlying Web and Grpc will also take Request<Body> as input, but they’ll produce slightly different output: either Response<WebBody> or Response<GrpcBody>. We’re going to need to somehow unify those body representations. As mentioned above, we’re going to use trait objects for error handling, so no unification there is necessary.

type Response = Response<HybridBody<WebBody, GrpcBody>>;
type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
type Future = HybridFuture<Web::Future, Grpc::Future>;

The associated Response type is going to be a Response<...> as well, but its body is going to be the HybridBody<WebBody, GrpcBody> type. We’ll get to that later. Similarly, we have two different Futures that may get called, depending on the kind of request. We need to unify over that with a HybridFuture type.

Next, let’s look at poll_ready. We need to check for both Web and Grpc being ready for a new request. And each check can result in one of three cases: Pending, Ready(Err), or Ready(Ok). This function is all about pattern matching and unifying the error representation using .into():

fn poll_ready(
    &mut self,
    cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
    match self.web.poll_ready(cx) {
        Poll::Ready(Ok(())) => match self.grpc.poll_ready(cx) {
            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
            Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
            Poll::Pending => Poll::Pending,
        },
        Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
        Poll::Pending => Poll::Pending,
    }
}

And finally, we can see call, where the real logic we’re trying to accomplish lives. This is where we get to look at the request and determine where to route it:

fn call(&mut self, req: Request<Body>) -> Self::Future {
    if req.headers().get("content-type").map(|x| x.as_bytes()) == Some(b"application/grpc") {
        HybridFuture::Grpc(self.grpc.call(req))
    } else {
        HybridFuture::Web(self.web.call(req))
    }
}

Amazing. All of this work for essentially 5 lines of meaningful code!

HybridFuture

That’s it, we’re at the end! The final type we’re going to analyze in this series is HybridFuture. (There’s also a HybridBody type, but it’s similar enough to HybridFuture that it doesn’t warrant its own explanation.) The struct‘s definition is:

#[pin_project(project = HybridFutureProj)]
enum HybridFuture<WebFuture, GrpcFuture> {
    Web(#[pin] WebFuture),
    Grpc(#[pin] GrpcFuture),
}

Like before, we’re using pin_project. This time, let’s explore why. The interface for the Future trait requires pinned pointers in memory. Specifically, the first argument to poll is self: Pin<&mut Self>. Rust itself never gives any guarantees about object permanence, and that’s absolutely critical to writing an async runtime system.

The poll method on HybridFuture is therefore going to receive an argument of type Pin<&mut HybridFuture>. The problem is that we need to call the poll method on the underlying WebBody or GrpcBody. Assuming we have the Web variant, the problem we face is that pattern matching on HybridFuture will give us a &WebFuture or &mut WebFuture. It won’t give us a Pin<&mut WebFuture>, which is what we need!

pin_project makes a projected data type, and provides a method .project() on the original that gives us those pinned mutable references instead. This allows us to implement the Future trait for HybridFuture correctly, like so:

impl<WebFuture, GrpcFuture, WebBody, GrpcBody, WebError, GrpcError> Future
    for HybridFuture<WebFuture, GrpcFuture>
where
    WebFuture: Future<Output = Result<Response<WebBody>, WebError>>,
    GrpcFuture: Future<Output = Result<Response<GrpcBody>, GrpcError>>,
    WebError: Into<Box<dyn std::error::Error + Send + Sync + 'static>>,
    GrpcError: Into<Box<dyn std::error::Error + Send + Sync + 'static>>,
{
    type Output = Result<
        Response<HybridBody<WebBody, GrpcBody>>,
        Box<dyn std::error::Error + Send + Sync + 'static>,
    >;

    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context) -> Poll<Self::Output> {
        match self.project() {
            HybridFutureProj::Web(a) => match a.poll(cx) {
                Poll::Ready(Ok(res)) => Poll::Ready(Ok(res.map(HybridBody::Web))),
                Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
                Poll::Pending => Poll::Pending,
            },
            HybridFutureProj::Grpc(b) => match b.poll(cx) {
                Poll::Ready(Ok(res)) => Poll::Ready(Ok(res.map(HybridBody::Grpc))),
                Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
                Poll::Pending => Poll::Pending,
            },
        }
    }
}

We unify together the successful response bodies with the HybridBody enum and use a trait object for error handling. And now we’re presenting a single unified type for both types of requests. Huzzah!

Conclusions

Thank you dear reader for getting through these posts. I hope it was helpful. I definitely felt more comfortable with the Tower/Hyper ecosystem after diving into these details like this. Let’s sum up some highlights from this series:

If you want to review it, remember that a complete project is available on GitHub at https://github.com/snoyberg/tonic-example.

Finally, some more subjective takeaways from me:

If this kind of content was helpful, and you’re interested in more in the future, please consider subscribing to our blog. Let me know (on Twitter or elsewhere) if you have any requests for additional content like this.

If you’re looking for more Rust content, check out:

Subscribe to our blog via email

Email subscriptions come from our Atom feed and are handled by Blogtrottr. You will only receive notifications of blog posts, and can unsubscribe any time.

Tagged