Add rest and grpc example (#967)

* Add rest and grpc example * remove needless dependencies * Apply suggestions from code review Co-authored-by: David Pedersen <[email protected]> * Update examples/rest-grpc-multiplex/src/main.rs Co-authored-by: David Pedersen <[email protected]> * Update examples/rest-grpc-multiplex/src/main.rs Co-authored-by: David Pedersen <[email protected]> * Update examples/rest-grpc-multiplex/src/multiplex_service.rs Co-authored-by: David Pedersen <[email protected]> * Update examples/rest-grpc-multiplex/Cargo.toml Co-authored-by: David Pedersen <[email protected]> * Update examples/rest-grpc-multiplex/src/multiplex_service.rs Co-authored-by: David Pedersen <[email protected]> * clean noisy code * fix nitpicks * missing newline Co-authored-by: David Pedersen <[email protected]>
2026-08-27 00:00:24 +02:00 · 2022-05-02 15:52:55 +00:00
parent a223fd8b85
commit 4ff5e409e3
5 changed files with 216 additions and 0 deletions
@@ -0,0 +1,85 @@
+use axum::{body::BoxBody, response::IntoResponse};
+use futures::{future::BoxFuture, ready};
+use hyper::{Body, Request, Response};
+use std::{
+    convert::Infallible,
+    task::{Context, Poll},
+};
+use tower::Service;
+
+#[derive(Clone)]
+pub struct MultiplexService<A, B> {
+    rest: A,
+    rest_ready: bool,
+    grpc: B,
+    grpc_ready: bool,
+}
+
+impl<A, B> MultiplexService<A, B> {
+    pub fn new(rest: A, grpc: B) -> Self {
+        Self {
+            rest,
+            rest_ready: false,
+            grpc,
+            grpc_ready: false,
+        }
+    }
+}
+
+impl<A, B> Service<Request<Body>> for MultiplexService<A, B>
+where
+    A: Service<Request<Body>, Error = Infallible>,
+    A::Response: IntoResponse,
+    A::Future: Send + 'static,
+    B: Service<Request<Body>, Error = Infallible>,
+    B::Response: IntoResponse,
+    B::Future: Send + 'static,
+{
+    type Response = Response<BoxBody>;
+    type Error = Infallible;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        loop {
+            match (self.rest_ready, self.grpc_ready) {
+                (true, true) => {
+                    return Ok(()).into();
+                }
+                (false, _) => {
+                    ready!(self.rest.poll_ready(cx))?;
+                    self.rest_ready = false;
+                }
+                (_, false) => {
+                    ready!(self.grpc.poll_ready(cx))?;
+                    self.grpc_ready = true;
+                }
+            }
+        }
+    }
+
+    fn call(&mut self, req: Request<Body>) -> Self::Future {
+        if is_grpc_request(&req) {
+            self.grpc_ready = false;
+            let future = self.grpc.call(req);
+            Box::pin(async move {
+                let res = future.await?;
+                Ok(res.into_response())
+            })
+        } else {
+            self.rest_ready = false;
+            let future = self.rest.call(req);
+            Box::pin(async move {
+                let res = future.await?;
+                Ok(res.into_response())
+            })
+        }
+    }
+}
+
+fn is_grpc_request<B>(req: &Request<B>) -> bool {
+    req.headers()
+        .get("content-type")
+        .map(|content_type| content_type.as_bytes())
+        .filter(|content_type| content_type.starts_with(b"application/grpc"))
+        .is_some()
+}