use futures::Stream;
use tokio::sync::mpsc;
use super::types::RscStreamChunk;
type CleanupCallback = Box<dyn FnOnce() + Send + 'static>;
pub struct RscStream {
receiver: mpsc::Receiver<RscStreamChunk>,
_request_context_guard:
Option<std::sync::Arc<crate::server::middleware::request_context::RequestContext>>,
cleanup: Option<CleanupCallback>,
}
impl RscStream {
pub fn new(receiver: mpsc::Receiver<RscStreamChunk>) -> Self {
Self { receiver, _request_context_guard: None, cleanup: None }
}
pub fn with_request_context(
mut self,
request_context: std::sync::Arc<crate::server::middleware::request_context::RequestContext>,
) -> Self {
self._request_context_guard = Some(request_context);
self
}
pub fn with_cleanup<F>(mut self, cleanup: F) -> Self
where
F: FnOnce() + Send + 'static,
{
if let Some(existing_cleanup) = self.cleanup.take() {
self.cleanup = Some(Box::new(move || {
if let Err(e) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
existing_cleanup();
})) {
tracing::error!(
"Panic in existing cleanup handler: {:?}",
e.downcast_ref::<&str>()
.copied()
.or_else(|| e.downcast_ref::<String>().map(|s| s.as_str()))
.unwrap_or("unknown panic")
);
}
if let Err(e) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
cleanup();
})) {
tracing::error!(
"Panic in cleanup handler: {:?}",
e.downcast_ref::<&str>()
.copied()
.or_else(|| e.downcast_ref::<String>().map(|s| s.as_str()))
.unwrap_or("unknown panic")
);
}
}));
} else {
self.cleanup = Some(Box::new(cleanup));
}
self
}
pub async fn next_chunk(&mut self) -> Option<RscStreamChunk> {
self.receiver.recv().await
}
pub fn is_complete(&self) -> bool {
self.receiver.is_closed()
}
}
impl Stream for RscStream {
type Item = Result<Vec<u8>, String>;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
use std::task::Poll;
match self.receiver.poll_recv(cx) {
Poll::Ready(Some(chunk)) => Poll::Ready(Some(Ok(chunk.data))),
Poll::Ready(None) => Poll::Ready(None),
Poll::Pending => Poll::Pending,
}
}
}
impl Drop for RscStream {
fn drop(&mut self) {
if let Some(cleanup) = self.cleanup.take()
&& let Err(e) = std::panic::catch_unwind(std::panic::AssertUnwindSafe(cleanup))
{
tracing::error!("RscStream cleanup callback panicked: {:?}", e);
}
}
}