use std::collections::VecDeque;
use std::sync::Arc;
use std::time::Duration;
use super::{AsyncOutboxStore, ClaimOutboxMessages, OutboxClaimRef};
use crate::bus::{AsyncMessageSource, Message, ReceivedMessage, TransportError};
use crate::outbox::OutboxMessage;
pub const DEFAULT_OUTBOX_SOURCE_LEASE: Duration = Duration::from_secs(30);
pub const DEFAULT_OUTBOX_SOURCE_BATCH: usize = 16;
pub struct OutboxSource<S> {
store: Arc<S>,
worker_id: String,
lease: Duration,
max_attempts: u32,
batch_size: usize,
destination: Option<String>,
buffer: VecDeque<OutboxMessage>,
}
impl<S> OutboxSource<S>
where
S: AsyncOutboxStore,
{
pub fn new(store: Arc<S>, worker_id: impl Into<String>, max_attempts: u32) -> Self {
Self {
store,
worker_id: worker_id.into(),
lease: DEFAULT_OUTBOX_SOURCE_LEASE,
max_attempts,
batch_size: DEFAULT_OUTBOX_SOURCE_BATCH,
destination: None,
buffer: VecDeque::new(),
}
}
pub fn with_lease(mut self, lease: Duration) -> Self {
assert!(
!lease.is_zero(),
"OutboxSource lease must be greater than zero"
);
self.lease = lease;
self
}
pub fn with_batch_size(mut self, batch_size: usize) -> Self {
assert!(
batch_size > 0,
"OutboxSource batch_size must be greater than zero"
);
self.batch_size = batch_size;
self
}
pub fn with_destination(mut self, destination: impl Into<String>) -> Self {
self.destination = Some(destination.into());
self
}
fn claim_request(&self) -> ClaimOutboxMessages {
let mut request =
ClaimOutboxMessages::new(self.worker_id.clone(), self.batch_size, self.lease);
if let Some(destination) = &self.destination {
request = request.to_destination(destination.clone());
}
request
}
}
impl<S> AsyncMessageSource for OutboxSource<S>
where
S: AsyncOutboxStore,
{
type Received = ReceivedOutboxMessage<S>;
async fn recv(&mut self) -> Result<Option<Self::Received>, TransportError> {
if self.buffer.is_empty() {
let claimed = self.store.claim_async(self.claim_request()).await?;
self.buffer.extend(claimed);
}
match self.buffer.pop_front() {
Some(row) => {
let claim = OutboxClaimRef::from_message(&row)?;
Ok(Some(ReceivedOutboxMessage {
store: self.store.clone(),
message: Message::from(&row),
claim,
max_attempts: self.max_attempts,
}))
}
None => Ok(None),
}
}
}
pub struct ReceivedOutboxMessage<S> {
store: Arc<S>,
message: Message,
claim: OutboxClaimRef,
max_attempts: u32,
}
impl<S> ReceivedMessage for ReceivedOutboxMessage<S>
where
S: AsyncOutboxStore,
{
fn message(&self) -> &Message {
&self.message
}
async fn ack(self) -> Result<(), TransportError> {
self.store.complete_async(&self.claim).await?;
Ok(())
}
async fn nack(self, reason: &str) -> Result<(), TransportError> {
self.store
.record_failure_async(&self.claim, reason, self.max_attempts)
.await?;
Ok(())
}
async fn dead_letter(self, reason: &str) -> Result<(), TransportError> {
self.store.fail_async(&self.claim, reason).await?;
Ok(())
}
async fn park(self, reason: &str) -> Result<(), TransportError> {
self.store.fail_async(&self.claim, reason).await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::bus::{run_source, RunOptions};
use crate::microsvc::Service;
use crate::{
CommitBatch, HashMapRepository, OutboxMessage, OutboxMessageStatus, OutboxStore,
TransactionalCommit,
};
use serde_json::json;
use std::future::Future;
fn block_on<F: Future>(future: F) -> F::Output {
use std::ptr;
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
const VTABLE: RawWakerVTable = RawWakerVTable::new(
|_| RawWaker::new(ptr::null(), &VTABLE),
|_| {},
|_| {},
|_| {},
);
let waker = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &VTABLE)) };
let mut cx = Context::from_waker(&waker);
let mut future = std::pin::pin!(future);
loop {
if let Poll::Ready(output) = future.as_mut().poll(&mut cx) {
return output;
}
}
}
fn store_row(repo: &HashMapRepository, id: &str, name: &str) {
let message = OutboxMessage::create(id, name, b"{}".to_vec()).unwrap();
let mut batch = CommitBatch::empty();
batch.outbox_messages.push(message);
block_on(repo.commit_batch(batch)).unwrap();
}
fn status(repo: &HashMapRepository, id: &str) -> Option<OutboxMessageStatus> {
let store = repo.outbox_store();
[
OutboxMessageStatus::Pending,
OutboxMessageStatus::InFlight,
OutboxMessageStatus::Published,
OutboxMessageStatus::Failed,
]
.into_iter()
.find(|status| {
store
.messages_by_status(status.clone())
.unwrap()
.iter()
.any(|m| m.id() == id)
})
}
fn source(repo: &HashMapRepository) -> OutboxSource<crate::HashMapOutboxStore> {
OutboxSource::new(Arc::new(repo.outbox_store()), "pg-transport", 3)
}
#[test]
#[should_panic(expected = "lease must be greater than zero")]
fn with_lease_zero_panics() {
let repo = HashMapRepository::new();
let _ = source(&repo).with_lease(Duration::ZERO);
}
#[test]
#[should_panic(expected = "batch_size must be greater than zero")]
fn with_batch_size_zero_panics() {
let repo = HashMapRepository::new();
let _ = source(&repo).with_batch_size(0);
}
#[test]
fn recv_yields_claimed_rows_then_drains_to_none() {
let repo = HashMapRepository::new();
store_row(&repo, "m1", "evt");
store_row(&repo, "m2", "evt");
let mut src = source(&repo);
let first = block_on(src.recv()).unwrap().expect("first row");
let second = block_on(src.recv()).unwrap().expect("second row");
let third = block_on(src.recv()).unwrap();
assert!(third.is_none(), "drains to None once nothing is claimable");
let mut ids = vec![
first.message().id().unwrap().to_string(),
second.message().id().unwrap().to_string(),
];
ids.sort();
assert_eq!(ids, vec!["m1".to_string(), "m2".to_string()]);
}
#[test]
fn ack_completes_the_row() {
let repo = HashMapRepository::new();
store_row(&repo, "m1", "evt");
let mut src = source(&repo);
let received = block_on(src.recv()).unwrap().unwrap();
block_on(received.ack()).unwrap();
assert_eq!(status(&repo, "m1"), Some(OutboxMessageStatus::Published));
}
#[test]
fn nack_releases_for_retry() {
let repo = HashMapRepository::new();
store_row(&repo, "m1", "evt");
let mut src = source(&repo);
let received = block_on(src.recv()).unwrap().unwrap();
block_on(received.nack("transient")).unwrap();
assert_eq!(status(&repo, "m1"), Some(OutboxMessageStatus::Pending));
}
#[test]
fn dead_letter_fails_the_row() {
let repo = HashMapRepository::new();
store_row(&repo, "m1", "evt");
let mut src = source(&repo);
let received = block_on(src.recv()).unwrap().unwrap();
block_on(received.dead_letter("poison")).unwrap();
assert_eq!(status(&repo, "m1"), Some(OutboxMessageStatus::Failed));
}
#[test]
fn run_source_drains_outbox_and_completes() {
let repo = HashMapRepository::new();
store_row(&repo, "m1", "evt");
store_row(&repo, "m2", "evt");
let handled = Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
let h = handled.clone();
let service = Arc::new(Service::new().event("evt").handle(
move |ctx: &crate::microsvc::Context<()>| {
let h = h.clone();
let id = ctx.message().id().unwrap_or_default().to_string();
async move {
h.lock().unwrap().push(id);
Ok(json!({}))
}
},
));
block_on(run_source(service, source(&repo), RunOptions::idempotent())).unwrap();
let mut ids = handled.lock().unwrap().clone();
ids.sort();
assert_eq!(ids, vec!["m1".to_string(), "m2".to_string()]);
assert_eq!(status(&repo, "m1"), Some(OutboxMessageStatus::Published));
assert_eq!(status(&repo, "m2"), Some(OutboxMessageStatus::Published));
}
#[test]
fn unhandled_outbox_message_is_acked_and_completed() {
let repo = HashMapRepository::new();
store_row(&repo, "m1", "unrelated");
let service: Arc<Service<()>> = Arc::new(
Service::new()
.event("evt")
.handle(|_: &crate::microsvc::Context<()>| async move { Ok(json!({})) }),
);
block_on(run_source(service, source(&repo), RunOptions::idempotent())).unwrap();
assert_eq!(status(&repo, "m1"), Some(OutboxMessageStatus::Published));
}
}