knafeh 1.0.0

use tokio::sync::mpsc;

use crate::error::KnafehError;

/// A stream of incoming RPC message chunks (client→server or bidi).
pub struct RpcStreamReceiver {
    rx: mpsc::Receiver<Result<Vec<u8>, KnafehError>>,
}

/// A stream sender for outgoing RPC message chunks (server→client or bidi).
pub struct RpcStreamSender {
    tx: mpsc::Sender<Result<Vec<u8>, KnafehError>>,
}

/// A complete server-streaming response: metadata + a stream of body chunks.
pub struct RpcStreamResponse {
    pub metadata: super::message::Metadata,
    pub receiver: RpcStreamReceiver,
}

/// A complete client-streaming request: the incoming stream of body chunks.
pub struct RpcStreamRequest {
    pub metadata: super::message::Metadata,
    pub receiver: RpcStreamReceiver,
}

/// Create a paired (sender, receiver) for streaming RPC data.
pub fn rpc_stream_channel(buffer: usize) -> (RpcStreamSender, RpcStreamReceiver) {
    let (tx, rx) = mpsc::channel(buffer);
    (RpcStreamSender { tx }, RpcStreamReceiver { rx })
}

impl RpcStreamReceiver {
    /// Receive the next chunk, or `None` if the stream is finished.
    pub async fn next(&mut self) -> Option<Result<Vec<u8>, KnafehError>> {
        self.rx.recv().await
    }
}

impl RpcStreamSender {
    /// Send a chunk of data on the stream.
    pub async fn send(&self, data: Vec<u8>) -> Result<(), KnafehError> {
        self.tx
            .send(Ok(data))
            .await
            .map_err(|_| KnafehError::ConnectionClosed)
    }

    /// Signal an error on the stream.
    pub async fn send_error(&self, err: KnafehError) -> Result<(), KnafehError> {
        self.tx
            .send(Err(err))
            .await
            .map_err(|_| KnafehError::ConnectionClosed)
    }

    /// Close the stream (drops the sender).
    pub fn close(self) {
        drop(self);
    }
}