use std::{
ops::DerefMut,
pin::Pin,
sync::{Arc, Mutex},
task::{Context, Poll, ready},
};
use http_body::Body;
use snafu::Snafu;
use tokio::sync::oneshot;
pub struct RemainStream<S> {
rx: Arc<Mutex<oneshot::Receiver<S>>>,
}
impl<S> Clone for RemainStream<S> {
fn clone(&self) -> Self {
Self {
rx: self.rx.clone(),
}
}
}
impl<S> RemainStream<S> {
pub fn pending() -> (oneshot::Sender<S>, Self) {
let (tx, rx) = oneshot::channel();
(
tx,
Self {
rx: Arc::new(Mutex::new(rx)),
},
)
}
pub fn immediately(stream: S) -> Self {
let (tx, this) = Self::pending();
_ = tx.send(stream);
this
}
}
impl<S> Future for RemainStream<S> {
type Output = Option<S>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut rx = self.get_mut().rx.lock().expect("poisoned mutex");
Pin::new(rx.deref_mut()).poll(cx).map(Result::ok)
}
}
pub struct TakeoverSlot<T> {
state: Arc<Mutex<TakeoverState<T>>>,
}
impl<T> Clone for TakeoverSlot<T> {
fn clone(&self) -> Self {
Self {
state: self.state.clone(),
}
}
}
enum TakeoverState<T> {
Available { stream: RemainStream<T> },
Taken,
}
impl<T> TakeoverSlot<T> {
pub fn new(stream: RemainStream<T>) -> Self {
Self {
state: Arc::new(Mutex::new(TakeoverState::Available { stream })),
}
}
pub fn poll_take(&self, cx: &mut Context<'_>) -> Poll<Result<T, TakeoverError>> {
let mut guard = self.state.lock().expect("poisoned mutex");
match &mut *guard {
TakeoverState::Available { stream } => match ready!(Pin::new(stream).poll(cx)) {
Some(value) => {
*guard = TakeoverState::Taken;
Poll::Ready(Ok(value))
}
None => {
*guard = TakeoverState::Taken;
Poll::Ready(Err(TakeoverError::Aborted))
}
},
TakeoverState::Taken => Poll::Ready(Err(TakeoverError::AlreadyTaken)),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Snafu)]
pub enum TakeoverError {
#[snafu(display("takeover stream is unsupported on this message"))]
Unsupported,
#[snafu(display("stream already taken"))]
AlreadyTaken,
#[snafu(display("stream provider was dropped before delivery"))]
Aborted,
#[snafu(display("message body could not be released"))]
BodyNotReleased,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Snafu)]
#[snafu(module, visibility(pub))]
pub enum UpgradeError {
#[snafu(display("failed to take over stream"))]
Takeover { source: TakeoverError },
#[snafu(display("upgrade incomplete, missing {missing} stream"))]
Incomplete { missing: MissingStream },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MissingStream {
Read,
Write,
Both,
}
impl std::fmt::Display for MissingStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MissingStream::Read => f.write_str("read"),
MissingStream::Write => f.write_str("write"),
MissingStream::Both => f.write_str("read and write"),
}
}
}
impl From<TakeoverError> for UpgradeError {
fn from(source: TakeoverError) -> Self {
UpgradeError::Takeover { source }
}
}
mod sealed {
use super::*;
pub trait SealedTakeover<T> {
fn poll_takeover(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TakeoverError>>;
}
}
fn poll_release_body<B: Body + Unpin>(
body: &mut B,
cx: &mut Context<'_>,
) -> Poll<Result<(), TakeoverError>> {
loop {
match ready!(Pin::new(&mut *body).poll_frame(cx)) {
Some(Ok(_)) => continue,
Some(Err(_)) => return Poll::Ready(Err(TakeoverError::BodyNotReleased)),
None => return Poll::Ready(Ok(())),
}
}
}
fn poll_take_slot<T: Send + 'static>(
extensions: &http::Extensions,
cx: &mut Context<'_>,
) -> Poll<Result<T, TakeoverError>> {
let slot = extensions
.get::<TakeoverSlot<T>>()
.ok_or(TakeoverError::Unsupported)?;
slot.poll_take(cx)
}
pub trait HasTakeover<T> {
fn poll_takeover(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TakeoverError>>;
}
impl<M, T> HasTakeover<T> for M
where
M: sealed::SealedTakeover<T>,
{
fn poll_takeover(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TakeoverError>> {
sealed::SealedTakeover::poll_takeover(self, cx)
}
}
impl<B: Body + Unpin, T: Send + 'static> sealed::SealedTakeover<T> for http::Request<B> {
fn poll_takeover(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TakeoverError>> {
ready!(poll_release_body(self.body_mut(), cx))?;
poll_take_slot(self.extensions(), cx)
}
}
impl<B: Body + Unpin, T: Send + 'static> sealed::SealedTakeover<T> for http::Response<B> {
fn poll_takeover(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TakeoverError>> {
ready!(poll_release_body(self.body_mut(), cx))?;
poll_take_slot(self.extensions(), cx)
}
}
impl<B: Body + Unpin, T: Send + 'static> sealed::SealedTakeover<T> for &mut http::Request<B> {
fn poll_takeover(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TakeoverError>> {
ready!(poll_release_body(self.body_mut(), cx))?;
poll_take_slot(self.extensions(), cx)
}
}
impl<B: Body + Unpin, T: Send + 'static> sealed::SealedTakeover<T> for &mut http::Response<B> {
fn poll_takeover(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, TakeoverError>> {
ready!(poll_release_body(self.body_mut(), cx))?;
poll_take_slot(self.extensions(), cx)
}
}
#[cfg(test)]
mod tests {
use std::{
future::poll_fn,
pin::Pin,
task::{Context, Poll},
};
use bytes::Bytes;
use http_body::{Body, Frame};
use super::*;
use crate::message::stream::ReadStream;
#[derive(Debug, Clone)]
struct ErrorBody;
impl Body for ErrorBody {
type Data = Bytes;
type Error = std::io::Error;
fn poll_frame(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
Poll::Ready(Some(Err(std::io::Error::other("body error"))))
}
}
#[tokio::test]
async fn takeover_returns_unsupported_when_slot_missing() {
let mut request = http::Request::new(http_body_util::Empty::<Bytes>::new());
let result = poll_fn(|cx| HasTakeover::<ReadStream>::poll_takeover(&mut request, cx)).await;
assert!(matches!(result, Err(TakeoverError::Unsupported)));
}
#[tokio::test]
async fn takeover_returns_ready_when_slot_available() {
let mut request = http::Request::new(http_body_util::Empty::<Bytes>::new());
let (read_tx, read) = RemainStream::<ReadStream>::pending();
request.extensions_mut().insert(TakeoverSlot::new(read));
drop(read_tx);
let result = poll_fn(|cx| HasTakeover::<ReadStream>::poll_takeover(&mut request, cx)).await;
assert!(matches!(result, Err(TakeoverError::Aborted)));
}
#[tokio::test]
async fn takeover_returns_already_taken_when_taken_twice() {
let mut request = http::Request::new(http_body_util::Empty::<Bytes>::new());
let read = RemainStream::<u32>::immediately(42);
request.extensions_mut().insert(TakeoverSlot::new(read));
let first = poll_fn(|cx| HasTakeover::<u32>::poll_takeover(&mut request, cx)).await;
assert!(first.is_ok());
let second = poll_fn(|cx| HasTakeover::<u32>::poll_takeover(&mut request, cx)).await;
assert!(matches!(second, Err(TakeoverError::AlreadyTaken)));
}
#[tokio::test]
async fn takeover_returns_aborted_when_sender_dropped() {
let mut request = http::Request::new(http_body_util::Empty::<Bytes>::new());
let (read_tx, read) = RemainStream::<ReadStream>::pending();
request.extensions_mut().insert(TakeoverSlot::new(read));
drop(read_tx);
let result = poll_fn(|cx| HasTakeover::<ReadStream>::poll_takeover(&mut request, cx)).await;
assert!(matches!(result, Err(TakeoverError::Aborted)));
}
#[tokio::test]
async fn takeover_returns_body_not_released_on_body_error() {
let mut request = http::Request::new(ErrorBody);
let result = poll_fn(|cx| HasTakeover::<ReadStream>::poll_takeover(&mut request, cx)).await;
assert!(matches!(result, Err(TakeoverError::BodyNotReleased)));
}
}