use bytes::Bytes;
use camel_api::error::CamelError;
use futures::stream;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::sync::Notify;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::{CancellationToken, PollSender};
use tracing::Instrument;
use wasmtime::StoreContextMut;
use wasmtime::component::{
Accessor, FutureConsumer, FutureReader, Lift, Source, StreamConsumer, StreamReader,
StreamResult,
};
use crate::error::WasmError;
pub(crate) const DEFAULT_DRAIN_CHANNEL_BOUND: usize = 8;
pub(crate) type TerminalSlot = Arc<std::sync::Mutex<Option<Result<(), CamelError>>>>;
pub(crate) fn write_terminal(slot: &TerminalSlot, result: Result<(), CamelError>) {
match slot.lock() {
Ok(mut guard) => *guard = Some(result),
Err(poisoned) => *poisoned.into_inner() = Some(result),
}
}
#[derive(Clone)]
pub(crate) struct DrainCoord {
pub(crate) cancel: CancellationToken,
pub(crate) progress: Arc<Notify>,
pub(crate) receiver_gone: Arc<Notify>,
pub(crate) terminal_slot: TerminalSlot,
}
#[derive(Debug)]
pub(crate) enum DrainEvent {
Chunk(Bytes),
Error(CamelError),
}
pub(crate) trait StreamReturnable<E: 'static> {
fn take_stream(&mut self) -> Option<GuestStreamParts<E>>;
}
pub(crate) type GuestStreamParts<E> = (
StreamReader<u8>,
FutureReader<Result<(), E>>,
camel_api::StreamMetadata,
);
impl StreamReturnable<crate::bindings::camel::plugin::types::WasmError>
for crate::bindings::camel::plugin::types::WasmExchange
{
fn take_stream(
&mut self,
) -> Option<GuestStreamParts<crate::bindings::camel::plugin::types::WasmError>> {
use crate::bindings::camel::plugin::types::WasmBody;
let msg = self.output.as_mut()?;
if !matches!(msg.body, WasmBody::Stream(_)) {
return None;
}
if let WasmBody::Stream(handle) = std::mem::replace(&mut msg.body, WasmBody::Empty) {
let metadata = camel_api::StreamMetadata {
size_hint: handle.size_hint,
content_type: handle.content_type,
origin: handle.origin,
};
Some((handle.r#stream, handle.terminal, metadata))
} else {
unreachable!()
}
}
}
impl StreamReturnable<crate::bean_bindings::camel::plugin::types::WasmError>
for crate::bean_bindings::camel::plugin::types::WasmExchange
{
fn take_stream(
&mut self,
) -> Option<GuestStreamParts<crate::bean_bindings::camel::plugin::types::WasmError>> {
use crate::bean_bindings::camel::plugin::types::WasmBody;
let msg = &mut self.input;
if !matches!(msg.body, WasmBody::Stream(_)) {
return None;
}
if let WasmBody::Stream(handle) = std::mem::replace(&mut msg.body, WasmBody::Empty) {
let metadata = camel_api::StreamMetadata {
size_hint: handle.size_hint,
content_type: handle.content_type,
origin: handle.origin,
};
Some((handle.r#stream, handle.terminal, metadata))
} else {
unreachable!()
}
}
}
impl StreamReturnable<crate::source_bindings::camel::plugin::types::WasmError>
for crate::source_bindings::camel::plugin::types::WasmExchange
{
fn take_stream(
&mut self,
) -> Option<GuestStreamParts<crate::source_bindings::camel::plugin::types::WasmError>> {
use crate::source_bindings::camel::plugin::types::WasmBody;
let msg = &mut self.input;
if !matches!(msg.body, WasmBody::Stream(_)) {
return None;
}
if let WasmBody::Stream(handle) = std::mem::replace(&mut msg.body, WasmBody::Empty) {
let metadata = camel_api::StreamMetadata {
size_hint: handle.size_hint,
content_type: handle.content_type,
origin: handle.origin,
};
Some((handle.r#stream, handle.terminal, metadata))
} else {
unreachable!()
}
}
}
pub(crate) fn receiver_to_body_stream(
mut drain_rx: DrainReceiver,
) -> futures::stream::BoxStream<'static, Result<Bytes, CamelError>> {
Box::pin(stream::poll_fn(move |cx| match drain_rx.rx.poll_recv(cx) {
Poll::Ready(Some(DrainEvent::Chunk(b))) => Poll::Ready(Some(Ok(b))),
Poll::Ready(Some(DrainEvent::Error(e))) => Poll::Ready(Some(Err(e))),
Poll::Ready(None) => {
let mut guard = drain_rx.terminal.lock().unwrap_or_else(|p| p.into_inner());
match guard.take() {
Some(Err(e)) => Poll::Ready(Some(Err(e))),
_ => Poll::Ready(None),
}
}
Poll::Pending => Poll::Pending,
}))
}
pub(crate) struct ChannelConsumer<S> {
sender: PollSender<DrainEvent>,
cancel: CancellationToken,
progress: Arc<Notify>,
receiver_gone: Arc<Notify>, _marker: std::marker::PhantomData<S>,
}
impl<S> ChannelConsumer<S> {
pub(crate) fn new(
tx: mpsc::Sender<DrainEvent>,
cancel: CancellationToken,
progress: Arc<Notify>,
receiver_gone: Arc<Notify>,
) -> Self {
Self {
sender: PollSender::new(tx),
cancel,
progress,
receiver_gone,
_marker: std::marker::PhantomData,
}
}
}
impl<S> Unpin for ChannelConsumer<S> {}
impl<S: Send + 'static> StreamConsumer<S> for ChannelConsumer<S> {
type Item = u8;
fn poll_consume(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
store: StoreContextMut<'_, S>,
mut source: Source<'_, Self::Item>,
finish: bool,
) -> Poll<wasmtime::Result<StreamResult>> {
if self.cancel.is_cancelled() {
return Poll::Ready(Ok(StreamResult::Dropped));
}
match self.sender.poll_reserve(cx) {
Poll::Ready(Ok(())) => {
let mut buf = Vec::with_capacity(65536);
match source.read(store, &mut buf) {
Ok(()) => {}
Err(e) => {
let _ =
self.sender
.send_item(DrainEvent::Error(CamelError::ProcessorError(
e.to_string(),
)));
self.receiver_gone.notify_one();
return Poll::Ready(Ok(StreamResult::Dropped));
}
}
if buf.is_empty() {
if finish {
return Poll::Ready(Ok(StreamResult::Completed));
}
return Poll::Pending;
}
self.sender
.send_item(DrainEvent::Chunk(Bytes::from(buf)))
.expect("reserved slot must accept send"); self.progress.notify_one();
Poll::Ready(Ok(StreamResult::Completed))
}
Poll::Ready(Err(_)) => {
self.receiver_gone.notify_one();
Poll::Ready(Ok(StreamResult::Dropped))
}
Poll::Pending => Poll::Pending, }
}
}
pub(crate) async fn drain_guest_stream<E, S>(
accessor: &Accessor<S>,
stream_reader: StreamReader<u8>,
terminal: FutureReader<Result<(), E>>,
tx: mpsc::Sender<DrainEvent>,
coord: DrainCoord,
) where
E: Into<CamelError> + Send + Sync + 'static,
S: Send + 'static,
Result<(), E>: Lift,
{
let (terminal_tx, terminal_rx) = tokio::sync::oneshot::channel::<Result<(), E>>();
accessor.with(|mut access| {
stream_reader
.pipe(
&mut access,
ChannelConsumer::new(
tx.clone(),
coord.cancel.clone(),
coord.progress.clone(),
coord.receiver_gone.clone(),
),
)
.expect("pipe registration");
terminal
.pipe(
&mut access,
TerminalFutureConsumer {
tx: Some(terminal_tx),
_marker: std::marker::PhantomData,
},
)
.expect("terminal pipe registration"); });
let outcome = match terminal_rx.await {
Ok(v) => v,
Err(_) => {
write_terminal(
&coord.terminal_slot,
Err(CamelError::ProcessorError(
"wasm guest terminal future dropped without resolving".into(),
)),
);
drop(tx);
return;
}
};
if let Err(e) = outcome {
write_terminal(&coord.terminal_slot, Err(e.into()));
}
drop(tx);
}
struct TerminalFutureConsumer<E> {
tx: Option<tokio::sync::oneshot::Sender<Result<(), E>>>,
_marker: std::marker::PhantomData<E>,
}
impl<E> Unpin for TerminalFutureConsumer<E> {}
impl<E> std::fmt::Debug for TerminalFutureConsumer<E> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TerminalFutureConsumer")
.field("tx", &self.tx.as_ref().map(|_| "Some(Sender)"))
.finish()
}
}
impl<E: Send + Sync + 'static, S> FutureConsumer<S> for TerminalFutureConsumer<E>
where
Result<(), E>: Lift,
{
type Item = Result<(), E>;
fn poll_consume(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
store: StoreContextMut<'_, S>,
mut source: Source<'_, Self::Item>,
_finish: bool,
) -> Poll<wasmtime::Result<()>> {
let this = self.get_mut();
let mut buf: Option<Result<(), E>> = None;
source.read(store, &mut buf)?;
if let Some(value) = buf
&& let Some(tx) = this.tx.take()
{
let _ = tx.send(value);
}
Poll::Ready(Ok(()))
}
}
pub(crate) type StreamHandoff<W> =
Result<(W, Option<DrainReceiver>, camel_api::StreamMetadata), WasmError>;
pub(crate) type StreamHandoffSender<W> = oneshot::Sender<StreamHandoff<W>>;
pub(crate) struct DrainReceiver {
pub(crate) rx: mpsc::Receiver<DrainEvent>,
pub(crate) terminal: TerminalSlot,
}
pub(crate) fn take_stream_handoff_sender<W>(
handoff: &Arc<std::sync::Mutex<Option<StreamHandoffSender<W>>>>,
) -> Option<StreamHandoffSender<W>> {
handoff
.lock()
.expect("stream handoff lock not poisoned") .take()
}
pub(crate) async fn spawn_return_drain<W, F, Fut>(
pending_permit: Option<tokio::sync::OwnedSemaphorePermit>,
cancel: CancellationToken,
no_progress_timeout: Duration,
invoke_stall_timeout: Option<Duration>,
completion_notify: Option<Arc<Notify>>,
make_drive: F,
) -> Result<(W, Option<DrainReceiver>, camel_api::StreamMetadata), WasmError>
where
W: Send + 'static,
F: FnOnce(
Arc<std::sync::Mutex<Option<StreamHandoffSender<W>>>>,
mpsc::Sender<DrainEvent>,
mpsc::Receiver<DrainEvent>,
Arc<Notify>,
DrainCoord,
) -> Fut
+ Send
+ 'static,
Fut: Future<Output = Result<(), WasmError>> + Send + 'static,
{
let (handoff_tx, handoff_rx) = oneshot::channel::<StreamHandoff<W>>();
let handoff_shared = Arc::new(std::sync::Mutex::new(Some(handoff_tx)));
let cancel_child = cancel.child_token();
let progress_notify = Arc::new(Notify::new());
let receiver_gone = Arc::new(Notify::new());
let drain_started = Arc::new(Notify::new());
let terminal: TerminalSlot = Arc::new(std::sync::Mutex::new(None));
tokio::spawn(async move {
let _permit = pending_permit; let (dtx, drx) = mpsc::channel::<DrainEvent>(DEFAULT_DRAIN_CHANNEL_BOUND);
let handoff_drive = handoff_shared.clone();
let coord = DrainCoord {
cancel: cancel_child.clone(),
progress: progress_notify.clone(),
receiver_gone: receiver_gone.clone(),
terminal_slot: terminal.clone(),
};
let drive = make_drive(handoff_drive, dtx, drx, drain_started.clone(), coord);
let span = tracing::info_span!("wasm return-stream drain");
let drain_result = async {
tokio::select! {
r = crate::runtime::WasmRuntime::drive_with_drain_watchdog(
drive, &progress_notify, &drain_started, no_progress_timeout,
invoke_stall_timeout,
) => r,
_ = cancel_child.cancelled() => Err(WasmError::Cancelled(
"invoke cancelled by caller drop".into(),
)),
}
}
.instrument(span.clone())
.await;
let outcome = match &drain_result {
Ok(()) => {
if cancel_child.is_cancelled() {
"cancelled"
} else {
"eof"
}
}
Err(_) => "error",
};
tracing::debug!(outcome, "wasm return-stream drain ended");
match drain_result {
Ok(_) => {}
Err(e) => {
tracing::warn!(error = %e, "wasm return-stream drain failed");
if let Some(tx) = take_stream_handoff_sender(&handoff_shared) {
let _ = tx.send(Err(e.clone()));
}
write_terminal(&terminal, Err(e.into()));
}
}
if let Some(notify) = completion_notify {
notify.notify_one();
}
});
let (exchange_out, drain_rx, metadata) = match handoff_rx.await {
Ok(Ok((exchange_out, drain_rx, metadata))) => (exchange_out, drain_rx, metadata),
Ok(Err(e)) => return Err(e),
Err(_) => {
return Err(WasmError::InstantiationFailed("drain task panicked".into()));
}
};
Ok((exchange_out, drain_rx, metadata))
}
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
use camel_api::error::CamelError;
use futures::StreamExt;
use tokio::sync::mpsc;
#[tokio::test]
async fn receiver_to_body_stream_yields_chunks_then_eof() {
let (tx, rx) = mpsc::channel::<DrainEvent>(8);
tx.send(DrainEvent::Chunk(Bytes::from_static(b"hello")))
.await
.unwrap();
tx.send(DrainEvent::Chunk(Bytes::from_static(b"world")))
.await
.unwrap();
drop(tx);
let mut s = receiver_to_body_stream(crate::return_stream::DrainReceiver {
rx,
terminal: Arc::new(std::sync::Mutex::new(None)),
});
assert_eq!(
s.next().await.unwrap().unwrap(),
Bytes::from_static(b"hello")
);
assert_eq!(
s.next().await.unwrap().unwrap(),
Bytes::from_static(b"world")
);
assert!(s.next().await.is_none()); }
#[tokio::test]
async fn receiver_to_body_stream_surfaces_error_as_final_item() {
let (tx, rx) = mpsc::channel::<DrainEvent>(8);
tx.send(DrainEvent::Chunk(Bytes::from_static(b"partial")))
.await
.unwrap();
tx.send(DrainEvent::Error(CamelError::Io("guest boom".into())))
.await
.unwrap();
drop(tx);
let mut s = receiver_to_body_stream(crate::return_stream::DrainReceiver {
rx,
terminal: Arc::new(std::sync::Mutex::new(None)),
});
assert_eq!(
s.next().await.unwrap().unwrap(),
Bytes::from_static(b"partial")
);
assert!(s.next().await.unwrap().is_err()); assert!(s.next().await.is_none());
}
#[test]
fn write_terminal_writes_result() {
let slot: TerminalSlot = Arc::new(std::sync::Mutex::new(None));
write_terminal(&slot, Err(CamelError::ProcessorError("boom".into())));
let guard = slot.lock().unwrap();
let stored = guard.as_ref().expect("terminal slot must be Some");
assert!(
matches!(stored, Err(CamelError::ProcessorError(msg)) if msg == "boom"),
"expected ProcessorError('boom'), got {stored:?}"
);
}
#[test]
fn write_terminal_recovers_from_poison() {
let slot: TerminalSlot = Arc::new(std::sync::Mutex::new(None));
let slot2 = slot.clone();
let handle = std::thread::spawn(move || {
let _guard = slot2.lock().unwrap();
panic!("intentional poison");
});
let _ = handle.join();
write_terminal(&slot, Err(CamelError::ProcessorError("recovered".into())));
let guard = slot.lock().unwrap_or_else(|p| p.into_inner());
let stored = guard.as_ref().expect("terminal slot must be Some");
assert!(
matches!(stored, Err(CamelError::ProcessorError(msg)) if msg == "recovered"),
"expected ProcessorError('recovered'), got {stored:?}"
);
}
}