use async_trait::async_trait;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use camel_api::{CamelError, Exchange};
pub struct ExchangeEnvelope {
pub exchange: Exchange,
pub reply_tx: Option<oneshot::Sender<Result<Exchange, CamelError>>>,
}
#[derive(Clone)]
pub struct ConsumerContext {
sender: mpsc::Sender<ExchangeEnvelope>,
cancel_token: CancellationToken,
}
impl ConsumerContext {
pub fn new(sender: mpsc::Sender<ExchangeEnvelope>, cancel_token: CancellationToken) -> Self {
Self {
sender,
cancel_token,
}
}
pub async fn cancelled(&self) {
self.cancel_token.cancelled().await
}
pub fn is_cancelled(&self) -> bool {
self.cancel_token.is_cancelled()
}
pub fn cancel_token(&self) -> CancellationToken {
self.cancel_token.clone()
}
pub fn sender(&self) -> mpsc::Sender<ExchangeEnvelope> {
self.sender.clone()
}
pub async fn send(&self, exchange: Exchange) -> Result<(), CamelError> {
self.sender
.send(ExchangeEnvelope {
exchange,
reply_tx: None,
})
.await
.map_err(|_| CamelError::ChannelClosed)
}
pub async fn send_and_wait(&self, exchange: Exchange) -> Result<Exchange, CamelError> {
let (reply_tx, reply_rx) = oneshot::channel();
self.sender
.send(ExchangeEnvelope {
exchange,
reply_tx: Some(reply_tx),
})
.await
.map_err(|_| CamelError::ChannelClosed)?;
reply_rx.await.map_err(|_| CamelError::ChannelClosed)?
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConcurrencyModel {
Sequential,
Concurrent { max: Option<usize> },
}
#[async_trait]
pub trait Consumer: Send + Sync {
async fn start(&mut self, context: ConsumerContext) -> Result<(), CamelError>;
async fn stop(&mut self) -> Result<(), CamelError>;
fn concurrency_model(&self) -> ConcurrencyModel {
ConcurrencyModel::Sequential
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_consumer_context_cancelled() {
let (tx, _rx) = mpsc::channel(16);
let token = CancellationToken::new();
let ctx = ConsumerContext::new(tx, token.clone());
assert!(!ctx.is_cancelled());
token.cancel();
ctx.cancelled().await;
assert!(ctx.is_cancelled());
}
#[test]
fn test_concurrency_model_default_is_sequential() {
use super::ConcurrencyModel;
struct DummyConsumer;
#[async_trait::async_trait]
impl super::Consumer for DummyConsumer {
async fn start(&mut self, _ctx: super::ConsumerContext) -> Result<(), CamelError> {
Ok(())
}
async fn stop(&mut self) -> Result<(), CamelError> {
Ok(())
}
}
let consumer = DummyConsumer;
assert_eq!(consumer.concurrency_model(), ConcurrencyModel::Sequential);
}
#[test]
fn test_concurrency_model_concurrent_override() {
use super::ConcurrencyModel;
struct ConcurrentConsumer;
#[async_trait::async_trait]
impl super::Consumer for ConcurrentConsumer {
async fn start(&mut self, _ctx: super::ConsumerContext) -> Result<(), CamelError> {
Ok(())
}
async fn stop(&mut self) -> Result<(), CamelError> {
Ok(())
}
fn concurrency_model(&self) -> ConcurrencyModel {
ConcurrencyModel::Concurrent { max: Some(16) }
}
}
let consumer = ConcurrentConsumer;
assert_eq!(
consumer.concurrency_model(),
ConcurrencyModel::Concurrent { max: Some(16) }
);
}
}