use tokio::sync::mpsc;
const DEFAULT_CHANNEL_CAPACITY: usize = 8;
pub(crate) enum StreamEvent {
Chunk(Vec<u8>),
End,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StreamClosed;
impl std::fmt::Display for StreamClosed {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("streaming response receiver has been closed")
}
}
impl std::error::Error for StreamClosed {}
pub struct StreamingResponse {
pub status: u16,
pub(crate) content_type: Option<&'static str>,
pub(crate) rx: mpsc::Receiver<StreamEvent>,
}
impl StreamingResponse {
#[must_use]
pub fn new(status: u16) -> (Self, BodyWriter) {
Self::channel(status, None, DEFAULT_CHANNEL_CAPACITY)
}
#[must_use]
pub fn channel(
status: u16,
content_type: Option<&'static str>,
capacity: usize,
) -> (Self, BodyWriter) {
let (tx, rx) = mpsc::channel(capacity.max(1));
(
Self {
status,
content_type,
rx,
},
BodyWriter { tx },
)
}
}
#[derive(Clone)]
pub struct BodyWriter {
tx: mpsc::Sender<StreamEvent>,
}
impl BodyWriter {
pub async fn send(&self, data: Vec<u8>) -> Result<(), StreamClosed> {
self.tx
.send(StreamEvent::Chunk(data))
.await
.map_err(|_| StreamClosed)
}
pub async fn finish(self) -> Result<(), StreamClosed> {
self.tx
.send(StreamEvent::End)
.await
.map_err(|_| StreamClosed)
}
}
pub(crate) enum RecvOutcome {
Chunk(Vec<u8>),
End,
Aborted,
}
impl StreamingResponse {
pub(crate) async fn recv(&mut self) -> RecvOutcome {
match self.rx.recv().await {
Some(StreamEvent::Chunk(data)) => RecvOutcome::Chunk(data),
Some(StreamEvent::End) => RecvOutcome::End,
None => RecvOutcome::Aborted,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn send_then_finish_yields_chunks_then_end() {
let (mut response, writer) = StreamingResponse::channel(200, None, 4);
writer.send(b"a".to_vec()).await.unwrap();
writer.send(b"b".to_vec()).await.unwrap();
writer.finish().await.unwrap();
match response.recv().await {
RecvOutcome::Chunk(data) => assert_eq!(data, b"a"),
_ => panic!("expected first chunk"),
}
match response.recv().await {
RecvOutcome::Chunk(data) => assert_eq!(data, b"b"),
_ => panic!("expected second chunk"),
}
match response.recv().await {
RecvOutcome::End => {}
_ => panic!("expected End"),
}
}
#[tokio::test]
async fn drop_without_finish_yields_aborted() {
let (mut response, writer) = StreamingResponse::channel(200, None, 4);
writer.send(b"only".to_vec()).await.unwrap();
drop(writer);
match response.recv().await {
RecvOutcome::Chunk(data) => assert_eq!(data, b"only"),
_ => panic!("expected chunk before abort"),
}
match response.recv().await {
RecvOutcome::Aborted => {}
_ => panic!("expected Aborted after writer drop without finish"),
}
}
#[tokio::test]
async fn send_after_receiver_dropped_returns_stream_closed() {
let (response, writer) = StreamingResponse::channel(200, None, 4);
drop(response);
let err = writer.send(b"x".to_vec()).await.unwrap_err();
assert_eq!(err, StreamClosed);
}
#[tokio::test]
async fn zero_capacity_is_rounded_up_to_one() {
let (mut response, writer) = StreamingResponse::channel(200, None, 0);
writer.send(b"ok".to_vec()).await.unwrap();
match response.recv().await {
RecvOutcome::Chunk(data) => assert_eq!(data, b"ok"),
_ => panic!("expected chunk"),
}
}
#[test]
fn stream_closed_display_and_error_trait() {
fn assert_error<E: std::error::Error>() {}
assert_error::<StreamClosed>();
assert_eq!(
StreamClosed.to_string(),
"streaming response receiver has been closed"
);
}
}