use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use crate::time::now_unix_ms;
use async_trait::async_trait;
use futures::StreamExt;
use serde::{Deserialize, Serialize};
use crate::kv::KvStore;
use crate::{PutMeta, Storage};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClaimedMessage {
pub id: String,
pub topic: String,
pub payload: Vec<u8>,
pub attempts: u32,
pub group: String,
}
pub use boatramp_types::config::StartPosition;
#[derive(Debug, Clone, thiserror::Error)]
pub enum MessagingError {
#[error("messaging backend error: {0}")]
Backend(String),
#[error("messaging decode error: {0}")]
Decode(String),
}
impl MessagingError {
fn backend<E: std::fmt::Display>(err: E) -> Self {
Self::Backend(err.to_string())
}
}
#[async_trait]
pub trait Messaging: Send + Sync {
async fn publish(&self, topic: &str, payload: &[u8]) -> Result<(), MessagingError>;
async fn claim(
&self,
topic: &str,
lease: Duration,
max_batch: usize,
max_attempts: u32,
) -> Result<Vec<ClaimedMessage>, MessagingError>;
async fn claim_grouped(
&self,
topic: &str,
group: &str,
_start: StartPosition,
lease: Duration,
max_batch: usize,
max_attempts: u32,
) -> Result<Vec<ClaimedMessage>, MessagingError> {
if group.is_empty() {
return self.claim(topic, lease, max_batch, max_attempts).await;
}
Err(MessagingError::Backend(
"this messaging backend does not support consumer groups".into(),
))
}
async fn ack(&self, msg: &ClaimedMessage) -> Result<(), MessagingError>;
async fn nack(&self, msg: &ClaimedMessage) -> Result<(), MessagingError>;
async fn backlog(&self, _topic: &str) -> Result<usize, MessagingError> {
Ok(0)
}
async fn dead_letter_count(&self, _topic: &str) -> Result<usize, MessagingError> {
Ok(0)
}
async fn purge_dead_letters(&self, _topic: &str) -> Result<usize, MessagingError> {
Ok(0)
}
async fn redrive_dead_letters(&self, _topic: &str) -> Result<usize, MessagingError> {
Ok(0)
}
async fn retention_sweep(&self, _topic: &str) -> Result<usize, MessagingError> {
Ok(0)
}
fn subscribe(
&self,
_topic: &str,
_after: Option<&str>,
) -> futures::stream::BoxStream<'static, StreamEvent> {
futures::stream::empty().boxed()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamEvent {
pub id: String,
pub payload: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct Record {
#[serde(default = "crate::schema_version")]
pub version: u32,
pub attempts: u32,
pub lease_until_ms: u64,
}
impl Record {
pub fn fresh() -> Self {
Self {
version: crate::SCHEMA_VERSION,
attempts: 0,
lease_until_ms: 0,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClaimAction {
Lease {
id: String,
record: Record,
},
DeadLetter {
id: String,
record: Record,
},
}
pub fn plan_claim(
mut records: Vec<(String, Record)>,
now_ms: u64,
lease_ms: u64,
max_batch: usize,
max_attempts: u32,
) -> Vec<ClaimAction> {
records.sort_by(|a, b| a.0.cmp(&b.0));
let mut actions = Vec::new();
let mut leased = 0;
for (id, mut record) in records {
if leased >= max_batch {
break;
}
if record.lease_until_ms > now_ms {
continue; }
if record.attempts >= max_attempts {
actions.push(ClaimAction::DeadLetter { id, record });
continue;
}
record.attempts += 1;
record.lease_until_ms = now_ms + lease_ms;
actions.push(ClaimAction::Lease { id, record });
leased += 1;
}
actions
}
pub fn meta_key(topic: &str, id: &str) -> String {
format!("mq/{topic}/{id}")
}
pub fn meta_prefix(topic: &str) -> String {
format!("mq/{topic}/")
}
pub fn payload_key(topic: &str, id: &str) -> String {
format!("mqp/{topic}/{id}")
}
pub fn dead_key(topic: &str, id: &str) -> String {
format!("mqdead/{topic}/{id}")
}
pub fn dead_prefix(topic: &str) -> String {
format!("mqdead/{topic}/")
}
pub fn glog_key(topic: &str, id: &str) -> String {
format!("mqglog/{topic}/{id}")
}
pub fn glog_prefix(topic: &str) -> String {
format!("mqglog/{topic}/")
}
pub fn gpayload_key(topic: &str, id: &str) -> String {
format!("mqgp/{topic}/{id}")
}
pub fn gstate_key(topic: &str, group: &str) -> String {
format!("mqgstate/{topic}/{group}")
}
pub fn gstate_prefix(topic: &str) -> String {
format!("mqgstate/{topic}/")
}
pub fn logmax_key(topic: &str) -> String {
format!("mqlogmax/{topic}")
}
pub fn gdead_key(topic: &str, group: &str, id: &str) -> String {
format!("mqgd/{topic}/{group}/{id}")
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct InFlight {
pub id: String,
pub attempts: u32,
pub lease_until_ms: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct GroupState {
#[serde(default = "crate::schema_version")]
pub version: u32,
pub hwm: String,
pub in_flight: Vec<InFlight>,
}
impl GroupState {
pub fn new(hwm: String) -> Self {
Self {
version: crate::SCHEMA_VERSION,
hwm,
in_flight: Vec::new(),
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct GroupedClaim {
pub leased: Vec<(String, u32)>,
pub dead: Vec<(String, u32)>,
}
pub fn plan_claim_grouped(
state: &mut GroupState,
now_ms: u64,
lease_ms: u64,
max_batch: usize,
max_attempts: u32,
new_ids: &[String],
) -> GroupedClaim {
let mut out = GroupedClaim::default();
let mut budget = max_batch;
let mut in_flight = std::mem::take(&mut state.in_flight);
in_flight.sort_by(|a, b| a.id.cmp(&b.id));
let mut kept = Vec::with_capacity(in_flight.len());
for mut entry in in_flight {
if entry.lease_until_ms > now_ms {
kept.push(entry); continue;
}
if entry.attempts >= max_attempts {
out.dead.push((entry.id.clone(), entry.attempts)); continue;
}
if budget == 0 {
kept.push(entry); continue;
}
entry.attempts += 1;
entry.lease_until_ms = now_ms + lease_ms;
budget -= 1;
out.leased.push((entry.id.clone(), entry.attempts));
kept.push(entry);
}
state.in_flight = kept;
for id in new_ids {
if budget == 0 {
break;
}
if id.as_str() <= state.hwm.as_str() {
continue; }
state.hwm = id.clone();
state.in_flight.push(InFlight {
id: id.clone(),
attempts: 1,
lease_until_ms: now_ms + lease_ms,
});
out.leased.push((id.clone(), 1));
budget -= 1;
}
out
}
pub fn grouped_message_needed(states: &[GroupState], id: &str) -> bool {
states
.iter()
.any(|s| id > s.hwm.as_str() || s.in_flight.iter().any(|f| f.id == id))
}
pub fn is_direct_child(key: &str, prefix: &str) -> bool {
key.len() > prefix.len() && !key[prefix.len()..].contains('/')
}
#[derive(Default)]
struct TopicHub {
recent: std::collections::VecDeque<StreamEvent>,
subscribers: Vec<futures::channel::mpsc::Sender<StreamEvent>>,
}
const STREAM_RING: usize = 64;
#[derive(Default)]
pub struct StreamHubs {
live: std::sync::Mutex<HashMap<String, TopicHub>>,
}
impl StreamHubs {
pub fn new() -> Self {
Self::default()
}
pub fn broadcast(&self, topic: &str, id: &str, payload: &[u8]) {
let event = StreamEvent {
id: id.to_string(),
payload: payload.to_vec(),
};
let mut live = self.live.lock().unwrap();
let Some(hub) = live.get_mut(topic) else {
return; };
hub.subscribers
.retain_mut(|tx| match tx.try_send(event.clone()) {
Ok(()) => true,
Err(err) => !err.is_disconnected(), });
hub.recent.push_back(event);
while hub.recent.len() > STREAM_RING {
hub.recent.pop_front();
}
if hub.subscribers.is_empty() {
live.remove(topic);
}
}
pub fn subscribe(
&self,
topic: &str,
after: Option<&str>,
) -> futures::stream::BoxStream<'static, StreamEvent> {
let (tx, rx) = futures::channel::mpsc::channel(64);
let replay: Vec<StreamEvent> = {
let mut live = self.live.lock().unwrap();
let hub = live.entry(topic.to_string()).or_default();
let replay = match after {
Some(after) => hub
.recent
.iter()
.filter(|event| event.id.as_str() > after)
.cloned()
.collect(),
None => Vec::new(),
};
hub.subscribers.push(tx);
replay
};
if replay.is_empty() {
rx.boxed()
} else {
futures::stream::iter(replay).chain(rx).boxed()
}
}
}
pub struct LogMessaging {
storage: Arc<dyn Storage>,
kv: Arc<dyn KvStore>,
claim_lock: futures::lock::Mutex<()>,
seq: AtomicU64,
hubs: StreamHubs,
grouped_topics: std::sync::Mutex<Option<std::collections::HashSet<String>>>,
}
pub const GROUP_RETENTION_MS: u64 = 24 * 60 * 60 * 1000;
pub fn id_millis(id: &str) -> u64 {
id.split('-')
.next()
.and_then(|m| m.parse().ok())
.unwrap_or(0)
}
impl LogMessaging {
pub fn new(storage: Arc<dyn Storage>, kv: Arc<dyn KvStore>) -> Self {
Self {
storage,
kv,
claim_lock: futures::lock::Mutex::new(()),
seq: AtomicU64::new(0),
hubs: StreamHubs::new(),
grouped_topics: std::sync::Mutex::new(None),
}
}
async fn topic_has_groups(&self, topic: &str) -> bool {
{
let cache = self.grouped_topics.lock().unwrap();
if let Some(set) = cache.as_ref() {
return set.contains(topic);
}
}
let keys = self.kv.list_prefix("mqgstate/").await.unwrap_or_default();
let mut set = std::collections::HashSet::new();
for key in keys {
if let Some(rest) = key.strip_prefix("mqgstate/") {
if let Some(slash) = rest.rfind('/') {
set.insert(rest[..slash].to_string());
}
}
}
let has = set.contains(topic);
*self.grouped_topics.lock().unwrap() = Some(set);
has
}
fn mark_grouped(&self, topic: &str) {
let mut cache = self.grouped_topics.lock().unwrap();
cache
.get_or_insert_with(std::collections::HashSet::new)
.insert(topic.to_string());
}
async fn read_payload(&self, topic: &str, id: &str) -> Result<Vec<u8>, MessagingError> {
self.read_storage(&payload_key(topic, id)).await
}
async fn read_gpayload(&self, topic: &str, id: &str) -> Result<Vec<u8>, MessagingError> {
self.read_storage(&gpayload_key(topic, id)).await
}
async fn read_storage(&self, key: &str) -> Result<Vec<u8>, MessagingError> {
let object = self
.storage
.get(key)
.await
.map_err(MessagingError::backend)?;
let mut body = object.body;
let mut buf = Vec::new();
while let Some(chunk) = body.next().await {
buf.extend_from_slice(&chunk.map_err(MessagingError::backend)?);
}
Ok(buf)
}
async fn count_direct(&self, prefix: &str) -> Result<usize, MessagingError> {
let keys = self
.kv
.list_prefix(prefix)
.await
.map_err(MessagingError::backend)?;
Ok(keys.iter().filter(|k| is_direct_child(k, prefix)).count())
}
async fn read_logmax(&self, topic: &str) -> Result<String, MessagingError> {
Ok(self
.kv
.get(&logmax_key(topic))
.await
.map_err(MessagingError::backend)?
.map(|raw| String::from_utf8_lossy(&raw).into_owned())
.unwrap_or_default())
}
async fn put_group_state(
&self,
topic: &str,
group: &str,
state: &GroupState,
) -> Result<(), MessagingError> {
let json = serde_json::to_vec(state).map_err(MessagingError::backend)?;
self.kv
.put(&gstate_key(topic, group), json)
.await
.map_err(MessagingError::backend)
}
async fn get_group_state(
&self,
topic: &str,
group: &str,
) -> Result<Option<GroupState>, MessagingError> {
let Some(raw) = self
.kv
.get(&gstate_key(topic, group))
.await
.map_err(MessagingError::backend)?
else {
return Ok(None);
};
serde_json::from_slice(&raw)
.map(Some)
.map_err(|e| MessagingError::Decode(e.to_string()))
}
async fn log_ids_after(
&self,
topic: &str,
after: &str,
limit: usize,
) -> Result<Vec<String>, MessagingError> {
if limit == 0 {
return Ok(Vec::new());
}
let prefix = glog_prefix(topic);
let mut out = Vec::new();
let mut cursor = after.to_string();
loop {
let batch = self
.kv
.list_from(&prefix, &cursor, limit)
.await
.map_err(MessagingError::backend)?;
let Some(last) = batch.last().cloned() else {
break; };
let scanned = batch.len();
for key in batch {
if is_direct_child(&key, &prefix) {
out.push(key[prefix.len()..].to_string());
if out.len() >= limit {
return Ok(out);
}
}
}
cursor = last[prefix.len()..].to_string();
if scanned < limit {
break;
}
}
Ok(out)
}
pub async fn gc_grouped(&self, topic: &str) -> Result<usize, MessagingError> {
let _guard = self.claim_lock.lock().await;
let now = now_unix_ms();
let state_prefix = gstate_prefix(topic);
let state_keys = self
.kv
.list_prefix(&state_prefix)
.await
.map_err(MessagingError::backend)?;
let mut states = Vec::new();
for key in state_keys {
if !is_direct_child(&key, &state_prefix) {
continue;
}
let group = &key[state_prefix.len()..];
if let Some(state) = self.get_group_state(topic, group).await? {
states.push(state);
}
}
let log_prefix = glog_prefix(topic);
let log_keys = self
.kv
.list_prefix(&log_prefix)
.await
.map_err(MessagingError::backend)?;
let mut reclaimed = 0;
for key in log_keys {
if !is_direct_child(&key, &log_prefix) {
continue;
}
let id = &key[log_prefix.len()..];
let needed = grouped_message_needed(&states, id);
let expired = id_millis(id) + GROUP_RETENTION_MS < now;
if !needed || expired {
let _ = self.storage.delete(&gpayload_key(topic, id)).await;
let _ = self.kv.delete(&glog_key(topic, id)).await;
reclaimed += 1;
}
}
Ok(reclaimed)
}
}
#[async_trait]
impl Messaging for LogMessaging {
async fn publish(&self, topic: &str, payload: &[u8]) -> Result<(), MessagingError> {
let id = format!(
"{:013}-{:016x}",
now_unix_ms(),
self.seq.fetch_add(1, Ordering::Relaxed)
);
let bytes = bytes::Bytes::copy_from_slice(payload);
let body = futures::stream::once(async move { Ok(bytes) }).boxed();
self.storage
.put(&payload_key(topic, &id), body, PutMeta::default())
.await
.map_err(MessagingError::backend)?;
let json = serde_json::to_vec(&Record::fresh()).map_err(MessagingError::backend)?;
self.kv
.put(&meta_key(topic, &id), json)
.await
.map_err(MessagingError::backend)?;
if self.topic_has_groups(topic).await {
let bytes = bytes::Bytes::copy_from_slice(payload);
let body = futures::stream::once(async move { Ok(bytes) }).boxed();
self.storage
.put(&gpayload_key(topic, &id), body, PutMeta::default())
.await
.map_err(MessagingError::backend)?;
self.kv
.put(&glog_key(topic, &id), Vec::new())
.await
.map_err(MessagingError::backend)?;
let cur = self
.kv
.get(&logmax_key(topic))
.await
.map_err(MessagingError::backend)?
.map(|v| String::from_utf8_lossy(&v).into_owned())
.unwrap_or_default();
if id.as_str() > cur.as_str() {
self.kv
.put(&logmax_key(topic), id.clone().into_bytes())
.await
.map_err(MessagingError::backend)?;
}
}
self.hubs.broadcast(topic, &id, payload);
Ok(())
}
async fn claim(
&self,
topic: &str,
lease: Duration,
max_batch: usize,
max_attempts: u32,
) -> Result<Vec<ClaimedMessage>, MessagingError> {
let _guard = self.claim_lock.lock().await;
let now = now_unix_ms();
let prefix = meta_prefix(topic);
let keys = self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?;
let mut records = Vec::new();
for key in keys {
if !is_direct_child(&key, &prefix) {
continue; }
let Some(raw) = self.kv.get(&key).await.map_err(MessagingError::backend)? else {
continue; };
let record: Record =
serde_json::from_slice(&raw).map_err(|e| MessagingError::Decode(e.to_string()))?;
records.push((key[prefix.len()..].to_string(), record));
}
let actions = plan_claim(
records,
now,
lease.as_millis() as u64,
max_batch,
max_attempts,
);
let mut claimed = Vec::new();
for action in actions {
match action {
ClaimAction::Lease { id, record } => {
let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
self.kv
.put(&meta_key(topic, &id), json)
.await
.map_err(MessagingError::backend)?;
let payload = self.read_payload(topic, &id).await?;
claimed.push(ClaimedMessage {
id,
topic: topic.to_string(),
payload,
attempts: record.attempts,
group: String::new(),
});
}
ClaimAction::DeadLetter { id, record } => {
let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
self.kv
.put(&dead_key(topic, &id), json)
.await
.map_err(MessagingError::backend)?;
self.kv
.delete(&meta_key(topic, &id))
.await
.map_err(MessagingError::backend)?;
}
}
}
Ok(claimed)
}
async fn claim_grouped(
&self,
topic: &str,
group: &str,
start: StartPosition,
lease: Duration,
max_batch: usize,
max_attempts: u32,
) -> Result<Vec<ClaimedMessage>, MessagingError> {
if group.is_empty() {
return self.claim(topic, lease, max_batch, max_attempts).await;
}
let _guard = self.claim_lock.lock().await;
let now = now_unix_ms();
let lease_ms = lease.as_millis() as u64;
let (mut state, existed) = match self.get_group_state(topic, group).await? {
Some(state) => (state, true),
None => {
self.mark_grouped(topic);
let hwm = match start {
StartPosition::Latest => self.read_logmax(topic).await?,
StartPosition::Earliest => String::new(),
};
(GroupState::new(hwm), false)
}
};
let new_ids = if state.hwm.as_str() < self.read_logmax(topic).await?.as_str() {
self.log_ids_after(topic, &state.hwm, max_batch).await?
} else {
Vec::new()
};
let plan = plan_claim_grouped(&mut state, now, lease_ms, max_batch, max_attempts, &new_ids);
for (id, attempts) in &plan.dead {
let record = Record {
version: crate::SCHEMA_VERSION,
attempts: *attempts,
lease_until_ms: 0,
};
let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
self.kv
.put(&gdead_key(topic, group, id), json)
.await
.map_err(MessagingError::backend)?;
}
let changed = !existed || !plan.leased.is_empty() || !plan.dead.is_empty();
let mut claimed = Vec::new();
for (id, attempts) in plan.leased {
match self.read_gpayload(topic, &id).await {
Ok(payload) => claimed.push(ClaimedMessage {
id,
topic: topic.to_string(),
payload,
attempts,
group: group.to_string(),
}),
Err(_) => continue,
}
}
if changed {
self.put_group_state(topic, group, &state).await?;
}
Ok(claimed)
}
async fn ack(&self, msg: &ClaimedMessage) -> Result<(), MessagingError> {
if !msg.group.is_empty() {
let _guard = self.claim_lock.lock().await;
let Some(mut state) = self.get_group_state(&msg.topic, &msg.group).await? else {
return Ok(()); };
let before = state.in_flight.len();
state.in_flight.retain(|f| f.id != msg.id);
if state.in_flight.len() != before {
self.put_group_state(&msg.topic, &msg.group, &state).await?;
}
return Ok(());
}
self.kv
.delete(&meta_key(&msg.topic, &msg.id))
.await
.map_err(MessagingError::backend)?;
self.storage
.delete(&payload_key(&msg.topic, &msg.id))
.await
.map_err(MessagingError::backend)?;
Ok(())
}
async fn backlog(&self, topic: &str) -> Result<usize, MessagingError> {
self.count_direct(&meta_prefix(topic)).await
}
async fn dead_letter_count(&self, topic: &str) -> Result<usize, MessagingError> {
self.count_direct(&dead_prefix(topic)).await
}
async fn nack(&self, msg: &ClaimedMessage) -> Result<(), MessagingError> {
if !msg.group.is_empty() {
let _guard = self.claim_lock.lock().await;
let Some(mut state) = self.get_group_state(&msg.topic, &msg.group).await? else {
return Ok(()); };
let mut changed = false;
for entry in &mut state.in_flight {
if entry.id == msg.id {
entry.lease_until_ms = 0;
changed = true;
break;
}
}
if changed {
self.put_group_state(&msg.topic, &msg.group, &state).await?;
}
return Ok(());
}
let key = meta_key(&msg.topic, &msg.id);
let Some(raw) = self.kv.get(&key).await.map_err(MessagingError::backend)? else {
return Ok(()); };
let mut record: Record =
serde_json::from_slice(&raw).map_err(|e| MessagingError::Decode(e.to_string()))?;
record.lease_until_ms = 0; let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
self.kv
.put(&key, json)
.await
.map_err(MessagingError::backend)?;
Ok(())
}
async fn purge_dead_letters(&self, topic: &str) -> Result<usize, MessagingError> {
let prefix = dead_prefix(topic);
let keys = self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?;
let mut purged = 0;
for key in keys {
if !is_direct_child(&key, &prefix) {
continue; }
let id = &key[prefix.len()..];
self.storage
.delete(&payload_key(topic, id))
.await
.map_err(MessagingError::backend)?;
self.kv
.delete(&key)
.await
.map_err(MessagingError::backend)?;
purged += 1;
}
Ok(purged)
}
async fn redrive_dead_letters(&self, topic: &str) -> Result<usize, MessagingError> {
let prefix = dead_prefix(topic);
let keys = self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?;
let mut redriven = 0;
for key in keys {
if !is_direct_child(&key, &prefix) {
continue;
}
let id = &key[prefix.len()..];
let json = serde_json::to_vec(&Record::fresh()).map_err(MessagingError::backend)?;
self.kv
.put(&meta_key(topic, id), json)
.await
.map_err(MessagingError::backend)?;
self.kv
.delete(&key)
.await
.map_err(MessagingError::backend)?;
redriven += 1;
}
Ok(redriven)
}
async fn retention_sweep(&self, topic: &str) -> Result<usize, MessagingError> {
self.gc_grouped(topic).await
}
fn subscribe(
&self,
topic: &str,
after: Option<&str>,
) -> futures::stream::BoxStream<'static, StreamEvent> {
self.hubs.subscribe(topic, after)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::kv::MemoryKv;
use crate::{ByteStream, GetObject, ObjectMeta, StorageError};
use std::collections::HashMap;
use std::sync::Mutex;
#[derive(Default)]
struct MemStorage {
objects: Mutex<HashMap<String, Vec<u8>>>,
}
#[async_trait]
impl Storage for MemStorage {
async fn get(&self, key: &str) -> Result<GetObject, StorageError> {
let bytes = self
.objects
.lock()
.unwrap()
.get(key)
.cloned()
.ok_or_else(|| StorageError::NotFound(key.to_string()))?;
let size = bytes.len() as u64;
let body: ByteStream =
futures::stream::once(async move { Ok(bytes::Bytes::from(bytes)) }).boxed();
Ok(GetObject {
meta: ObjectMeta {
key: key.to_string(),
size: Some(size),
..Default::default()
},
body,
})
}
async fn get_range(
&self,
key: &str,
_: u64,
_: Option<u64>,
) -> Result<GetObject, StorageError> {
self.get(key).await
}
async fn put(
&self,
key: &str,
mut body: ByteStream,
_: PutMeta,
) -> Result<ObjectMeta, StorageError> {
let mut buf = Vec::new();
while let Some(chunk) = body.next().await {
buf.extend_from_slice(&chunk?);
}
let size = buf.len() as u64;
self.objects.lock().unwrap().insert(key.to_string(), buf);
Ok(ObjectMeta {
key: key.to_string(),
size: Some(size),
..Default::default()
})
}
async fn head(&self, key: &str) -> Result<ObjectMeta, StorageError> {
let map = self.objects.lock().unwrap();
let bytes = map
.get(key)
.ok_or_else(|| StorageError::NotFound(key.to_string()))?;
Ok(ObjectMeta {
key: key.to_string(),
size: Some(bytes.len() as u64),
..Default::default()
})
}
async fn delete(&self, key: &str) -> Result<(), StorageError> {
self.objects.lock().unwrap().remove(key);
Ok(())
}
async fn list(&self, prefix: &str) -> Result<Vec<ObjectMeta>, StorageError> {
Ok(self
.objects
.lock()
.unwrap()
.keys()
.filter(|k| k.starts_with(prefix))
.map(|k| ObjectMeta {
key: k.clone(),
..Default::default()
})
.collect())
}
}
fn mq() -> LogMessaging {
LogMessaging::new(Arc::new(MemStorage::default()), Arc::new(MemoryKv::new()))
}
const LEASE: Duration = Duration::from_secs(30);
fn payloads(msgs: &[ClaimedMessage]) -> Vec<Vec<u8>> {
msgs.iter().map(|m| m.payload.clone()).collect()
}
#[tokio::test]
async fn consumer_groups_fan_out_and_ack_independently() {
let mq = mq();
let t = "bus/orders";
assert!(mq
.claim_grouped(t, "billing", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
assert!(mq
.claim_grouped(t, "audit", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
mq.publish(t, b"a").await.unwrap();
mq.publish(t, b"b").await.unwrap();
let billing = mq
.claim_grouped(t, "billing", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap();
assert_eq!(payloads(&billing), vec![b"a".to_vec(), b"b".to_vec()]);
let audit = mq
.claim_grouped(t, "audit", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap();
assert_eq!(payloads(&audit), vec![b"a".to_vec(), b"b".to_vec()]);
for m in &billing {
mq.ack(m).await.unwrap();
}
assert!(mq
.claim_grouped(t, "billing", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
for m in &audit {
mq.nack(m).await.unwrap();
}
let audit_again = mq
.claim_grouped(t, "audit", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap();
assert_eq!(payloads(&audit_again), vec![b"a".to_vec(), b"b".to_vec()]);
}
#[tokio::test]
async fn consumer_group_start_position_latest_vs_earliest() {
let mq = mq();
let t = "bus/events";
assert!(mq
.claim_grouped(t, "seed", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
mq.publish(t, b"a").await.unwrap();
mq.publish(t, b"b").await.unwrap();
let replay = mq
.claim_grouped(t, "replay", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap();
assert_eq!(payloads(&replay), vec![b"a".to_vec(), b"b".to_vec()]);
let live = mq
.claim_grouped(t, "live", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap();
assert!(live.is_empty());
mq.publish(t, b"c").await.unwrap();
let live_after = mq
.claim_grouped(t, "live", StartPosition::Latest, LEASE, 10, 5)
.await
.unwrap();
assert_eq!(payloads(&live_after), vec![b"c".to_vec()]);
}
#[tokio::test]
async fn consumer_group_batches_backlog_by_max_batch() {
let mq = mq();
let t = "bus/jobs";
assert!(mq
.claim_grouped(t, "worker", StartPosition::Earliest, LEASE, 2, 5)
.await
.unwrap()
.is_empty());
for n in 0..5u8 {
mq.publish(t, &[b'0' + n]).await.unwrap();
}
let first = mq
.claim_grouped(t, "worker", StartPosition::Earliest, LEASE, 2, 5)
.await
.unwrap();
assert_eq!(payloads(&first), vec![b"0".to_vec(), b"1".to_vec()]);
let second = mq
.claim_grouped(t, "worker", StartPosition::Earliest, LEASE, 2, 5)
.await
.unwrap();
assert_eq!(payloads(&second), vec![b"2".to_vec(), b"3".to_vec()]);
let third = mq
.claim_grouped(t, "worker", StartPosition::Earliest, LEASE, 2, 5)
.await
.unwrap();
assert_eq!(payloads(&third), vec![b"4".to_vec()]);
assert!(mq
.claim_grouped(t, "worker", StartPosition::Earliest, LEASE, 2, 5)
.await
.unwrap()
.is_empty());
}
#[tokio::test]
async fn consumer_group_dead_letters_after_max_attempts() {
let mq = mq();
let t = "bus/flaky";
assert!(mq
.claim_grouped(t, "g", StartPosition::Earliest, LEASE, 10, 2)
.await
.unwrap()
.is_empty());
mq.publish(t, b"x").await.unwrap();
for expected in 1..=2 {
let batch = mq
.claim_grouped(t, "g", StartPosition::Earliest, Duration::ZERO, 10, 2)
.await
.unwrap();
assert_eq!(batch.len(), 1, "attempt {expected}");
assert_eq!(batch[0].attempts, expected);
}
assert!(mq
.claim_grouped(t, "g", StartPosition::Earliest, Duration::ZERO, 10, 2)
.await
.unwrap()
.is_empty());
assert!(mq
.claim_grouped(t, "g", StartPosition::Earliest, Duration::ZERO, 10, 2)
.await
.unwrap()
.is_empty());
}
#[tokio::test]
async fn consumer_group_survives_restart() {
let storage: Arc<dyn Storage> = Arc::new(MemStorage::default());
let kv: Arc<dyn KvStore> = Arc::new(MemoryKv::new());
let t = "bus/resume";
{
let mq = LogMessaging::new(storage.clone(), kv.clone());
assert!(mq
.claim_grouped(t, "g", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
mq.publish(t, b"a").await.unwrap();
mq.publish(t, b"b").await.unwrap();
let batch = mq
.claim_grouped(t, "g", StartPosition::Earliest, Duration::ZERO, 10, 5)
.await
.unwrap();
assert_eq!(payloads(&batch), vec![b"a".to_vec(), b"b".to_vec()]);
mq.ack(&batch[0]).await.unwrap(); }
let mq = LogMessaging::new(storage, kv);
let redelivered = mq
.claim_grouped(t, "g", StartPosition::Earliest, Duration::ZERO, 10, 5)
.await
.unwrap();
assert_eq!(payloads(&redelivered), vec![b"b".to_vec()]);
assert_eq!(
redelivered[0].attempts, 2,
"redelivery re-charges the attempt"
);
}
#[tokio::test]
async fn gc_grouped_reclaims_only_fully_consumed_messages() {
let storage: Arc<dyn Storage> = Arc::new(MemStorage::default());
let kv: Arc<dyn KvStore> = Arc::new(MemoryKv::new());
let mq = LogMessaging::new(storage.clone(), kv);
let t = "bus/retain";
for g in ["one", "two"] {
assert!(mq
.claim_grouped(t, g, StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
}
mq.publish(t, b"a").await.unwrap();
mq.publish(t, b"b").await.unwrap();
let one = mq
.claim_grouped(t, "one", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap();
for m in &one {
mq.ack(m).await.unwrap();
}
assert_eq!(mq.gc_grouped(t).await.unwrap(), 0);
let two = mq
.claim_grouped(t, "two", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap();
assert_eq!(payloads(&two), vec![b"a".to_vec(), b"b".to_vec()]);
for m in &two {
mq.ack(m).await.unwrap();
}
assert_eq!(mq.gc_grouped(t).await.unwrap(), 2);
let ids: Vec<String> = one.iter().map(|m| m.id.clone()).collect();
for id in &ids {
assert!(
storage.head(&gpayload_key(t, id)).await.is_err(),
"reclaimed payload for {id}"
);
}
assert!(mq
.claim_grouped(t, "one", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
}
#[tokio::test]
async fn publish_claim_ack_roundtrip_and_fifo() {
let mq = mq();
mq.publish("orders/created", b"a").await.unwrap();
mq.publish("orders/created", b"b").await.unwrap();
let batch = mq.claim("orders/created", LEASE, 10, 5).await.unwrap();
assert_eq!(batch.len(), 2);
assert_eq!(batch[0].payload, b"a");
assert_eq!(batch[1].payload, b"b");
assert_eq!(batch[0].attempts, 1);
assert!(mq
.claim("orders/created", LEASE, 10, 5)
.await
.unwrap()
.is_empty());
for m in &batch {
mq.ack(m).await.unwrap();
}
assert!(mq
.claim("orders/created", LEASE, 10, 5)
.await
.unwrap()
.is_empty());
}
#[tokio::test]
async fn topic_scoping_excludes_subtopics() {
let mq = mq();
mq.publish("orders", b"top").await.unwrap();
mq.publish("orders/created", b"sub").await.unwrap();
let batch = mq.claim("orders", LEASE, 10, 5).await.unwrap();
assert_eq!(batch.len(), 1);
assert_eq!(batch[0].payload, b"top");
}
#[tokio::test]
async fn lease_expiry_redelivers() {
let mq = mq();
mq.publish("t", b"x").await.unwrap();
let first = mq.claim("t", Duration::ZERO, 10, 5).await.unwrap();
assert_eq!(first.len(), 1);
assert_eq!(first[0].attempts, 1);
let second = mq.claim("t", LEASE, 10, 5).await.unwrap();
assert_eq!(second.len(), 1);
assert_eq!(second[0].attempts, 2); }
#[tokio::test]
async fn nack_makes_claimable_again() {
let mq = mq();
mq.publish("t", b"x").await.unwrap();
let m = mq.claim("t", LEASE, 10, 5).await.unwrap().pop().unwrap();
mq.nack(&m).await.unwrap();
let again = mq.claim("t", LEASE, 10, 5).await.unwrap();
assert_eq!(again.len(), 1);
assert_eq!(again[0].attempts, 2);
}
#[tokio::test]
async fn subscribe_receives_live_broadcast() {
use futures::StreamExt;
let mq = mq();
let mut sub = mq.subscribe("events", None);
mq.publish("events", b"hello").await.unwrap();
mq.publish("events", b"world").await.unwrap();
assert_eq!(sub.next().await.unwrap().payload, b"hello");
assert_eq!(sub.next().await.unwrap().payload, b"world");
mq.publish("other", b"nope").await.unwrap();
mq.publish("events", b"again").await.unwrap();
assert_eq!(sub.next().await.unwrap().payload, b"again");
}
#[tokio::test]
async fn last_event_id_replays_recent_then_goes_live() {
use futures::StreamExt;
let mq = mq();
let mut keepalive = mq.subscribe("events", None);
mq.publish("events", b"one").await.unwrap();
mq.publish("events", b"two").await.unwrap();
mq.publish("events", b"three").await.unwrap();
let first = keepalive.next().await.unwrap();
assert_eq!(first.payload, b"one");
let mut resumed = mq.subscribe("events", Some(&first.id));
assert_eq!(resumed.next().await.unwrap().payload, b"two");
assert_eq!(resumed.next().await.unwrap().payload, b"three");
mq.publish("events", b"four").await.unwrap();
assert_eq!(resumed.next().await.unwrap().payload, b"four");
}
#[tokio::test]
async fn dropped_subscriber_is_pruned_without_error() {
let mq = mq();
{
let _sub = mq.subscribe("events", None);
} mq.publish("events", b"x").await.unwrap();
}
#[tokio::test]
async fn dead_letters_after_max_attempts() {
let mq = mq();
mq.publish("t", b"x").await.unwrap();
for expected in 1..=2 {
let m = mq.claim("t", Duration::ZERO, 10, 2).await.unwrap();
assert_eq!(m.len(), 1, "attempt {expected}");
assert_eq!(m[0].attempts, expected);
}
let exhausted = mq.claim("t", Duration::ZERO, 10, 2).await.unwrap();
assert!(
exhausted.is_empty(),
"should dead-letter, not deliver a 3rd time"
);
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 1);
}
#[tokio::test]
async fn purge_dead_letters_clears_records_and_payloads() {
let storage: Arc<dyn Storage> = Arc::new(MemStorage::default());
let kv: Arc<dyn KvStore> = Arc::new(MemoryKv::new());
let mq = LogMessaging::new(storage.clone(), kv);
mq.publish("t", b"x").await.unwrap();
let id = mq.claim("t", Duration::ZERO, 10, 1).await.unwrap()[0]
.id
.clone();
assert!(mq
.claim("t", Duration::ZERO, 10, 1)
.await
.unwrap()
.is_empty());
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 1);
let purged = mq.purge_dead_letters("t").await.unwrap();
assert_eq!(purged, 1);
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 0);
assert!(
storage.head(&payload_key("t", &id)).await.is_err(),
"purge frees the dead-lettered payload"
);
}
#[tokio::test]
async fn redrive_dead_letters_requeues_with_fresh_attempts() {
let mq = mq();
mq.publish("t", b"x").await.unwrap();
assert_eq!(mq.claim("t", Duration::ZERO, 10, 1).await.unwrap().len(), 1);
assert!(mq
.claim("t", Duration::ZERO, 10, 1)
.await
.unwrap()
.is_empty());
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 1);
let redriven = mq.redrive_dead_letters("t").await.unwrap();
assert_eq!(redriven, 1);
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 0);
assert_eq!(mq.backlog("t").await.unwrap(), 1);
let again = mq.claim("t", LEASE, 10, 5).await.unwrap();
assert_eq!(again.len(), 1);
assert_eq!(again[0].payload, b"x");
assert_eq!(again[0].attempts, 1, "fresh attempts after redrive");
}
#[tokio::test]
async fn survives_restart_over_shared_backends() {
let storage: Arc<dyn Storage> = Arc::new(MemStorage::default());
let kv: Arc<dyn KvStore> = Arc::new(MemoryKv::new());
{
let mq = LogMessaging::new(storage.clone(), kv.clone());
mq.publish("orders", b"a").await.unwrap();
mq.publish("orders", b"b").await.unwrap();
let batch = mq.claim("orders", Duration::ZERO, 10, 5).await.unwrap();
assert_eq!(batch.len(), 2);
mq.ack(&batch[0]).await.unwrap(); }
let mq = LogMessaging::new(storage, kv);
let batch = mq.claim("orders", LEASE, 10, 5).await.unwrap();
assert_eq!(batch.len(), 1, "only the un-acked message survives");
assert_eq!(batch[0].payload, b"b");
assert_eq!(batch[0].attempts, 2, "redelivery re-charges the attempt");
}
}