639 lines
19 KiB
Rust
639 lines
19 KiB
Rust
//! Negotiate a pool of services
|
|
//!
|
|
//! The negotiate pool allows for a service that can decide between two service
|
|
//! types based on an intermediate return value. It differs from typical
|
|
//! routing since it doesn't depend on the request, but the response.
|
|
//!
|
|
//! The original use case is support ALPN upgrades to HTTP/2, with a fallback
|
|
//! to HTTP/1.
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
|
//! # struct Conn;
|
|
//! # impl Conn { fn negotiated_protocol(&self) -> &[u8] { b"h2" } }
|
|
//! # let some_tls_connector = tower::service::service_fn(|_| async move {
|
|
//! # Ok::<_, std::convert::Infallible>(Conn)
|
|
//! # });
|
|
//! # let http1_layer = tower::layer::layer_fn(|s| s);
|
|
//! # let http2_layer = tower::layer::layer_fn(|s| s);
|
|
//! let mut pool = hyper_util::client::pool::negotiate::builder()
|
|
//! .connect(some_tls_connector)
|
|
//! .inspect(|c| c.negotiated_protocol() == b"h2")
|
|
//! .fallback(http1_layer)
|
|
//! .upgrade(http2_layer)
|
|
//! .build();
|
|
//!
|
|
//! // connect
|
|
//! let mut svc = pool.call(http::Uri::from_static("https://hyper.rs")).await?;
|
|
//! svc.ready().await;
|
|
//!
|
|
//! // http1 or http2 is now set up
|
|
//! # let some_http_req = http::Request::new(());
|
|
//! let resp = svc.call(some_http_req).await?;
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
|
|
pub use self::internal::builder;
|
|
|
|
#[cfg(docsrs)]
|
|
pub use self::internal::Builder;
|
|
#[cfg(docsrs)]
|
|
pub use self::internal::Negotiate;
|
|
#[cfg(docsrs)]
|
|
pub use self::internal::Negotiated;
|
|
|
|
mod internal {
|
|
use std::future::Future;
|
|
use std::pin::Pin;
|
|
use std::sync::{Arc, Mutex};
|
|
use std::task::{self, ready, Poll};
|
|
|
|
use pin_project_lite::pin_project;
|
|
use tower_layer::Layer;
|
|
use tower_service::Service;
|
|
|
|
type BoxError = Box<dyn std::error::Error + Send + Sync>;
|
|
|
|
/// A negotiating pool over an inner make service.
|
|
///
|
|
/// Created with [`builder()`].
|
|
///
|
|
/// # Unnameable
|
|
///
|
|
/// This type is normally unnameable, forbidding naming of the type within
|
|
/// code. The type is exposed in the documentation to show which methods
|
|
/// can be publicly called.
|
|
#[derive(Clone)]
|
|
pub struct Negotiate<L, R> {
|
|
left: L,
|
|
right: R,
|
|
}
|
|
|
|
/// A negotiated service returned by [`Negotiate`].
|
|
///
|
|
/// # Unnameable
|
|
///
|
|
/// This type is normally unnameable, forbidding naming of the type within
|
|
/// code. The type is exposed in the documentation to show which methods
|
|
/// can be publicly called.
|
|
#[derive(Clone, Debug)]
|
|
pub enum Negotiated<L, R> {
|
|
#[doc(hidden)]
|
|
Fallback(L),
|
|
#[doc(hidden)]
|
|
Upgraded(R),
|
|
}
|
|
|
|
pin_project! {
|
|
pub struct Negotiating<Dst, L, R>
|
|
where
|
|
L: Service<Dst>,
|
|
R: Service<()>,
|
|
{
|
|
#[pin]
|
|
state: State<Dst, L::Future, R::Future>,
|
|
left: L,
|
|
right: R,
|
|
}
|
|
}
|
|
|
|
pin_project! {
|
|
#[project = StateProj]
|
|
enum State<Dst, FL, FR> {
|
|
Eager {
|
|
#[pin]
|
|
future: FR,
|
|
dst: Option<Dst>,
|
|
},
|
|
Fallback {
|
|
#[pin]
|
|
future: FL,
|
|
},
|
|
Upgrade {
|
|
#[pin]
|
|
future: FR,
|
|
}
|
|
}
|
|
}
|
|
|
|
pin_project! {
|
|
#[project = NegotiatedProj]
|
|
pub enum NegotiatedFuture<L, R> {
|
|
Fallback {
|
|
#[pin]
|
|
future: L
|
|
},
|
|
Upgraded {
|
|
#[pin]
|
|
future: R
|
|
},
|
|
}
|
|
}
|
|
|
|
/// A builder to configure a `Negotiate`.
|
|
///
|
|
/// # Unnameable
|
|
///
|
|
/// This type is normally unnameable, forbidding naming of the type within
|
|
/// code. The type is exposed in the documentation to show which methods
|
|
/// can be publicly called.
|
|
#[derive(Debug)]
|
|
pub struct Builder<C, I, L, R> {
|
|
connect: C,
|
|
inspect: I,
|
|
fallback: L,
|
|
upgrade: R,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct WantsConnect;
|
|
#[derive(Debug)]
|
|
pub struct WantsInspect;
|
|
#[derive(Debug)]
|
|
pub struct WantsFallback;
|
|
#[derive(Debug)]
|
|
pub struct WantsUpgrade;
|
|
|
|
/// Start a builder to construct a `Negotiate` pool.
|
|
pub fn builder() -> Builder<WantsConnect, WantsInspect, WantsFallback, WantsUpgrade> {
|
|
Builder {
|
|
connect: WantsConnect,
|
|
inspect: WantsInspect,
|
|
fallback: WantsFallback,
|
|
upgrade: WantsUpgrade,
|
|
}
|
|
}
|
|
|
|
impl<C, I, L, R> Builder<C, I, L, R> {
|
|
/// Provide the initial connector.
|
|
pub fn connect<CC>(self, connect: CC) -> Builder<CC, I, L, R> {
|
|
Builder {
|
|
connect,
|
|
inspect: self.inspect,
|
|
fallback: self.fallback,
|
|
upgrade: self.upgrade,
|
|
}
|
|
}
|
|
|
|
/// Provide the inspector that determines the result of the negotiation.
|
|
pub fn inspect<II>(self, inspect: II) -> Builder<C, II, L, R> {
|
|
Builder {
|
|
connect: self.connect,
|
|
inspect,
|
|
fallback: self.fallback,
|
|
upgrade: self.upgrade,
|
|
}
|
|
}
|
|
|
|
/// Provide the layer to fallback to if negotiation fails.
|
|
pub fn fallback<LL>(self, fallback: LL) -> Builder<C, I, LL, R> {
|
|
Builder {
|
|
connect: self.connect,
|
|
inspect: self.inspect,
|
|
fallback,
|
|
upgrade: self.upgrade,
|
|
}
|
|
}
|
|
|
|
/// Provide the layer to upgrade to if negotiation succeeds.
|
|
pub fn upgrade<RR>(self, upgrade: RR) -> Builder<C, I, L, RR> {
|
|
Builder {
|
|
connect: self.connect,
|
|
inspect: self.inspect,
|
|
fallback: self.fallback,
|
|
upgrade,
|
|
}
|
|
}
|
|
|
|
/// Build the `Negotiate` pool.
|
|
pub fn build<Dst>(self) -> Negotiate<L::Service, R::Service>
|
|
where
|
|
C: Service<Dst>,
|
|
C::Error: Into<BoxError>,
|
|
L: Layer<Inspector<C, C::Response, I>>,
|
|
L::Service: Service<Dst> + Clone,
|
|
<L::Service as Service<Dst>>::Error: Into<BoxError>,
|
|
R: Layer<Inspected<C::Response>>,
|
|
R::Service: Service<()> + Clone,
|
|
<R::Service as Service<()>>::Error: Into<BoxError>,
|
|
I: Fn(&C::Response) -> bool + Clone,
|
|
{
|
|
let Builder {
|
|
connect,
|
|
inspect,
|
|
fallback,
|
|
upgrade,
|
|
} = self;
|
|
|
|
let slot = Arc::new(Mutex::new(None));
|
|
let wrapped = Inspector {
|
|
svc: connect,
|
|
inspect,
|
|
slot: slot.clone(),
|
|
};
|
|
let left = fallback.layer(wrapped);
|
|
|
|
let right = upgrade.layer(Inspected { slot });
|
|
|
|
Negotiate { left, right }
|
|
}
|
|
}
|
|
|
|
impl<L, R> Negotiate<L, R> {
|
|
/// Get a mutable reference to the fallback service.
|
|
pub fn fallback_mut(&mut self) -> &mut L {
|
|
&mut self.left
|
|
}
|
|
|
|
/// Get a mutable reference to the upgrade service.
|
|
pub fn upgrade_mut(&mut self) -> &mut R {
|
|
&mut self.right
|
|
}
|
|
}
|
|
|
|
impl<L, R, Target> Service<Target> for Negotiate<L, R>
|
|
where
|
|
L: Service<Target> + Clone,
|
|
L::Error: Into<BoxError>,
|
|
R: Service<()> + Clone,
|
|
R::Error: Into<BoxError>,
|
|
{
|
|
type Response = Negotiated<L::Response, R::Response>;
|
|
type Error = BoxError;
|
|
type Future = Negotiating<Target, L, R>;
|
|
|
|
fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
self.left.poll_ready(cx).map_err(Into::into)
|
|
}
|
|
|
|
fn call(&mut self, dst: Target) -> Self::Future {
|
|
let left = self.left.clone();
|
|
Negotiating {
|
|
state: State::Eager {
|
|
future: self.right.call(()),
|
|
dst: Some(dst),
|
|
},
|
|
// place clone, take original that we already polled-ready.
|
|
left: std::mem::replace(&mut self.left, left),
|
|
right: self.right.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<Dst, L, R> Future for Negotiating<Dst, L, R>
|
|
where
|
|
L: Service<Dst>,
|
|
L::Error: Into<BoxError>,
|
|
R: Service<()>,
|
|
R::Error: Into<BoxError>,
|
|
{
|
|
type Output = Result<Negotiated<L::Response, R::Response>, BoxError>;
|
|
|
|
fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
|
|
// States:
|
|
// - `Eager`: try the "right" path first; on `UseOther` sentinel, fall back to left.
|
|
// - `Fallback`: try the left path; on `UseOther` sentinel, upgrade back to right.
|
|
// - `Upgrade`: retry the right path after a fallback.
|
|
// If all fail, give up.
|
|
let mut me = self.project();
|
|
loop {
|
|
match me.state.as_mut().project() {
|
|
StateProj::Eager { future, dst } => match ready!(future.poll(cx)) {
|
|
Ok(out) => return Poll::Ready(Ok(Negotiated::Upgraded(out))),
|
|
Err(err) => {
|
|
let err = err.into();
|
|
if UseOther::is(&*err) {
|
|
let dst = dst.take().unwrap();
|
|
let f = me.left.call(dst);
|
|
me.state.set(State::Fallback { future: f });
|
|
continue;
|
|
} else {
|
|
return Poll::Ready(Err(err));
|
|
}
|
|
}
|
|
},
|
|
StateProj::Fallback { future } => match ready!(future.poll(cx)) {
|
|
Ok(out) => return Poll::Ready(Ok(Negotiated::Fallback(out))),
|
|
Err(err) => {
|
|
let err = err.into();
|
|
if UseOther::is(&*err) {
|
|
let f = me.right.call(());
|
|
me.state.set(State::Upgrade { future: f });
|
|
continue;
|
|
} else {
|
|
return Poll::Ready(Err(err));
|
|
}
|
|
}
|
|
},
|
|
StateProj::Upgrade { future } => match ready!(future.poll(cx)) {
|
|
Ok(out) => return Poll::Ready(Ok(Negotiated::Upgraded(out))),
|
|
Err(err) => return Poll::Ready(Err(err.into())),
|
|
},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<L, R> Negotiated<L, R> {
|
|
// Could be useful?
|
|
#[cfg(test)]
|
|
pub(super) fn is_fallback(&self) -> bool {
|
|
matches!(self, Negotiated::Fallback(_))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(super) fn is_upgraded(&self) -> bool {
|
|
matches!(self, Negotiated::Upgraded(_))
|
|
}
|
|
|
|
// TODO: are these the correct methods? Or .as_ref().fallback(), etc?
|
|
|
|
/// Get a reference to the fallback service if this is it.
|
|
pub fn fallback_ref(&self) -> Option<&L> {
|
|
if let Negotiated::Fallback(ref left) = self {
|
|
Some(left)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Get a mutable reference to the fallback service if this is it.
|
|
pub fn fallback_mut(&mut self) -> Option<&mut L> {
|
|
if let Negotiated::Fallback(ref mut left) = self {
|
|
Some(left)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Get a reference to the upgraded service if this is it.
|
|
pub fn upgraded_ref(&self) -> Option<&R> {
|
|
if let Negotiated::Upgraded(ref right) = self {
|
|
Some(right)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Get a mutable reference to the upgraded service if this is it.
|
|
pub fn upgraded_mut(&mut self) -> Option<&mut R> {
|
|
if let Negotiated::Upgraded(ref mut right) = self {
|
|
Some(right)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<L, R, Req, Res, E> Service<Req> for Negotiated<L, R>
|
|
where
|
|
L: Service<Req, Response = Res, Error = E>,
|
|
R: Service<Req, Response = Res, Error = E>,
|
|
{
|
|
type Response = Res;
|
|
type Error = E;
|
|
type Future = NegotiatedFuture<L::Future, R::Future>;
|
|
|
|
fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
match self {
|
|
Negotiated::Fallback(ref mut s) => s.poll_ready(cx),
|
|
Negotiated::Upgraded(ref mut s) => s.poll_ready(cx),
|
|
}
|
|
}
|
|
|
|
fn call(&mut self, req: Req) -> Self::Future {
|
|
match self {
|
|
Negotiated::Fallback(ref mut s) => NegotiatedFuture::Fallback {
|
|
future: s.call(req),
|
|
},
|
|
Negotiated::Upgraded(ref mut s) => NegotiatedFuture::Upgraded {
|
|
future: s.call(req),
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<L, R, Out> Future for NegotiatedFuture<L, R>
|
|
where
|
|
L: Future<Output = Out>,
|
|
R: Future<Output = Out>,
|
|
{
|
|
type Output = Out;
|
|
|
|
fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
|
|
match self.project() {
|
|
NegotiatedProj::Fallback { future } => future.poll(cx),
|
|
NegotiatedProj::Upgraded { future } => future.poll(cx),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ===== internal =====
|
|
|
|
pub struct Inspector<M, S, I> {
|
|
svc: M,
|
|
inspect: I,
|
|
slot: Arc<Mutex<Option<S>>>,
|
|
}
|
|
|
|
pin_project! {
|
|
pub struct InspectFuture<F, S, I> {
|
|
#[pin]
|
|
future: F,
|
|
inspect: I,
|
|
slot: Arc<Mutex<Option<S>>>,
|
|
}
|
|
}
|
|
|
|
impl<M: Clone, S, I: Clone> Clone for Inspector<M, S, I> {
|
|
fn clone(&self) -> Self {
|
|
Self {
|
|
svc: self.svc.clone(),
|
|
inspect: self.inspect.clone(),
|
|
slot: self.slot.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<M, S, I, Target> Service<Target> for Inspector<M, S, I>
|
|
where
|
|
M: Service<Target, Response = S>,
|
|
M::Error: Into<BoxError>,
|
|
I: Clone + Fn(&S) -> bool,
|
|
{
|
|
type Response = M::Response;
|
|
type Error = BoxError;
|
|
type Future = InspectFuture<M::Future, S, I>;
|
|
|
|
fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
self.svc.poll_ready(cx).map_err(Into::into)
|
|
}
|
|
|
|
fn call(&mut self, dst: Target) -> Self::Future {
|
|
InspectFuture {
|
|
future: self.svc.call(dst),
|
|
inspect: self.inspect.clone(),
|
|
slot: self.slot.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<F, I, S, E> Future for InspectFuture<F, S, I>
|
|
where
|
|
F: Future<Output = Result<S, E>>,
|
|
E: Into<BoxError>,
|
|
I: Fn(&S) -> bool,
|
|
{
|
|
type Output = Result<S, BoxError>;
|
|
|
|
fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
|
|
let me = self.project();
|
|
let s = ready!(me.future.poll(cx)).map_err(Into::into)?;
|
|
Poll::Ready(if (me.inspect)(&s) {
|
|
*me.slot.lock().unwrap() = Some(s);
|
|
Err(UseOther.into())
|
|
} else {
|
|
Ok(s)
|
|
})
|
|
}
|
|
}
|
|
|
|
pub struct Inspected<S> {
|
|
slot: Arc<Mutex<Option<S>>>,
|
|
}
|
|
|
|
impl<S, Target> Service<Target> for Inspected<S> {
|
|
type Response = S;
|
|
type Error = BoxError;
|
|
type Future = std::future::Ready<Result<S, BoxError>>;
|
|
|
|
fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
if self.slot.lock().unwrap().is_some() {
|
|
Poll::Ready(Ok(()))
|
|
} else {
|
|
Poll::Ready(Err(UseOther.into()))
|
|
}
|
|
}
|
|
|
|
fn call(&mut self, _dst: Target) -> Self::Future {
|
|
let s = self
|
|
.slot
|
|
.lock()
|
|
.unwrap()
|
|
.take()
|
|
.ok_or_else(|| UseOther.into());
|
|
std::future::ready(s)
|
|
}
|
|
}
|
|
|
|
impl<S> Clone for Inspected<S> {
|
|
fn clone(&self) -> Inspected<S> {
|
|
Inspected {
|
|
slot: self.slot.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct UseOther;
|
|
|
|
impl std::fmt::Display for UseOther {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.write_str("sentinel error; using other")
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for UseOther {}
|
|
|
|
impl UseOther {
|
|
fn is(err: &(dyn std::error::Error + 'static)) -> bool {
|
|
let mut source = Some(err);
|
|
while let Some(err) = source {
|
|
if err.is::<UseOther>() {
|
|
return true;
|
|
}
|
|
source = err.source();
|
|
}
|
|
false
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use futures_util::future;
|
|
use tower_service::Service;
|
|
use tower_test::assert_request_eq;
|
|
|
|
#[tokio::test]
|
|
async fn not_negotiated_falls_back_to_left() {
|
|
let (mock_svc, mut handle) = tower_test::mock::pair::<(), &'static str>();
|
|
|
|
let mut negotiate = super::builder()
|
|
.connect(mock_svc)
|
|
.inspect(|_: &&str| false)
|
|
.fallback(layer_fn(|s| s))
|
|
.upgrade(layer_fn(|s| s))
|
|
.build();
|
|
|
|
std::future::poll_fn(|cx| negotiate.poll_ready(cx))
|
|
.await
|
|
.unwrap();
|
|
|
|
let fut = negotiate.call(());
|
|
let nsvc = future::join(fut, async move {
|
|
assert_request_eq!(handle, ()).send_response("one");
|
|
})
|
|
.await
|
|
.0
|
|
.expect("call");
|
|
assert!(nsvc.is_fallback());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn negotiated_uses_right() {
|
|
let (mock_svc, mut handle) = tower_test::mock::pair::<(), &'static str>();
|
|
|
|
let mut negotiate = super::builder()
|
|
.connect(mock_svc)
|
|
.inspect(|_: &&str| true)
|
|
.fallback(layer_fn(|s| s))
|
|
.upgrade(layer_fn(|s| s))
|
|
.build();
|
|
|
|
std::future::poll_fn(|cx| negotiate.poll_ready(cx))
|
|
.await
|
|
.unwrap();
|
|
|
|
let fut = negotiate.call(());
|
|
let nsvc = future::join(fut, async move {
|
|
assert_request_eq!(handle, ()).send_response("one");
|
|
})
|
|
.await
|
|
.0
|
|
.expect("call");
|
|
|
|
assert!(nsvc.is_upgraded());
|
|
}
|
|
|
|
fn layer_fn<F>(f: F) -> LayerFn<F> {
|
|
LayerFn(f)
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct LayerFn<F>(F);
|
|
|
|
impl<F, S, Out> tower_layer::Layer<S> for LayerFn<F>
|
|
where
|
|
F: Fn(S) -> Out,
|
|
{
|
|
type Service = Out;
|
|
fn layer(&self, inner: S) -> Self::Service {
|
|
(self.0)(inner)
|
|
}
|
|
}
|
|
}
|