use std::{
pin::Pin,
task::{Context, Poll},
};
use bytes::Bytes;
use connectrpc::error::ConnectError;
use http_body::{Body, Frame};
use tokio::sync::mpsc;
pub(crate) const CHANNEL_DEPTH: usize = 32;
pub(crate) struct IpcRequestBody {
rx: mpsc::Receiver<Bytes>,
}
impl IpcRequestBody {
pub(crate) fn channel() -> (mpsc::Sender<Bytes>, Self) {
let (tx, rx) = mpsc::channel(CHANNEL_DEPTH);
(tx, Self { rx })
}
pub(crate) fn complete(bytes: Bytes) -> Self {
let (tx, rx) = mpsc::channel(1);
if !bytes.is_empty() {
let _ = tx.try_send(bytes);
}
drop(tx);
Self { rx }
}
}
impl Body for IpcRequestBody {
type Data = Bytes;
type Error = ConnectError;
fn poll_frame(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Bytes>, Self::Error>>> {
self.rx
.poll_recv(cx)
.map(|opt| opt.map(|bytes| Ok(Frame::data(bytes))))
}
}