mirror of
https://github.com/tokio-rs/axum.git
synced 2026-08-25 00:00:23 +02:00
86 lines
2.3 KiB
Rust
86 lines
2.3 KiB
Rust
|
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()
|
||
|
|
}
|