use std::sync::{Arc, Mutex};
use bytes::Bytes;
use foundation_core::valtron::{PipeReceiver, PipeSender};
use foundation_errstacks::ErrorTrace;
use foundation_netio::shared::http::SimpleHeaders;
use crate::shared::context::{CancelSignal, Peer, Spec};
use crate::shared::error::{ConnectError, ConnectResult};
use super::frame::{race, BoxFuture, Either, Frame, FramePipe, DEFAULT_PIPE_DEPTH};
fn canceled() -> ErrorTrace<ConnectError> {
ConnectError::canceled("call canceled").into()
}
fn pipe_closed() -> ErrorTrace<ConnectError> {
ConnectError::unavailable("peer closed the stream").into()
}
async fn next_frame(
rx: &PipeReceiver<Frame>,
cancel: &CancelSignal,
) -> ConnectResult<Option<Frame>> {
if cancel.is_canceled() {
return Err(canceled());
}
match race(rx.receive(), cancel.cancelled()).await {
Either::Left(frame) => Ok(frame),
Either::Right(()) => Err(canceled()),
}
}
async fn send_frame(
tx: &PipeSender<Frame>,
frame: Frame,
cancel: &CancelSignal,
) -> ConnectResult<()> {
if cancel.is_canceled() {
return Err(canceled());
}
match race(tx.send(frame), cancel.cancelled()).await {
Either::Left(Ok(())) => Ok(()),
Either::Left(Err(_)) => Err(pipe_closed()),
Either::Right(()) => Err(canceled()),
}
}
pub trait HandlerConn: Send {
fn spec(&self) -> &Spec;
fn peer(&self) -> &Peer;
fn request_headers(&self) -> &SimpleHeaders;
fn split(self: Box<Self>) -> (Box<dyn ConnReceiver>, Box<dyn ConnSender>);
}
pub trait ConnReceiver: Send {
fn receive(&mut self) -> BoxFuture<'_, ConnectResult<Option<Bytes>>>;
fn trailers(&self) -> &SimpleHeaders;
}
pub trait ConnSender: Send {
fn send_headers(&mut self, headers: SimpleHeaders) -> BoxFuture<'_, ConnectResult<()>>;
fn send(&mut self, frame: Bytes) -> BoxFuture<'_, ConnectResult<()>>;
fn close(
self: Box<Self>,
error: Option<ErrorTrace<ConnectError>>,
trailers: SimpleHeaders,
) -> BoxFuture<'static, ConnectResult<()>>;
}
pub struct HandlerConnEnds {
pub request_tx: PipeSender<Frame>,
pub response_rx: PipeReceiver<Frame>,
pub response_headers: Arc<Mutex<Option<SimpleHeaders>>>,
}
pub struct PipeHandlerConn {
spec: Spec,
peer: Peer,
request_headers: SimpleHeaders,
req_rx: PipeReceiver<Frame>,
res_tx: PipeSender<Frame>,
response_headers: Arc<Mutex<Option<SimpleHeaders>>>,
cancel: CancelSignal,
}
impl PipeHandlerConn {
#[must_use]
pub fn new(
spec: Spec,
peer: Peer,
request_headers: SimpleHeaders,
cancel: CancelSignal,
depth: usize,
) -> (Self, HandlerConnEnds) {
let (req_tx, req_rx) = FramePipe::with_depth(depth);
let (res_tx, res_rx) = FramePipe::with_depth(depth);
let response_headers = Arc::new(Mutex::new(None));
let conn = Self {
spec,
peer,
request_headers,
req_rx,
res_tx,
response_headers: Arc::clone(&response_headers),
cancel,
};
let ends = HandlerConnEnds {
request_tx: req_tx,
response_rx: res_rx,
response_headers,
};
(conn, ends)
}
#[must_use]
pub fn with_defaults(
spec: Spec,
peer: Peer,
request_headers: SimpleHeaders,
cancel: CancelSignal,
) -> (Self, HandlerConnEnds) {
Self::new(spec, peer, request_headers, cancel, DEFAULT_PIPE_DEPTH)
}
}
impl HandlerConn for PipeHandlerConn {
fn spec(&self) -> &Spec {
&self.spec
}
fn peer(&self) -> &Peer {
&self.peer
}
fn request_headers(&self) -> &SimpleHeaders {
&self.request_headers
}
fn split(self: Box<Self>) -> (Box<dyn ConnReceiver>, Box<dyn ConnSender>) {
let receiver = PipeConnReceiver {
rx: self.req_rx,
cancel: self.cancel.clone(),
trailers: SimpleHeaders::new(),
};
let sender = PipeConnSender {
tx: self.res_tx,
response_headers: self.response_headers,
cancel: self.cancel,
headers_sent: false,
};
(Box::new(receiver), Box::new(sender))
}
}
struct PipeConnReceiver {
rx: PipeReceiver<Frame>,
cancel: CancelSignal,
trailers: SimpleHeaders,
}
impl ConnReceiver for PipeConnReceiver {
fn receive(&mut self) -> BoxFuture<'_, ConnectResult<Option<Bytes>>> {
Box::pin(async move {
match next_frame(&self.rx, &self.cancel).await? {
Some(Frame::Message(bytes)) => Ok(Some(bytes)),
Some(Frame::EndStream { trailers, .. }) => {
self.trailers = trailers;
Ok(None)
}
None => Ok(None),
}
})
}
fn trailers(&self) -> &SimpleHeaders {
&self.trailers
}
}
struct PipeConnSender {
tx: PipeSender<Frame>,
response_headers: Arc<Mutex<Option<SimpleHeaders>>>,
cancel: CancelSignal,
headers_sent: bool,
}
impl ConnSender for PipeConnSender {
fn send_headers(&mut self, headers: SimpleHeaders) -> BoxFuture<'_, ConnectResult<()>> {
Box::pin(async move {
if !self.headers_sent {
*self.response_headers.lock().unwrap_or_else(|e| e.into_inner()) = Some(headers);
self.headers_sent = true;
}
Ok(())
})
}
fn send(&mut self, frame: Bytes) -> BoxFuture<'_, ConnectResult<()>> {
self.headers_sent = true;
Box::pin(async move { send_frame(&self.tx, Frame::Message(frame), &self.cancel).await })
}
fn close(
self: Box<Self>,
error: Option<ErrorTrace<ConnectError>>,
trailers: SimpleHeaders,
) -> BoxFuture<'static, ConnectResult<()>> {
Box::pin(async move {
send_frame(&self.tx, Frame::EndStream { error, trailers }, &self.cancel).await?;
self.tx.close();
Ok(())
})
}
}
pub trait ClientConn: Send {
fn spec(&self) -> &Spec;
fn request_headers_mut(&mut self) -> &mut SimpleHeaders;
fn split(self: Box<Self>) -> (Box<dyn ClientSender>, Box<dyn ClientReceiver>);
}
pub trait ClientSender: Send {
fn request_headers_mut(&mut self) -> &mut SimpleHeaders;
fn flush_headers(&mut self) -> BoxFuture<'_, ConnectResult<()>>;
fn send(&mut self, frame: Bytes) -> BoxFuture<'_, ConnectResult<()>>;
fn close_send(self: Box<Self>) -> BoxFuture<'static, ConnectResult<()>>;
}
pub trait ClientReceiver: Send {
fn receive(&mut self) -> BoxFuture<'_, ConnectResult<Option<Bytes>>>;
fn response_headers(&mut self) -> BoxFuture<'_, ConnectResult<&SimpleHeaders>>;
fn response_trailers(&mut self) -> BoxFuture<'_, ConnectResult<&SimpleHeaders>>;
}
pub struct ClientConnEnds {
pub request_rx: PipeReceiver<Frame>,
pub response_tx: PipeSender<Frame>,
pub response_headers: Arc<Mutex<Option<SimpleHeaders>>>,
}
pub struct PipeClientConn {
spec: Spec,
request_headers: SimpleHeaders,
req_tx: PipeSender<Frame>,
res_rx: PipeReceiver<Frame>,
response_headers: Arc<Mutex<Option<SimpleHeaders>>>,
cancel: CancelSignal,
}
impl PipeClientConn {
#[must_use]
pub fn new(
spec: Spec,
request_headers: SimpleHeaders,
cancel: CancelSignal,
depth: usize,
) -> (Self, ClientConnEnds) {
let (req_tx, req_rx) = FramePipe::with_depth(depth);
let (res_tx, res_rx) = FramePipe::with_depth(depth);
let response_headers = Arc::new(Mutex::new(None));
let conn = Self {
spec,
request_headers,
req_tx,
res_rx,
response_headers: Arc::clone(&response_headers),
cancel,
};
let ends = ClientConnEnds {
request_rx: req_rx,
response_tx: res_tx,
response_headers,
};
(conn, ends)
}
#[must_use]
pub fn with_defaults(
spec: Spec,
request_headers: SimpleHeaders,
cancel: CancelSignal,
) -> (Self, ClientConnEnds) {
Self::new(spec, request_headers, cancel, DEFAULT_PIPE_DEPTH)
}
}
impl ClientConn for PipeClientConn {
fn spec(&self) -> &Spec {
&self.spec
}
fn request_headers_mut(&mut self) -> &mut SimpleHeaders {
&mut self.request_headers
}
fn split(self: Box<Self>) -> (Box<dyn ClientSender>, Box<dyn ClientReceiver>) {
let sender = PipeClientSender {
tx: self.req_tx,
request_headers: self.request_headers,
cancel: self.cancel.clone(),
flushed: false,
};
let receiver = PipeClientReceiver {
rx: self.res_rx,
response_headers: self.response_headers,
headers_cache: None,
trailers: SimpleHeaders::new(),
cancel: self.cancel,
};
(Box::new(sender), Box::new(receiver))
}
}
struct PipeClientSender {
tx: PipeSender<Frame>,
request_headers: SimpleHeaders,
cancel: CancelSignal,
flushed: bool,
}
impl ClientSender for PipeClientSender {
fn request_headers_mut(&mut self) -> &mut SimpleHeaders {
&mut self.request_headers
}
fn flush_headers(&mut self) -> BoxFuture<'_, ConnectResult<()>> {
Box::pin(async move {
self.flushed = true;
Ok(())
})
}
fn send(&mut self, frame: Bytes) -> BoxFuture<'_, ConnectResult<()>> {
self.flushed = true;
Box::pin(async move { send_frame(&self.tx, Frame::Message(frame), &self.cancel).await })
}
fn close_send(self: Box<Self>) -> BoxFuture<'static, ConnectResult<()>> {
Box::pin(async move {
self.tx.close();
Ok(())
})
}
}
struct PipeClientReceiver {
rx: PipeReceiver<Frame>,
response_headers: Arc<Mutex<Option<SimpleHeaders>>>,
headers_cache: Option<SimpleHeaders>,
trailers: SimpleHeaders,
cancel: CancelSignal,
}
impl ClientReceiver for PipeClientReceiver {
fn receive(&mut self) -> BoxFuture<'_, ConnectResult<Option<Bytes>>> {
Box::pin(async move {
match next_frame(&self.rx, &self.cancel).await? {
Some(Frame::Message(bytes)) => Ok(Some(bytes)),
Some(Frame::EndStream { error, trailers }) => {
self.trailers = trailers;
match error {
Some(e) => Err(e),
None => Ok(None),
}
}
None => Ok(None),
}
})
}
fn response_headers(&mut self) -> BoxFuture<'_, ConnectResult<&SimpleHeaders>> {
Box::pin(async move {
if self.headers_cache.is_none() {
let guard = self.response_headers.lock().unwrap_or_else(|e| e.into_inner());
self.headers_cache = Some(guard.clone().unwrap_or_default());
}
Ok(self.headers_cache.as_ref().expect("cached above"))
})
}
fn response_trailers(&mut self) -> BoxFuture<'_, ConnectResult<&SimpleHeaders>> {
Box::pin(async move { Ok(&self.trailers) })
}
}