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, WriteOp};
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 signed_context: Option<String>,
pub inline: bool,
}
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),
#[error("messaging operation not supported by this backend: {0}")]
Unsupported(String),
#[error("publish rejected: topic {0:?} backlog is at its configured max_depth")]
DepthExceeded(String),
#[error("publish rejected: topic {0:?} exceeded its configured max_rate_per_sec")]
RateExceeded(String),
}
impl MessagingError {
fn backend<E: std::fmt::Display>(err: E) -> Self {
Self::Backend(err.to_string())
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TopicPolicy {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_depth: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_rate_per_sec: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_unflushed: Option<usize>,
}
#[async_trait]
pub trait Messaging: Send + Sync {
async fn publish(&self, topic: &str, payload: &[u8]) -> Result<(), MessagingError>;
async fn publish_ctx(
&self,
topic: &str,
payload: &[u8],
_signed_context: Option<&str>,
) -> Result<(), MessagingError> {
self.publish(topic, payload).await
}
async fn publish_batch_ctx(
&self,
messages: &[(String, Vec<u8>)],
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
for (topic, payload) in messages {
self.publish_ctx(topic, payload, signed_context).await?;
}
Ok(())
}
async fn publish_delayed_ctx(
&self,
topic: &str,
payload: &[u8],
_delay: Duration,
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
self.publish_ctx(topic, payload, signed_context).await
}
async fn publish_with_ttl_ctx(
&self,
topic: &str,
payload: &[u8],
_ttl: Duration,
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
self.publish_ctx(topic, payload, signed_context).await
}
async fn publish_with_priority_ctx(
&self,
topic: &str,
payload: &[u8],
_priority: u8,
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
self.publish_ctx(topic, payload, signed_context).await
}
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 nack_after(&self, msg: &ClaimedMessage, delay_ms: u64) -> Result<(), MessagingError> {
let _ = delay_ms;
self.nack(msg).await
}
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 oldest_pending_ms(&self, _topic: &str) -> Result<Option<u64>, MessagingError> {
Ok(None)
}
async fn in_flight_count(&self, _topic: &str) -> Result<usize, MessagingError> {
Ok(0)
}
async fn group_lag(&self, _topic: &str, _group: &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 set_last_error(
&self,
_msg: &ClaimedMessage,
_reason: &str,
) -> Result<(), MessagingError> {
Ok(())
}
async fn list_dead_letters(
&self,
_topic: &str,
_filter: &DeadLetterFilter,
) -> Result<Vec<DeadLetter>, MessagingError> {
Ok(Vec::new())
}
async fn show_dead_letter(
&self,
_topic: &str,
_group: &str,
_id: &str,
) -> Result<Option<DeadLetter>, MessagingError> {
Ok(None)
}
async fn redrive_dead_letters_filtered(
&self,
_topic: &str,
_filter: &DeadLetterFilter,
) -> Result<usize, MessagingError> {
Ok(0)
}
async fn discard_dead_letters(
&self,
_topic: &str,
_filter: &DeadLetterFilter,
) -> Result<usize, MessagingError> {
Ok(0)
}
async fn peek(
&self,
_topic: &str,
_limit: usize,
) -> Result<Vec<PeekedMessage>, MessagingError> {
Ok(Vec::new())
}
async fn replay(
&self,
_topic: &str,
_after: Option<&str>,
_limit: usize,
) -> Result<Vec<PeekedMessage>, MessagingError> {
Ok(Vec::new())
}
async fn list_groups(&self, _topic: &str) -> Result<Vec<GroupInfo>, MessagingError> {
Ok(Vec::new())
}
async fn reset_group(
&self,
_topic: &str,
_group: &str,
_start: StartPosition,
) -> Result<(), MessagingError> {
Err(MessagingError::Backend(
"this messaging backend does not support consumer groups".into(),
))
}
async fn delete_group(&self, _topic: &str, _group: &str) -> Result<(), MessagingError> {
Ok(())
}
async fn set_paused(&self, _topic: &str, _paused: bool) -> Result<(), MessagingError> {
Ok(())
}
async fn is_paused(&self, _topic: &str) -> Result<bool, MessagingError> {
Ok(false)
}
async fn set_topic_policy(
&self,
_topic: &str,
_policy: TopicPolicy,
) -> Result<(), MessagingError> {
Err(MessagingError::Unsupported(
"per-topic policy is not supported by this messaging backend".into(),
))
}
async fn topic_policy(&self, _topic: &str) -> Result<Option<TopicPolicy>, MessagingError> {
Ok(None)
}
async fn retention_sweep(
&self,
_topic: &str,
_retention_ms: u64,
) -> Result<usize, MessagingError> {
Ok(0)
}
fn supports_ready_set(&self) -> bool {
false
}
async fn ready_topics(&self) -> Result<Vec<String>, MessagingError> {
Err(MessagingError::Unsupported(
"this messaging backend has no ready-set (event-driven delivery)".into(),
))
}
async fn rebuild_ready_set(&self) -> Result<usize, MessagingError> {
Ok(0)
}
async fn due_topics(&self) -> Result<Vec<(String, u64)>, MessagingError> {
Ok(Vec::new())
}
async fn delivery_stats(&self, _this_node: &str) -> Result<DeliveryStats, MessagingError> {
Err(MessagingError::Unsupported(
"this messaging backend has no delivery-plane stats (event-driven delivery)".into(),
))
}
async fn shard_owns(&self, _topic: &str) -> bool {
true
}
async fn shard_owned(&self, topics: Vec<String>) -> Vec<String> {
topics
}
fn subscribe(
&self,
_topic: &str,
_after: Option<&str>,
) -> futures::stream::BoxStream<'static, StreamEvent> {
futures::stream::empty().boxed()
}
fn wake_handle(&self) -> Option<Wake> {
None
}
}
#[derive(Clone)]
pub struct Wake {
inner: Arc<WakeInner>,
}
struct WakeInner {
pending: std::sync::atomic::AtomicBool,
waker: futures::task::AtomicWaker,
}
impl Wake {
pub fn new() -> Self {
Self {
inner: Arc::new(WakeInner {
pending: std::sync::atomic::AtomicBool::new(false),
waker: futures::task::AtomicWaker::new(),
}),
}
}
pub fn notify(&self) {
self.inner
.pending
.store(true, std::sync::atomic::Ordering::Release);
self.inner.waker.wake();
}
pub async fn notified(&self) {
futures::future::poll_fn(|cx| {
if self
.inner
.pending
.swap(false, std::sync::atomic::Ordering::AcqRel)
{
return std::task::Poll::Ready(());
}
self.inner.waker.register(cx.waker());
if self
.inner
.pending
.swap(false, std::sync::atomic::Ordering::AcqRel)
{
std::task::Poll::Ready(())
} else {
std::task::Poll::Pending
}
})
.await;
}
}
impl Default for Wake {
fn default() -> Self {
Self::new()
}
}
#[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,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub signed_context: Option<String>,
#[serde(default, with = "inline_b64", skip_serializing_if = "Option::is_none")]
pub inline: Option<Vec<u8>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_error: Option<String>,
#[serde(default, skip_serializing_if = "is_zero_u64")]
pub expires_at_ms: u64,
#[serde(default, skip_serializing_if = "is_zero_u8")]
pub priority: u8,
}
fn is_zero_u8(v: &u8) -> bool {
*v == 0
}
fn is_zero_u64(v: &u64) -> bool {
*v == 0
}
mod inline_b64 {
use base64::Engine as _;
use serde::{Deserialize, Deserializer, Serializer};
pub fn serialize<S: Serializer>(v: &Option<Vec<u8>>, s: S) -> Result<S::Ok, S::Error> {
match v {
Some(bytes) => {
s.serialize_str(&base64::engine::general_purpose::STANDARD.encode(bytes))
}
None => s.serialize_none(),
}
}
pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Option<Vec<u8>>, D::Error> {
let opt = Option::<String>::deserialize(d)?;
match opt {
Some(text) => base64::engine::general_purpose::STANDARD
.decode(text.as_bytes())
.map(Some)
.map_err(serde::de::Error::custom),
None => Ok(None),
}
}
}
pub const INLINE_MAX: usize = 4096;
pub const INLINE_INFLIGHT_MAX_BYTES: usize = 32 * 1024 * 1024;
const GROUP_COMMIT_MAX: usize = 512;
struct PublishJob {
ops: Vec<WriteOp>,
msgs: usize,
max_unflushed: usize,
done: futures::channel::oneshot::Sender<Result<(), MessagingError>>,
}
impl Record {
pub fn fresh(signed_context: Option<String>) -> Self {
Self {
version: crate::SCHEMA_VERSION,
attempts: 0,
lease_until_ms: 0,
signed_context,
inline: None,
last_error: None,
expires_at_ms: 0,
priority: 0,
}
}
}
pub const LAST_ERROR_MAX: usize = 256;
pub fn sanitize_reason(reason: &str) -> String {
let mut out = String::new();
for c in reason.chars() {
let c = if c.is_control() { ' ' } else { c };
if out.len() + c.len_utf8() > LAST_ERROR_MAX {
break;
}
out.push(c);
}
out.trim().to_string()
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct DeadLetter {
pub id: String,
pub group: String,
pub attempts: u32,
pub last_error: Option<String>,
pub signed_context: Option<String>,
pub payload: Option<Vec<u8>>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct PeekedMessage {
pub id: String,
pub attempts: u32,
pub leased: bool,
pub signed_context: Option<String>,
pub payload: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct GroupInfo {
pub group: String,
pub hwm: String,
pub in_flight: usize,
pub lag: usize,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct DeliveryStats {
pub ready_set_size: usize,
pub due_heap_depth: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_rebuild_age_ms: Option<u64>,
#[serde(default, skip_serializing_if = "String::is_empty")]
pub this_node: String,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct DeadLetterFilter {
pub id: Option<String>,
pub group: Option<String>,
pub older_than_ms: Option<u64>,
pub match_last_error: Option<String>,
pub limit: Option<usize>,
}
impl DeadLetterFilter {
pub fn matches(&self, dl: &DeadLetter, now_ms: u64) -> bool {
if let Some(id) = &self.id {
if &dl.id != id {
return false;
}
}
if let Some(group) = &self.group {
if &dl.group != group {
return false;
}
}
if let Some(older) = self.older_than_ms {
match id_age_ms(&dl.id, now_ms) {
Some(age) if age >= older => {}
_ => return false,
}
}
if let Some(needle) = &self.match_last_error {
match &dl.last_error {
Some(err) if err.contains(needle.as_str()) => {}
_ => return false,
}
}
true
}
}
fn id_age_ms(id: &str, now_ms: u64) -> Option<u64> {
let millis: u64 = id.split('-').next()?.parse().ok()?;
Some(now_ms.saturating_sub(millis))
}
#[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| b.1.priority.cmp(&a.1.priority).then_with(|| 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.expires_at_ms != 0 && record.expires_at_ms <= now_ms {
record.last_error = Some("ttl-expired".to_string());
actions.push(ClaimAction::DeadLetter { id, record });
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 pause_key(topic: &str) -> String {
format!("mqpause/{topic}")
}
pub fn mqpolicy_key(topic: &str) -> String {
format!("mqpolicy/{topic}")
}
pub fn ready_key(topic: &str) -> String {
format!("mqready/{topic}")
}
pub const READY_PREFIX: &str = "mqready/";
pub fn topic_of_ready_key(key: &str) -> Option<&str> {
key.strip_prefix(READY_PREFIX).filter(|t| !t.is_empty())
}
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}")
}
pub fn gdead_topic_prefix(topic: &str) -> String {
format!("mqgd/{topic}/")
}
pub fn split_group_id(suffix: &str) -> Option<(&str, &str)> {
let (group, id) = suffix.split_once('/')?;
if group.is_empty() || id.is_empty() || id.contains('/') {
return None;
}
Some((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>>>,
inline_inflight_bytes: std::sync::atomic::AtomicUsize,
inline_budget_bytes: usize,
commit_queue: std::sync::Mutex<Vec<PublishJob>>,
commit_gate: futures::lock::Mutex<()>,
max_unflushed: usize,
unflushed: std::sync::atomic::AtomicUsize,
policy_cache: std::sync::Mutex<HashMap<String, Option<TopicPolicy>>>,
rate_buckets: std::sync::Mutex<HashMap<String, TokenBucket>>,
wake: Wake,
last_rebuild_ms: AtomicU64,
}
#[derive(Debug, Clone, Copy)]
struct TokenBucket {
tokens: f64,
last_refill_ms: u64,
}
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),
inline_inflight_bytes: std::sync::atomic::AtomicUsize::new(0),
inline_budget_bytes: INLINE_INFLIGHT_MAX_BYTES,
commit_queue: std::sync::Mutex::new(Vec::new()),
commit_gate: futures::lock::Mutex::new(()),
wake: Wake::new(),
last_rebuild_ms: AtomicU64::new(0),
max_unflushed: 0, unflushed: std::sync::atomic::AtomicUsize::new(0),
policy_cache: std::sync::Mutex::new(HashMap::new()),
rate_buckets: std::sync::Mutex::new(HashMap::new()),
}
}
#[must_use]
pub fn with_max_unflushed(mut self, max: usize) -> Self {
self.max_unflushed = max;
self
}
async fn resolve_policy(&self, topic: &str) -> Result<Option<TopicPolicy>, MessagingError> {
{
let cache = self.policy_cache.lock().unwrap();
if let Some(hit) = cache.get(topic) {
return Ok(hit.clone());
}
}
let policy = match self
.kv
.get(&mqpolicy_key(topic))
.await
.map_err(MessagingError::backend)?
{
Some(raw) => Some(
serde_json::from_slice::<TopicPolicy>(&raw)
.map_err(|e| MessagingError::Decode(e.to_string()))?,
),
None => None,
};
self.policy_cache
.lock()
.unwrap()
.insert(topic.to_string(), policy.clone());
Ok(policy)
}
async fn enforce_publish_policy(
&self,
topic: &str,
policy: &TopicPolicy,
n: usize,
) -> Result<(), MessagingError> {
if let Some(max_depth) = policy.max_depth {
let backlog = self.backlog(topic).await?;
if backlog >= max_depth {
return Err(MessagingError::DepthExceeded(topic.to_string()));
}
}
if let Some(rate) = policy.max_rate_per_sec {
if !self.try_take_tokens(topic, rate, n) {
return Err(MessagingError::RateExceeded(topic.to_string()));
}
}
Ok(())
}
async fn enforce_and_resolve_budget(
&self,
topic: &str,
n: usize,
) -> Result<usize, MessagingError> {
match self.resolve_policy(topic).await? {
Some(p) => {
self.enforce_publish_policy(topic, &p, n).await?;
Ok(p.max_unflushed.unwrap_or(self.max_unflushed))
}
None => Ok(self.max_unflushed),
}
}
fn try_take_tokens(&self, topic: &str, rate: u32, n: usize) -> bool {
let now = now_unix_ms();
let cap = f64::from(rate);
let mut buckets = self.rate_buckets.lock().unwrap();
let bucket = buckets.entry(topic.to_string()).or_insert(TokenBucket {
tokens: cap,
last_refill_ms: now,
});
let elapsed_ms = now.saturating_sub(bucket.last_refill_ms);
if elapsed_ms > 0 {
bucket.tokens = (bucket.tokens + (elapsed_ms as f64) * cap / 1000.0).min(cap);
bucket.last_refill_ms = now;
}
let need = n as f64;
if bucket.tokens >= need {
bucket.tokens -= need;
true
} else {
false
}
}
async fn group_commit(
&self,
ops: Vec<WriteOp>,
msgs: usize,
max_unflushed: usize,
) -> Result<(), MessagingError> {
use futures::future::{select, Either};
let (done_tx, mut done_rx) = futures::channel::oneshot::channel();
self.commit_queue.lock().unwrap().push(PublishJob {
ops,
msgs,
max_unflushed,
done: done_tx,
});
let gate = self.commit_gate.lock();
futures::pin_mut!(gate);
match select(gate, &mut done_rx).await {
Either::Right((res, _gate)) => {
return res
.map_err(|_| MessagingError::backend("group-commit dropped before durable"))?
}
Either::Left((_turn, _done)) => loop {
let batch: Vec<PublishJob> = {
let mut q = self.commit_queue.lock().unwrap();
if q.is_empty() {
break;
}
let mut n = 0;
let mut op_count = 0;
while n < q.len() {
if n > 0 && op_count + q[n].ops.len() > GROUP_COMMIT_MAX {
break;
}
op_count += q[n].ops.len();
n += 1;
}
q.drain(..n).collect()
};
let mut all_ops = Vec::new();
let mut dones = Vec::with_capacity(batch.len());
let mut msgs = 0usize;
let mut batch_max_unflushed = usize::MAX;
for mut job in batch {
all_ops.append(&mut job.ops);
msgs += job.msgs;
batch_max_unflushed = batch_max_unflushed.min(job.max_unflushed);
dones.push(job.done);
}
let outcome = self.commit_group(all_ops, msgs, batch_max_unflushed).await;
for done in dones {
let _ = done.send(outcome.clone());
}
},
}
done_rx
.await
.map_err(|_| MessagingError::backend("group-commit dropped before durable"))?
}
async fn commit_group(
&self,
ops: Vec<WriteOp>,
msgs: usize,
max_unflushed: usize,
) -> Result<(), MessagingError> {
use std::sync::atomic::Ordering;
let prior = self.unflushed.load(Ordering::Relaxed);
if prior + msgs <= max_unflushed {
self.kv
.write_batch_relaxed(ops)
.await
.map_err(MessagingError::backend)?;
self.unflushed.fetch_add(msgs, Ordering::Relaxed);
} else {
self.kv
.write_batch(ops)
.await
.map_err(MessagingError::backend)?;
self.unflushed.store(0, Ordering::Relaxed);
}
Ok(())
}
#[must_use]
pub fn with_inline_budget(mut self, bytes: usize) -> Self {
self.inline_budget_bytes = bytes;
self
}
fn signal_ready(&self) {
self.wake.notify();
}
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
}
async fn build_publish_ops(
&self,
topic: &str,
payload: &[u8],
signed_context: Option<&str>,
not_before_ms: u64,
expires_at_ms: u64,
priority: u8,
) -> Result<(String, Vec<WriteOp>), MessagingError> {
let id = format!(
"{:013}-{:016x}",
now_unix_ms(),
self.seq.fetch_add(1, Ordering::Relaxed)
);
let retain = self.topic_has_groups(topic).await;
let inline = !retain
&& payload.len() <= INLINE_MAX
&& self
.inline_inflight_bytes
.load(std::sync::atomic::Ordering::Relaxed)
.saturating_add(payload.len())
<= self.inline_budget_bytes;
if inline {
self.inline_inflight_bytes
.fetch_add(payload.len(), std::sync::atomic::Ordering::Relaxed);
}
if !inline {
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 mut ops: Vec<WriteOp> = Vec::with_capacity(3);
let mut record = Record::fresh(signed_context.map(str::to_owned));
record.lease_until_ms = not_before_ms;
record.expires_at_ms = expires_at_ms;
record.priority = priority;
if inline {
record.inline = Some(payload.to_vec());
}
ops.push(WriteOp::Put(
meta_key(topic, &id),
serde_json::to_vec(&record).map_err(MessagingError::backend)?,
));
ops.push(WriteOp::Put(ready_key(topic), Vec::new()));
if retain {
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)?;
ops.push(WriteOp::Put(glog_key(topic, &id), Vec::new()));
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() {
ops.push(WriteOp::Put(logmax_key(topic), id.clone().into_bytes()));
}
}
Ok((id, ops))
}
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_ctx(&self, topic: &str, id: &str) -> Option<String> {
let raw = self.kv.get(&meta_key(topic, id)).await.ok()??;
let record: Record = serde_json::from_slice(&raw).ok()?;
record.signed_context
}
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 collect_dead_letters(&self, topic: &str) -> Result<Vec<DeadLetter>, MessagingError> {
let mut out = Vec::new();
let wq_prefix = dead_prefix(topic);
for key in self
.kv
.list_prefix(&wq_prefix)
.await
.map_err(MessagingError::backend)?
{
if !is_direct_child(&key, &wq_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()))?;
out.push(DeadLetter {
id: key[wq_prefix.len()..].to_string(),
group: String::new(),
attempts: record.attempts,
last_error: record.last_error,
signed_context: record.signed_context,
payload: None,
});
}
let gprefix = gdead_topic_prefix(topic);
for key in self
.kv
.list_prefix(&gprefix)
.await
.map_err(MessagingError::backend)?
{
let Some((group, id)) = split_group_id(&key[gprefix.len()..]) else {
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()))?;
out.push(DeadLetter {
id: id.to_string(),
group: group.to_string(),
attempts: record.attempts,
last_error: record.last_error,
signed_context: record.signed_context,
payload: None,
});
}
out.sort_by(|a, b| a.id.cmp(&b.id));
Ok(out)
}
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,
retention_ms: u64,
) -> 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 gdead_prefix = gdead_topic_prefix(topic);
let mut dead_ids = std::collections::HashSet::new();
for key in self
.kv
.list_prefix(&gdead_prefix)
.await
.map_err(MessagingError::backend)?
{
if let Some((_, id)) = split_group_id(&key[gdead_prefix.len()..]) {
dead_ids.insert(id.to_string());
}
}
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 pinned = dead_ids.contains(id);
let needed = grouped_message_needed(&states, id);
let expired = id_millis(id) + retention_ms < now;
if !pinned && (!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> {
self.publish_ctx(topic, payload, None).await
}
async fn publish_ctx(
&self,
topic: &str,
payload: &[u8],
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
let max_unflushed = self.enforce_and_resolve_budget(topic, 1).await?;
let (id, ops) = self
.build_publish_ops(topic, payload, signed_context, 0, 0, 0)
.await?;
self.group_commit(ops, 1, max_unflushed).await?;
self.signal_ready();
self.hubs.broadcast(topic, &id, payload);
Ok(())
}
async fn publish_delayed_ctx(
&self,
topic: &str,
payload: &[u8],
delay: Duration,
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
let max_unflushed = self.enforce_and_resolve_budget(topic, 1).await?;
let not_before_ms = if delay.is_zero() {
0
} else {
now_unix_ms().saturating_add(delay.as_millis() as u64)
};
let (id, ops) = self
.build_publish_ops(topic, payload, signed_context, not_before_ms, 0, 0)
.await?;
self.group_commit(ops, 1, max_unflushed).await?;
self.signal_ready(); self.hubs.broadcast(topic, &id, payload);
Ok(())
}
async fn publish_with_ttl_ctx(
&self,
topic: &str,
payload: &[u8],
ttl: Duration,
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
let max_unflushed = self.enforce_and_resolve_budget(topic, 1).await?;
let expires_at_ms = if ttl.is_zero() {
0
} else {
now_unix_ms().saturating_add(ttl.as_millis() as u64)
};
let (id, ops) = self
.build_publish_ops(topic, payload, signed_context, 0, expires_at_ms, 0)
.await?;
self.group_commit(ops, 1, max_unflushed).await?;
self.signal_ready(); self.hubs.broadcast(topic, &id, payload);
Ok(())
}
async fn publish_with_priority_ctx(
&self,
topic: &str,
payload: &[u8],
priority: u8,
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
let max_unflushed = self.enforce_and_resolve_budget(topic, 1).await?;
let (id, ops) = self
.build_publish_ops(topic, payload, signed_context, 0, 0, priority)
.await?;
self.group_commit(ops, 1, max_unflushed).await?;
self.signal_ready(); self.hubs.broadcast(topic, &id, payload);
Ok(())
}
async fn publish_batch_ctx(
&self,
messages: &[(String, Vec<u8>)],
signed_context: Option<&str>,
) -> Result<(), MessagingError> {
if messages.is_empty() {
return Ok(());
}
let mut per_topic_count: HashMap<&str, usize> = HashMap::new();
for (topic, _) in messages {
*per_topic_count.entry(topic.as_str()).or_insert(0) += 1;
}
let mut batch_max_unflushed = usize::MAX;
for (topic, count) in &per_topic_count {
match self.resolve_policy(topic).await? {
Some(p) => {
self.enforce_publish_policy(topic, &p, *count).await?;
batch_max_unflushed =
batch_max_unflushed.min(p.max_unflushed.unwrap_or(self.max_unflushed));
}
None => batch_max_unflushed = batch_max_unflushed.min(self.max_unflushed),
}
}
let mut all_ops: Vec<WriteOp> = Vec::with_capacity(messages.len());
let mut broadcasts: Vec<(&str, String, &[u8])> = Vec::with_capacity(messages.len());
for (topic, payload) in messages {
let (id, ops) = self
.build_publish_ops(topic, payload, signed_context, 0, 0, 0)
.await?;
all_ops.extend(ops);
broadcasts.push((topic.as_str(), id, payload.as_slice()));
}
self.group_commit(all_ops, messages.len(), batch_max_unflushed)
.await?;
self.signal_ready(); for (topic, id, payload) in &broadcasts {
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> {
if self.is_paused(topic).await? {
return Ok(Vec::new());
}
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 lease_ms = lease.as_millis() as u64;
let claimable_now = |rec: &Record| -> bool {
rec.lease_until_ms <= now
&& !(rec.expires_at_ms != 0 && rec.expires_at_ms <= now)
&& rec.attempts < max_attempts
};
let claimable_before = records.iter().filter(|(_, r)| claimable_now(r)).count();
let actions = plan_claim(records, now, lease_ms, max_batch, max_attempts);
let mut ops: Vec<WriteOp> = Vec::new();
let mut claimed = Vec::new();
let mut leased = 0usize;
for action in actions {
match action {
ClaimAction::Lease { id, record } => {
let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
ops.push(WriteOp::Put(meta_key(topic, &id), json));
let inline = record.inline.is_some();
let payload = match record.inline {
Some(bytes) => bytes,
None => self.read_payload(topic, &id).await?,
};
claimed.push(ClaimedMessage {
id,
topic: topic.to_string(),
payload,
attempts: record.attempts,
group: String::new(),
signed_context: record.signed_context,
inline,
});
leased += 1;
}
ClaimAction::DeadLetter { id, record } => {
let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
ops.push(WriteOp::Put(dead_key(topic, &id), json));
ops.push(WriteOp::Delete(meta_key(topic, &id)));
}
}
}
if !ops.is_empty() {
self.kv
.write_batch(ops)
.await
.map_err(MessagingError::backend)?;
}
let residual_claimable = claimable_before.saturating_sub(leased);
let wq_drained = residual_claimable == 0;
if wq_drained && !self.topic_has_groups(topic).await {
let _commit = self.commit_gate.lock().await;
let prefix = meta_prefix(topic);
let has_any_record = self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?
.iter()
.any(|k| is_direct_child(k, &prefix));
if !has_any_record {
self.kv
.write_batch(vec![WriteOp::Delete(ready_key(topic))])
.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;
}
if self.is_paused(topic).await? {
return Ok(Vec::new());
}
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 signed_context = self.read_ctx(topic, id).await;
let record = Record {
version: crate::SCHEMA_VERSION,
attempts: *attempts,
lease_until_ms: 0,
signed_context,
inline: None,
last_error: None,
expires_at_ms: 0,
priority: 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) => {
let signed_context = self.read_ctx(topic, &id).await;
claimed.push(ClaimedMessage {
id,
topic: topic.to_string(),
payload,
attempts,
group: group.to_string(),
signed_context,
inline: false,
});
}
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)?;
if msg.inline {
let _ = self.inline_inflight_bytes.fetch_update(
std::sync::atomic::Ordering::Relaxed,
std::sync::atomic::Ordering::Relaxed,
|v| Some(v.saturating_sub(msg.payload.len())),
);
} else {
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 oldest_pending_ms(&self, topic: &str) -> Result<Option<u64>, MessagingError> {
let prefix = meta_prefix(topic);
let mut oldest: Option<u64> = None;
for key in self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?
{
if !is_direct_child(&key, &prefix) {
continue;
}
let ms = id_millis(&key[prefix.len()..]);
oldest = Some(oldest.map_or(ms, |o| o.min(ms)));
}
Ok(oldest.map(|ms| now_unix_ms().saturating_sub(ms)))
}
async fn in_flight_count(&self, topic: &str) -> Result<usize, MessagingError> {
let now = now_unix_ms();
let prefix = meta_prefix(topic);
let mut count = 0;
for key in self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?
{
if !is_direct_child(&key, &prefix) {
continue;
}
if let Some(raw) = self.kv.get(&key).await.map_err(MessagingError::backend)? {
if let Ok(rec) = serde_json::from_slice::<Record>(&raw) {
if rec.lease_until_ms > now {
count += 1;
}
}
}
}
let gprefix = gstate_prefix(topic);
for key in self
.kv
.list_prefix(&gprefix)
.await
.map_err(MessagingError::backend)?
{
if !is_direct_child(&key, &gprefix) {
continue;
}
let group = &key[gprefix.len()..];
if let Some(state) = self.get_group_state(topic, group).await? {
count += state
.in_flight
.iter()
.filter(|f| f.lease_until_ms > now)
.count();
}
}
Ok(count)
}
async fn group_lag(&self, topic: &str, group: &str) -> Result<usize, MessagingError> {
let Some(state) = self.get_group_state(topic, group).await? else {
return Ok(0);
};
let prefix = glog_prefix(topic);
let mut lag = 0;
for key in self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?
{
if !is_direct_child(&key, &prefix) {
continue;
}
if key[prefix.len()..] > *state.hwm {
lag += 1;
}
}
Ok(lag)
}
async fn dead_letter_count(&self, topic: &str) -> Result<usize, MessagingError> {
let wq = self.count_direct(&dead_prefix(topic)).await?;
let gprefix = gdead_topic_prefix(topic);
let grouped = self
.kv
.list_prefix(&gprefix)
.await
.map_err(MessagingError::backend)?
.into_iter()
.filter(|k| split_group_id(&k[gprefix.len()..]).is_some())
.count();
Ok(wq + grouped)
}
async fn nack(&self, msg: &ClaimedMessage) -> Result<(), MessagingError> {
self.nack_after(msg, 0).await
}
async fn nack_after(&self, msg: &ClaimedMessage, delay_ms: u64) -> Result<(), MessagingError> {
let until = if delay_ms == 0 {
0
} else {
now_unix_ms().saturating_add(delay_ms)
};
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 = until;
changed = true;
break;
}
}
if changed {
let json = serde_json::to_vec(&state).map_err(MessagingError::backend)?;
self.kv
.write_batch(vec![
WriteOp::Put(gstate_key(&msg.topic, &msg.group), json),
WriteOp::Put(ready_key(&msg.topic), Vec::new()),
])
.await
.map_err(MessagingError::backend)?;
self.signal_ready(); }
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 = until;
let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
self.kv
.write_batch(vec![
WriteOp::Put(key, json),
WriteOp::Put(ready_key(&msg.topic), Vec::new()),
])
.await
.map_err(MessagingError::backend)?;
self.signal_ready(); Ok(())
}
async fn purge_dead_letters(&self, topic: &str) -> Result<usize, MessagingError> {
let mut purged = 0;
let prefix = dead_prefix(topic);
for key in self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?
{
if !is_direct_child(&key, &prefix) {
continue; }
let id = &key[prefix.len()..];
let inline_len = self
.kv
.get(&key)
.await
.map_err(MessagingError::backend)?
.and_then(|raw| serde_json::from_slice::<Record>(&raw).ok())
.and_then(|r| r.inline.map(|p| p.len()));
match inline_len {
Some(len) => {
let _ = self.inline_inflight_bytes.fetch_update(
std::sync::atomic::Ordering::Relaxed,
std::sync::atomic::Ordering::Relaxed,
|v| Some(v.saturating_sub(len)),
);
}
None => {
self.storage
.delete(&payload_key(topic, id))
.await
.map_err(MessagingError::backend)?;
}
}
self.kv
.delete(&key)
.await
.map_err(MessagingError::backend)?;
purged += 1;
}
let gprefix = gdead_topic_prefix(topic);
for key in self
.kv
.list_prefix(&gprefix)
.await
.map_err(MessagingError::backend)?
{
if split_group_id(&key[gprefix.len()..]).is_none() {
continue;
}
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 mut redriven = 0;
let prefix = dead_prefix(topic);
for key in self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?
{
if !is_direct_child(&key, &prefix) {
continue;
}
let id = &key[prefix.len()..];
let mut record = self
.kv
.get(&key)
.await
.map_err(MessagingError::backend)?
.and_then(|raw| serde_json::from_slice::<Record>(&raw).ok())
.unwrap_or_else(|| Record::fresh(None));
record.attempts = 0;
record.lease_until_ms = 0;
record.last_error = None; record.expires_at_ms = 0; let json = serde_json::to_vec(&record).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;
}
let gprefix = gdead_topic_prefix(topic);
let gkeys = self
.kv
.list_prefix(&gprefix)
.await
.map_err(MessagingError::backend)?;
if !gkeys.is_empty() {
let _guard = self.claim_lock.lock().await;
for key in gkeys {
let Some((group, id)) = split_group_id(&key[gprefix.len()..]) else {
continue;
};
let mut state = match self.get_group_state(topic, group).await? {
Some(state) => state,
None => {
self.mark_grouped(topic);
GroupState::new(self.read_logmax(topic).await?)
}
};
if !state.in_flight.iter().any(|f| f.id == id) {
state.in_flight.push(InFlight {
id: id.to_string(),
attempts: 0,
lease_until_ms: 0,
});
}
self.put_group_state(topic, group, &state).await?;
self.kv
.delete(&key)
.await
.map_err(MessagingError::backend)?;
redriven += 1;
}
}
Ok(redriven)
}
async fn set_last_error(
&self,
msg: &ClaimedMessage,
reason: &str,
) -> Result<(), MessagingError> {
if !msg.group.is_empty() {
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.last_error = Some(sanitize_reason(reason));
let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
self.kv
.put(&key, json)
.await
.map_err(MessagingError::backend)?;
Ok(())
}
async fn list_dead_letters(
&self,
topic: &str,
filter: &DeadLetterFilter,
) -> Result<Vec<DeadLetter>, MessagingError> {
let now = now_unix_ms();
let mut matched: Vec<DeadLetter> = self
.collect_dead_letters(topic)
.await?
.into_iter()
.filter(|dl| filter.matches(dl, now))
.collect();
if let Some(limit) = filter.limit {
matched.truncate(limit); }
Ok(matched)
}
async fn show_dead_letter(
&self,
topic: &str,
group: &str,
id: &str,
) -> Result<Option<DeadLetter>, MessagingError> {
let key = if group.is_empty() {
dead_key(topic, id)
} else {
gdead_key(topic, group, id)
};
let Some(raw) = self.kv.get(&key).await.map_err(MessagingError::backend)? else {
return Ok(None);
};
let record: Record =
serde_json::from_slice(&raw).map_err(|e| MessagingError::Decode(e.to_string()))?;
let payload = if let Some(inline) = record.inline.clone() {
inline
} else if group.is_empty() {
self.read_payload(topic, id).await.unwrap_or_default()
} else {
self.read_gpayload(topic, id).await.unwrap_or_default()
};
Ok(Some(DeadLetter {
id: id.to_string(),
group: group.to_string(),
attempts: record.attempts,
last_error: record.last_error,
signed_context: record.signed_context,
payload: Some(payload),
}))
}
async fn redrive_dead_letters_filtered(
&self,
topic: &str,
filter: &DeadLetterFilter,
) -> Result<usize, MessagingError> {
let now = now_unix_ms();
let matched: Vec<DeadLetter> = {
let mut m: Vec<DeadLetter> = self
.collect_dead_letters(topic)
.await?
.into_iter()
.filter(|dl| filter.matches(dl, now))
.collect();
if let Some(limit) = filter.limit {
m.truncate(limit);
}
m
};
let mut redriven = 0;
for dl in matched.iter().filter(|dl| dl.group.is_empty()) {
let dead = dead_key(topic, &dl.id);
let Some(raw) = self.kv.get(&dead).await.map_err(MessagingError::backend)? else {
continue;
};
let mut record: Record =
serde_json::from_slice(&raw).map_err(|e| MessagingError::Decode(e.to_string()))?;
record.attempts = 0;
record.lease_until_ms = 0;
record.last_error = None; record.expires_at_ms = 0; let json = serde_json::to_vec(&record).map_err(MessagingError::backend)?;
self.kv
.put(&meta_key(topic, &dl.id), json)
.await
.map_err(MessagingError::backend)?;
self.kv
.delete(&dead)
.await
.map_err(MessagingError::backend)?;
redriven += 1;
}
let grouped: Vec<&DeadLetter> = matched.iter().filter(|dl| !dl.group.is_empty()).collect();
if !grouped.is_empty() {
let _guard = self.claim_lock.lock().await;
for dl in grouped {
let dead = gdead_key(topic, &dl.group, &dl.id);
if self
.kv
.get(&dead)
.await
.map_err(MessagingError::backend)?
.is_none()
{
continue;
}
let mut state = match self.get_group_state(topic, &dl.group).await? {
Some(state) => state,
None => {
self.mark_grouped(topic);
GroupState::new(self.read_logmax(topic).await?)
}
};
if !state.in_flight.iter().any(|f| f.id == dl.id) {
state.in_flight.push(InFlight {
id: dl.id.clone(),
attempts: 0,
lease_until_ms: 0,
});
}
self.put_group_state(topic, &dl.group, &state).await?;
self.kv
.delete(&dead)
.await
.map_err(MessagingError::backend)?;
redriven += 1;
}
}
Ok(redriven)
}
async fn discard_dead_letters(
&self,
topic: &str,
filter: &DeadLetterFilter,
) -> Result<usize, MessagingError> {
let now = now_unix_ms();
let matched: Vec<DeadLetter> = {
let mut m: Vec<DeadLetter> = self
.collect_dead_letters(topic)
.await?
.into_iter()
.filter(|dl| filter.matches(dl, now))
.collect();
if let Some(limit) = filter.limit {
m.truncate(limit);
}
m
};
let mut discarded = 0;
for dl in &matched {
if dl.group.is_empty() {
let dead = dead_key(topic, &dl.id);
let inline_len = self
.kv
.get(&dead)
.await
.map_err(MessagingError::backend)?
.and_then(|raw| serde_json::from_slice::<Record>(&raw).ok())
.and_then(|r| r.inline.map(|p| p.len()));
match inline_len {
Some(len) => {
let _ = self.inline_inflight_bytes.fetch_update(
std::sync::atomic::Ordering::Relaxed,
std::sync::atomic::Ordering::Relaxed,
|v| Some(v.saturating_sub(len)),
);
}
None => {
self.storage
.delete(&payload_key(topic, &dl.id))
.await
.map_err(MessagingError::backend)?;
}
}
self.kv
.delete(&dead)
.await
.map_err(MessagingError::backend)?;
} else {
self.kv
.delete(&gdead_key(topic, &dl.group, &dl.id))
.await
.map_err(MessagingError::backend)?;
}
discarded += 1;
}
Ok(discarded)
}
async fn peek(&self, topic: &str, limit: usize) -> Result<Vec<PeekedMessage>, MessagingError> {
let now = now_unix_ms();
let prefix = meta_prefix(topic);
let mut keys: Vec<String> = self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?
.into_iter()
.filter(|k| is_direct_child(k, &prefix))
.collect();
keys.sort(); let mut out = Vec::new();
for key in keys.into_iter().take(limit) {
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()))?;
let id = key[prefix.len()..].to_string();
let payload = match &record.inline {
Some(bytes) => bytes.clone(),
None => self.read_payload(topic, &id).await.unwrap_or_default(),
};
out.push(PeekedMessage {
id,
attempts: record.attempts,
leased: record.lease_until_ms > now,
signed_context: record.signed_context,
payload,
});
}
Ok(out)
}
async fn replay(
&self,
topic: &str,
after: Option<&str>,
limit: usize,
) -> Result<Vec<PeekedMessage>, MessagingError> {
let prefix = glog_prefix(topic);
let mut ids: Vec<String> = self
.kv
.list_prefix(&prefix)
.await
.map_err(MessagingError::backend)?
.into_iter()
.filter(|k| is_direct_child(k, &prefix))
.map(|k| k[prefix.len()..].to_string())
.collect();
ids.sort(); let mut out = Vec::new();
for id in ids {
if let Some(after) = after {
if id.as_str() <= after {
continue;
}
}
if out.len() >= limit {
break;
}
let payload = self.read_gpayload(topic, &id).await.unwrap_or_default();
let signed_context = self.read_ctx(topic, &id).await;
out.push(PeekedMessage {
id,
attempts: 0, leased: false, signed_context,
payload,
});
}
Ok(out)
}
async fn list_groups(&self, topic: &str) -> Result<Vec<GroupInfo>, MessagingError> {
let gprefix = gstate_prefix(topic);
let group_keys = self
.kv
.list_prefix(&gprefix)
.await
.map_err(MessagingError::backend)?;
let log_prefix = glog_prefix(topic);
let log_ids: Vec<String> = self
.kv
.list_prefix(&log_prefix)
.await
.map_err(MessagingError::backend)?
.into_iter()
.filter(|k| is_direct_child(k, &log_prefix))
.map(|k| k[log_prefix.len()..].to_string())
.collect();
let mut out = Vec::new();
for key in group_keys {
if !is_direct_child(&key, &gprefix) {
continue;
}
let group = key[gprefix.len()..].to_string();
let Some(state) = self.get_group_state(topic, &group).await? else {
continue;
};
let lag = log_ids
.iter()
.filter(|id| id.as_str() > state.hwm.as_str())
.count();
out.push(GroupInfo {
group,
hwm: state.hwm,
in_flight: state.in_flight.len(),
lag,
});
}
out.sort_by(|a, b| a.group.cmp(&b.group));
Ok(out)
}
async fn reset_group(
&self,
topic: &str,
group: &str,
start: StartPosition,
) -> Result<(), MessagingError> {
let _guard = self.claim_lock.lock().await;
if self.get_group_state(topic, group).await?.is_none() {
return Err(MessagingError::Backend(format!(
"no such consumer group {group:?} on topic {topic:?}"
)));
}
let hwm = match start {
StartPosition::Earliest => String::new(),
StartPosition::Latest => self.read_logmax(topic).await?,
};
self.put_group_state(topic, group, &GroupState::new(hwm))
.await?;
Ok(())
}
async fn delete_group(&self, topic: &str, group: &str) -> Result<(), MessagingError> {
let _guard = self.claim_lock.lock().await;
let dprefix = format!("mqgd/{topic}/{group}/");
for key in self
.kv
.list_prefix(&dprefix)
.await
.map_err(MessagingError::backend)?
{
self.kv
.delete(&key)
.await
.map_err(MessagingError::backend)?;
}
self.kv
.delete(&gstate_key(topic, group))
.await
.map_err(MessagingError::backend)?;
Ok(())
}
async fn set_paused(&self, topic: &str, paused: bool) -> Result<(), MessagingError> {
let key = pause_key(topic);
if paused {
self.kv
.put(&key, Vec::new())
.await
.map_err(MessagingError::backend)?;
} else {
self.kv
.delete(&key)
.await
.map_err(MessagingError::backend)?;
}
Ok(())
}
async fn is_paused(&self, topic: &str) -> Result<bool, MessagingError> {
Ok(self
.kv
.get(&pause_key(topic))
.await
.map_err(MessagingError::backend)?
.is_some())
}
async fn set_topic_policy(
&self,
topic: &str,
policy: TopicPolicy,
) -> Result<(), MessagingError> {
let json = serde_json::to_vec(&policy).map_err(MessagingError::backend)?;
self.kv
.put(&mqpolicy_key(topic), json)
.await
.map_err(MessagingError::backend)?;
self.policy_cache.lock().unwrap().remove(topic);
Ok(())
}
async fn topic_policy(&self, topic: &str) -> Result<Option<TopicPolicy>, MessagingError> {
self.resolve_policy(topic).await
}
async fn retention_sweep(
&self,
topic: &str,
retention_ms: u64,
) -> Result<usize, MessagingError> {
self.gc_grouped(topic, retention_ms).await
}
fn subscribe(
&self,
topic: &str,
after: Option<&str>,
) -> futures::stream::BoxStream<'static, StreamEvent> {
self.hubs.subscribe(topic, after)
}
fn supports_ready_set(&self) -> bool {
self.kv.atomic_write_batch()
}
async fn ready_topics(&self) -> Result<Vec<String>, MessagingError> {
let keys = self
.kv
.list_prefix(READY_PREFIX)
.await
.map_err(MessagingError::backend)?;
Ok(keys
.iter()
.filter_map(|k| topic_of_ready_key(k).map(str::to_owned))
.collect())
}
async fn rebuild_ready_set(&self) -> Result<usize, MessagingError> {
let now = now_unix_ms();
let mut with_work: std::collections::HashSet<String> = std::collections::HashSet::new();
for key in self
.kv
.list_prefix("mq/")
.await
.map_err(MessagingError::backend)?
{
if let Some(rest) = key.strip_prefix("mq/") {
if let Some(slash) = rest.rfind('/') {
with_work.insert(rest[..slash].to_string());
}
}
}
let gstate_all = self
.kv
.list_prefix("mqgstate/")
.await
.map_err(MessagingError::backend)?;
for key in &gstate_all {
let Some(rest) = key.strip_prefix("mqgstate/") else {
continue;
};
let Some(slash) = rest.rfind('/') else {
continue;
};
let topic = &rest[..slash];
let group = &rest[slash + 1..];
if topic.is_empty() || group.is_empty() {
continue;
}
if with_work.contains(topic) {
continue; }
let Some(state) = self.get_group_state(topic, group).await? else {
continue;
};
let has_inflight = !state.in_flight.is_empty();
let has_backlog = state.hwm.as_str() < self.read_logmax(topic).await?.as_str();
if has_inflight || has_backlog {
with_work.insert(topic.to_string());
}
}
let existing: std::collections::HashSet<String> = self
.kv
.list_prefix(READY_PREFIX)
.await
.map_err(MessagingError::backend)?
.iter()
.filter_map(|k| topic_of_ready_key(k).map(str::to_owned))
.collect();
let mut added_any = false;
for topic in with_work.difference(&existing) {
self.kv
.put(&ready_key(topic), Vec::new())
.await
.map_err(MessagingError::backend)?;
added_any = true;
}
self.last_rebuild_ms.store(now, Ordering::Relaxed);
if added_any {
self.wake.notify();
}
Ok(with_work.len())
}
async fn due_topics(&self) -> Result<Vec<(String, u64)>, MessagingError> {
let now = now_unix_ms();
let mut due: HashMap<String, u64> = HashMap::new();
let mut note = |topic: &str, until: u64| {
if until > now {
due.entry(topic.to_string())
.and_modify(|e| *e = (*e).min(until))
.or_insert(until);
}
};
for key in self
.kv
.list_prefix("mq/")
.await
.map_err(MessagingError::backend)?
{
let Some(rest) = key.strip_prefix("mq/") else {
continue;
};
let Some(slash) = rest.rfind('/') else {
continue;
};
let topic = rest[..slash].to_string();
if let Some(raw) = self.kv.get(&key).await.map_err(MessagingError::backend)? {
if let Ok(rec) = serde_json::from_slice::<Record>(&raw) {
note(&topic, rec.lease_until_ms);
}
}
}
for key in self
.kv
.list_prefix("mqgstate/")
.await
.map_err(MessagingError::backend)?
{
let Some(rest) = key.strip_prefix("mqgstate/") else {
continue;
};
let Some(slash) = rest.rfind('/') else {
continue;
};
let topic = rest[..slash].to_string();
let group = &rest[slash + 1..];
if let Some(state) = self.get_group_state(&topic, group).await? {
for f in &state.in_flight {
note(&topic, f.lease_until_ms);
}
}
}
Ok(due.into_iter().collect())
}
async fn delivery_stats(&self, this_node: &str) -> Result<DeliveryStats, MessagingError> {
let ready_set_size = self.ready_topics().await?.len();
let due_heap_depth = self.due_topics().await?.len();
let last = self.last_rebuild_ms.load(Ordering::Relaxed);
let last_rebuild_age_ms = (last != 0).then(|| now_unix_ms().saturating_sub(last));
Ok(DeliveryStats {
ready_set_size,
due_heap_depth,
last_rebuild_age_ms,
this_node: this_node.to_string(),
})
}
fn wake_handle(&self) -> Option<Wake> {
Some(self.wake.clone())
}
}
#[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()))
}
struct FailingKv;
#[async_trait]
impl KvStore for FailingKv {
async fn get(&self, _: &str) -> Result<Option<Vec<u8>>, crate::kv::KvError> {
Ok(None)
}
async fn put(&self, _: &str, _: Vec<u8>) -> Result<(), crate::kv::KvError> {
Err(crate::kv::KvError::backend("commit failed"))
}
async fn delete(&self, _: &str) -> Result<(), crate::kv::KvError> {
Ok(())
}
async fn list_prefix(&self, _: &str) -> Result<Vec<String>, crate::kv::KvError> {
Ok(Vec::new())
}
async fn write_batch(&self, _: Vec<crate::kv::WriteOp>) -> Result<(), crate::kv::KvError> {
Err(crate::kv::KvError::backend("commit failed"))
}
}
struct CountingKv {
inner: MemoryKv,
batches: std::sync::atomic::AtomicUsize,
relaxed: std::sync::atomic::AtomicUsize,
}
#[async_trait]
impl KvStore for CountingKv {
async fn get(&self, k: &str) -> Result<Option<Vec<u8>>, crate::kv::KvError> {
self.inner.get(k).await
}
async fn put(&self, k: &str, v: Vec<u8>) -> Result<(), crate::kv::KvError> {
self.inner.put(k, v).await
}
async fn delete(&self, k: &str) -> Result<(), crate::kv::KvError> {
self.inner.delete(k).await
}
async fn list_prefix(&self, p: &str) -> Result<Vec<String>, crate::kv::KvError> {
self.inner.list_prefix(p).await
}
fn atomic_write_batch(&self) -> bool {
self.inner.atomic_write_batch()
}
async fn write_batch(
&self,
ops: Vec<crate::kv::WriteOp>,
) -> Result<(), crate::kv::KvError> {
self.batches
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.inner.write_batch(ops).await
}
async fn write_batch_relaxed(
&self,
ops: Vec<crate::kv::WriteOp>,
) -> Result<(), crate::kv::KvError> {
self.relaxed
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.inner.write_batch(ops).await
}
}
struct GateKv {
inner: MemoryKv,
calls: std::sync::atomic::AtomicUsize,
entered_tx: std::sync::Mutex<Option<futures::channel::oneshot::Sender<()>>>,
release_rx: std::sync::Mutex<Option<futures::channel::oneshot::Receiver<()>>>,
}
impl GateKv {
fn new() -> (
Arc<Self>,
futures::channel::oneshot::Receiver<()>,
futures::channel::oneshot::Sender<()>,
) {
let (entered_tx, entered_rx) = futures::channel::oneshot::channel();
let (release_tx, release_rx) = futures::channel::oneshot::channel();
let kv = Arc::new(Self {
inner: MemoryKv::new(),
calls: std::sync::atomic::AtomicUsize::new(0),
entered_tx: std::sync::Mutex::new(Some(entered_tx)),
release_rx: std::sync::Mutex::new(Some(release_rx)),
});
(kv, entered_rx, release_tx)
}
}
#[async_trait]
impl KvStore for GateKv {
async fn get(&self, k: &str) -> Result<Option<Vec<u8>>, crate::kv::KvError> {
self.inner.get(k).await
}
async fn put(&self, k: &str, v: Vec<u8>) -> Result<(), crate::kv::KvError> {
self.inner.put(k, v).await
}
async fn delete(&self, k: &str) -> Result<(), crate::kv::KvError> {
self.inner.delete(k).await
}
async fn list_prefix(&self, p: &str) -> Result<Vec<String>, crate::kv::KvError> {
self.inner.list_prefix(p).await
}
async fn write_batch(
&self,
ops: Vec<crate::kv::WriteOp>,
) -> Result<(), crate::kv::KvError> {
if self
.calls
.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
== 0
{
if let Some(tx) = self.entered_tx.lock().unwrap().take() {
let _ = tx.send(());
}
let rx = self.release_rx.lock().unwrap().take();
if let Some(rx) = rx {
let _ = rx.await;
}
}
self.inner.write_batch(ops).await
}
}
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, GROUP_RETENTION_MS).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, GROUP_RETENTION_MS).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 delayed_publish_defers_until_its_not_before() {
let mq = mq();
mq.publish_delayed_ctx("t", b"later", Duration::from_secs(3600), None)
.await
.unwrap();
mq.publish_delayed_ctx("t", b"now", Duration::ZERO, None)
.await
.unwrap();
assert_eq!(mq.backlog("t").await.unwrap(), 2, "both durably enqueued");
let ready = mq.claim("t", LEASE, 10, 5).await.unwrap();
assert_eq!(
payloads(&ready),
vec![b"now".to_vec()],
"delayed one deferred"
);
assert_eq!(ready[0].attempts, 1);
}
#[test]
fn plan_claim_dead_letters_an_expired_message() {
let expired = Record {
expires_at_ms: 100, ..Record::fresh(None)
};
let live = Record::fresh(None); let actions = plan_claim(
vec![
("0000000000100-a".to_string(), expired),
("0000000000200-b".to_string(), live),
],
1_000, 30_000,
10,
5,
);
assert!(
matches!(&actions[0], ClaimAction::DeadLetter { record, .. }
if record.last_error.as_deref() == Some("ttl-expired")),
"the expired message dead-letters with the ttl-expired reason"
);
assert!(
matches!(&actions[1], ClaimAction::Lease { .. }),
"the live message leases normally"
);
}
#[test]
fn plan_claim_orders_by_priority_then_fifo() {
let hi = Record {
priority: 5,
..Record::fresh(None)
};
let actions = plan_claim(
vec![
("0000000000001-a".to_string(), Record::fresh(None)), ("0000000000002-b".to_string(), hi), ("0000000000003-c".to_string(), Record::fresh(None)), ],
1_000,
30_000,
10,
5,
);
let leased: Vec<&str> = actions
.iter()
.filter_map(|a| match a {
ClaimAction::Lease { id, .. } => Some(id.as_str()),
_ => None,
})
.collect();
assert_eq!(
leased,
vec!["0000000000002-b", "0000000000001-a", "0000000000003-c"],
"high-priority b leases first, then a & c FIFO within the default priority"
);
}
#[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 grouped_dead_letters_are_visible_redrivable_and_purgeable() {
let mq = mq();
let t = "bus/sync";
assert!(mq
.claim_grouped(t, "worker", StartPosition::Latest, Duration::ZERO, 10, 2)
.await
.unwrap()
.is_empty());
mq.publish(t, b"poison").await.unwrap();
for expected in 1..=2 {
let m = mq
.claim_grouped(t, "worker", StartPosition::Latest, Duration::ZERO, 10, 2)
.await
.unwrap();
assert_eq!(m.len(), 1, "grouped attempt {expected}");
assert_eq!(m[0].attempts, expected);
}
assert!(mq
.claim_grouped(t, "worker", StartPosition::Latest, Duration::ZERO, 10, 2)
.await
.unwrap()
.is_empty());
assert_eq!(mq.dead_letter_count(t).await.unwrap(), 1);
assert_eq!(mq.redrive_dead_letters(t).await.unwrap(), 1);
assert_eq!(mq.dead_letter_count(t).await.unwrap(), 0);
let again = mq
.claim_grouped(t, "worker", StartPosition::Latest, Duration::ZERO, 10, 2)
.await
.unwrap();
assert_eq!(payloads(&again), vec![b"poison".to_vec()]);
assert_eq!(again[0].attempts, 1, "redrive reset the attempt count");
assert_eq!(
mq.claim_grouped(t, "worker", StartPosition::Latest, Duration::ZERO, 10, 2)
.await
.unwrap()
.len(),
1
);
assert!(mq
.claim_grouped(t, "worker", StartPosition::Latest, Duration::ZERO, 10, 2)
.await
.unwrap()
.is_empty());
assert_eq!(mq.dead_letter_count(t).await.unwrap(), 1);
assert_eq!(mq.purge_dead_letters(t).await.unwrap(), 1);
assert_eq!(mq.dead_letter_count(t).await.unwrap(), 0);
}
#[tokio::test]
async fn stats_expose_in_flight_and_oldest_pending() {
let mq = mq();
assert_eq!(mq.oldest_pending_ms("t").await.unwrap(), None);
assert_eq!(mq.in_flight_count("t").await.unwrap(), 0);
mq.publish("t", b"a").await.unwrap();
mq.publish("t", b"b").await.unwrap();
assert!(mq.oldest_pending_ms("t").await.unwrap().is_some());
assert_eq!(mq.in_flight_count("t").await.unwrap(), 0);
let claimed = mq.claim("t", Duration::from_secs(60), 10, 5).await.unwrap();
assert_eq!(claimed.len(), 2);
assert_eq!(mq.in_flight_count("t").await.unwrap(), 2);
assert_eq!(mq.backlog("t").await.unwrap(), 2);
mq.ack(&claimed[0]).await.unwrap();
assert_eq!(mq.in_flight_count("t").await.unwrap(), 1);
assert_eq!(mq.backlog("t").await.unwrap(), 1);
let g = "bus/lag";
assert!(mq
.claim_grouped(g, "w", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
mq.publish(g, b"x").await.unwrap();
mq.publish(g, b"y").await.unwrap();
assert_eq!(
mq.group_lag(g, "w").await.unwrap(),
2,
"two retained, none leased yet"
);
let got = mq
.claim_grouped(g, "w", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap();
assert_eq!(got.len(), 2);
assert_eq!(
mq.group_lag(g, "w").await.unwrap(),
0,
"caught up to the high-water"
);
}
#[tokio::test]
async fn small_work_queue_payload_is_inlined_large_takes_object_store() {
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"small").await.unwrap();
let m = mq.claim("t", Duration::from_secs(60), 10, 5).await.unwrap();
assert_eq!(m.len(), 1);
assert_eq!(m[0].payload, b"small");
assert!(m[0].inline, "small payload inlined into the record");
assert!(
storage.head(&payload_key("t", &m[0].id)).await.is_err(),
"inlined ⇒ no object-store object written"
);
mq.ack(&m[0]).await.unwrap();
assert_eq!(mq.backlog("t").await.unwrap(), 0);
let big = vec![7u8; INLINE_MAX + 1];
mq.publish("t", &big).await.unwrap();
let m = mq.claim("t", Duration::from_secs(60), 10, 5).await.unwrap();
assert_eq!(m[0].payload, big);
assert!(!m[0].inline, "large payload not inlined");
assert!(
storage.head(&payload_key("t", &m[0].id)).await.is_ok(),
"large payload lives in object storage"
);
}
#[tokio::test]
async fn inline_budget_falls_back_to_object_store_when_exhausted() {
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).with_inline_budget(10);
mq.publish("t", b"aaaaa").await.unwrap(); mq.publish("t", b"bbbbb").await.unwrap(); mq.publish("t", b"ccccc").await.unwrap(); let got = mq.claim("t", Duration::from_secs(60), 10, 5).await.unwrap();
assert_eq!(got.len(), 3);
let inline_count = got.iter().filter(|m| m.inline).count();
assert_eq!(
inline_count, 2,
"budget admitted exactly two inline messages"
);
let obj = got.iter().find(|m| !m.inline).unwrap();
assert!(
storage.head(&payload_key("t", &obj.id)).await.is_ok(),
"the over-budget message fell back to the object store"
);
let inline_msg = got.iter().find(|m| m.inline).unwrap().clone();
mq.ack(&inline_msg).await.unwrap();
mq.publish("t", b"ddddd").await.unwrap();
let more = mq.claim("t", Duration::from_secs(60), 10, 5).await.unwrap();
assert!(
more.iter().any(|m| m.payload == b"ddddd" && m.inline),
"after ack freed budget, the next small publish inlines again"
);
}
#[tokio::test]
async fn group_commit_coalesces_concurrent_publishes() {
let mq = Arc::new(mq());
let mut handles = Vec::new();
for i in 0..64u32 {
let mq = mq.clone();
handles.push(tokio::spawn(async move {
mq.publish("t", format!("m{i}").as_bytes()).await
}));
}
for h in handles {
h.await
.unwrap()
.expect("each publish returns Ok after its durable group commit");
}
let mut seen = 0;
loop {
let batch = mq
.claim("t", Duration::from_secs(60), 100, 5)
.await
.unwrap();
if batch.is_empty() {
break;
}
seen += batch.len();
}
assert_eq!(
seen, 64,
"all concurrent publishes were durably committed and claimable"
);
}
#[tokio::test]
async fn group_commit_fails_all_members_when_the_commit_fails() {
let mq = Arc::new(LogMessaging::new(
Arc::new(MemStorage::default()),
Arc::new(FailingKv),
));
assert!(
mq.publish("t", b"x").await.is_err(),
"a failed group commit fails the publish"
);
let mut handles = Vec::new();
for i in 0..16u32 {
let mq = mq.clone();
handles.push(tokio::spawn(async move {
mq.publish("t", format!("m{i}").as_bytes()).await
}));
}
for h in handles {
assert!(
h.await.unwrap().is_err(),
"every member of a failed group commit fails (fail-all)"
);
}
}
#[tokio::test]
async fn publish_batch_commits_all_messages_in_one_write_batch() {
use std::sync::atomic::Ordering as O;
let kv = Arc::new(CountingKv {
inner: MemoryKv::new(),
batches: std::sync::atomic::AtomicUsize::new(0),
relaxed: std::sync::atomic::AtomicUsize::new(0),
});
let mq = LogMessaging::new(Arc::new(MemStorage::default()), kv.clone());
let msgs: Vec<(String, Vec<u8>)> = (0..5)
.map(|i| (format!("t{}", i % 2), format!("m{i}").into_bytes()))
.collect();
mq.publish_batch_ctx(&msgs, Some("ctx-1")).await.unwrap();
assert_eq!(
kv.batches.load(O::Relaxed),
1,
"the whole batch was durably committed in ONE write_batch"
);
let t0 = mq.claim("t0", LEASE, 10, 5).await.unwrap();
let t1 = mq.claim("t1", LEASE, 10, 5).await.unwrap();
assert_eq!(
t0.len() + t1.len(),
5,
"all batch messages durably enqueued and claimable"
);
assert!(
t0.iter()
.chain(&t1)
.all(|m| m.signed_context.as_deref() == Some("ctx-1")),
"every message carries the one shared producer context"
);
}
#[tokio::test]
async fn publish_batch_fails_whole_when_the_commit_fails() {
let mq = LogMessaging::new(Arc::new(MemStorage::default()), Arc::new(FailingKv));
let msgs: Vec<(String, Vec<u8>)> =
(0..8).map(|i| ("t".to_string(), vec![i as u8])).collect();
assert!(
mq.publish_batch_ctx(&msgs, None).await.is_err(),
"a failed commit fails the whole batch"
);
}
#[tokio::test]
async fn publish_batch_preserves_publish_order() {
let mq = mq();
let msgs: Vec<(String, Vec<u8>)> = (0..20)
.map(|i| ("t".to_string(), format!("m{i:02}").into_bytes()))
.collect();
mq.publish_batch_ctx(&msgs, None).await.unwrap();
let got = mq.claim("t", LEASE, 100, 5).await.unwrap();
assert_eq!(
payloads(&got),
msgs.iter().map(|(_, p)| p.clone()).collect::<Vec<_>>(),
"the batch is claimable in publish order"
);
}
#[tokio::test]
async fn oversized_single_batch_commits_whole_in_one_write_batch() {
use std::sync::atomic::Ordering as O;
let kv = Arc::new(CountingKv {
inner: MemoryKv::new(),
batches: std::sync::atomic::AtomicUsize::new(0),
relaxed: std::sync::atomic::AtomicUsize::new(0),
});
let mq = LogMessaging::new(Arc::new(MemStorage::default()), kv.clone());
let n = GROUP_COMMIT_MAX + 100;
let msgs: Vec<(String, Vec<u8>)> =
(0..n).map(|i| ("t".to_string(), vec![i as u8])).collect();
mq.publish_batch_ctx(&msgs, None).await.unwrap();
assert_eq!(
kv.batches.load(O::Relaxed),
1,
"the oversized single batch committed in exactly one write_batch (drain took it whole)"
);
let mut seen = 0;
loop {
let b = mq.claim("t", LEASE, 10_000, 5).await.unwrap();
if b.is_empty() {
break;
}
seen += b.len();
}
assert_eq!(
seen, n,
"every message in the oversized batch is durably enqueued"
);
}
#[tokio::test]
async fn cancelled_publisher_is_still_committed_by_a_later_gate_holder() {
let (kv, entered_rx, release_tx) = GateKv::new();
let mq = Arc::new(LogMessaging::new(Arc::new(MemStorage::default()), kv));
let q = {
let mq = mq.clone();
tokio::spawn(async move { mq.publish("t", b"Q").await })
};
entered_rx.await.unwrap();
{
let mut p = Box::pin(mq.publish("t", b"P"));
let polled = futures::poll!(p.as_mut());
assert!(
polled.is_pending(),
"P pushed its job and is now blocked on the gate Q holds"
);
}
release_tx.send(()).unwrap();
q.await.unwrap().unwrap();
mq.publish("t", b"R").await.unwrap();
let mut seen = Vec::new();
loop {
let b = mq.claim("t", LEASE, 100, 5).await.unwrap();
if b.is_empty() {
break;
}
seen.extend(b.into_iter().map(|m| m.payload));
}
assert!(seen.contains(&b"Q".to_vec()), "Q committed");
assert!(seen.contains(&b"R".to_vec()), "R committed");
assert!(
seen.contains(&b"P".to_vec()),
"the cancelled publisher's message was still durably committed (never silently lost)"
);
}
#[tokio::test]
async fn group_commit_coalesces_piled_up_waiters_into_one_batch() {
use std::sync::atomic::Ordering;
let (kv, entered_rx, release_tx) = GateKv::new();
let mq = Arc::new(LogMessaging::new(
Arc::new(MemStorage::default()),
kv.clone(),
));
let l = {
let mq = mq.clone();
tokio::spawn(async move { mq.publish("t", b"L").await })
};
entered_rx.await.unwrap();
const K: usize = 8;
let mut waiters = Vec::new();
for i in 0..K {
let mut w = Box::pin({
let mq = mq.clone();
async move { mq.publish("t", format!("w{i}").as_bytes()).await }
});
assert!(
futures::poll!(w.as_mut()).is_pending(),
"waiter {i} pushed its job and parked without taking the gate L holds"
);
waiters.push(w);
}
release_tx.send(()).unwrap();
l.await.unwrap().expect("leader publish ok");
for w in waiters {
w.await
.expect("a piled-up waiter completes (never stranded)");
}
assert_eq!(
kv.calls.load(Ordering::Relaxed),
2,
"the K piled-up waiters coalesced into ONE write_batch after the leader's own commit"
);
let mut seen = 0;
loop {
let b = mq.claim("t", LEASE, 100, 5).await.unwrap();
if b.is_empty() {
break;
}
seen += b.len();
}
assert_eq!(seen, K + 1, "leader + all K waiters were durably committed");
}
#[tokio::test]
async fn max_unflushed_zero_is_strong_durability() {
use std::sync::atomic::Ordering as O;
let kv = Arc::new(CountingKv {
inner: MemoryKv::new(),
batches: std::sync::atomic::AtomicUsize::new(0),
relaxed: std::sync::atomic::AtomicUsize::new(0),
});
let mq = LogMessaging::new(Arc::new(MemStorage::default()), kv.clone());
for i in 0..20u32 {
mq.publish("t", format!("m{i}").as_bytes()).await.unwrap();
}
assert_eq!(
kv.relaxed.load(O::Relaxed),
0,
"strong default NEVER acks on the relaxed (memtable) path"
);
assert_eq!(
kv.batches.load(O::Relaxed),
20,
"every publish took the durable write_batch (== Option B)"
);
}
#[tokio::test]
async fn relaxed_durability_checkpoints_every_n_messages() {
use std::sync::atomic::Ordering as O;
let kv = Arc::new(CountingKv {
inner: MemoryKv::new(),
batches: std::sync::atomic::AtomicUsize::new(0),
relaxed: std::sync::atomic::AtomicUsize::new(0),
});
let mq =
LogMessaging::new(Arc::new(MemStorage::default()), kv.clone()).with_max_unflushed(4);
for i in 0..12u32 {
mq.publish("t", format!("m{i}").as_bytes()).await.unwrap();
}
let relaxed = kv.relaxed.load(O::Relaxed);
let durable = kv.batches.load(O::Relaxed);
assert_eq!(
relaxed + durable,
12,
"every message was committed exactly once"
);
assert!(
durable >= 2,
"durable checkpoints bounded the un-flushed window (got {durable} durable, {relaxed} relaxed)"
);
assert!(
relaxed >= 8,
"most publishes fast-acked on the relaxed path (got {relaxed} relaxed, {durable} durable)"
);
let mut seen = 0;
loop {
let b = mq.claim("t", LEASE, 100, 5).await.unwrap();
if b.is_empty() {
break;
}
seen += b.len();
}
assert_eq!(
seen, 12,
"every relaxed-and-checkpointed publish is claimable"
);
}
#[tokio::test]
async fn topic_policy_set_and_get_roundtrips_and_invalidates_cache() {
let mq = mq();
assert_eq!(
mq.topic_policy("t").await.unwrap(),
None,
"no policy initially"
);
mq.set_topic_policy(
"t",
TopicPolicy {
max_depth: Some(5),
max_rate_per_sec: None,
max_unflushed: Some(3),
},
)
.await
.unwrap();
let got = mq.topic_policy("t").await.unwrap().expect("policy set");
assert_eq!(got.max_depth, Some(5));
assert_eq!(got.max_unflushed, Some(3));
mq.set_topic_policy(
"t",
TopicPolicy {
max_depth: Some(9),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(
mq.topic_policy("t").await.unwrap().unwrap().max_depth,
Some(9),
"set invalidated the cache; the fresh cap is read"
);
}
#[tokio::test]
async fn max_depth_rejects_publish_at_cap() {
let mq = mq();
mq.set_topic_policy(
"t",
TopicPolicy {
max_depth: Some(3),
..Default::default()
},
)
.await
.unwrap();
for i in 0..3u32 {
mq.publish("t", format!("m{i}").as_bytes())
.await
.unwrap_or_else(|e| panic!("under-cap publish {i} should pass, got {e:?}"));
}
assert_eq!(mq.backlog("t").await.unwrap(), 3);
let err = mq.publish("t", b"overflow").await.unwrap_err();
assert!(
matches!(err, MessagingError::DepthExceeded(ref topic) if topic == "t"),
"publish at the cap is rejected with DepthExceeded, got {err:?}"
);
assert_eq!(
mq.backlog("t").await.unwrap(),
3,
"the rejected publish enqueued nothing (backlog unchanged)"
);
let claimed = mq.claim("t", LEASE, 1, 5).await.unwrap();
mq.ack(&claimed[0]).await.unwrap();
assert_eq!(mq.backlog("t").await.unwrap(), 2);
mq.publish("t", b"now-fits")
.await
.expect("under-cap again after a drain");
}
#[tokio::test]
async fn no_policy_leaves_publishing_unbounded() {
let mq = mq();
for i in 0..50u32 {
mq.publish("t", format!("m{i}").as_bytes()).await.unwrap();
}
assert_eq!(mq.backlog("t").await.unwrap(), 50);
}
#[tokio::test]
async fn max_rate_rejects_when_tokens_exhausted() {
let mq = mq();
mq.set_topic_policy(
"t",
TopicPolicy {
max_rate_per_sec: Some(3),
..Default::default()
},
)
.await
.unwrap();
let mut ok = 0;
let mut rate_rejected = 0;
for i in 0..8u32 {
match mq.publish("t", format!("m{i}").as_bytes()).await {
Ok(()) => ok += 1,
Err(MessagingError::RateExceeded(_)) => rate_rejected += 1,
Err(e) => panic!("unexpected error {e:?}"),
}
}
assert!(
ok >= 3,
"the initial burst up to the rate was admitted (got {ok})"
);
assert!(
rate_rejected > 0,
"once the bucket drained, further publishes were rate-rejected (got {rate_rejected})"
);
assert_eq!(ok + rate_rejected, 8);
}
#[tokio::test]
async fn per_topic_max_unflushed_override_forces_durable() {
use std::sync::atomic::Ordering as O;
let kv = Arc::new(CountingKv {
inner: MemoryKv::new(),
batches: std::sync::atomic::AtomicUsize::new(0),
relaxed: std::sync::atomic::AtomicUsize::new(0),
});
let mq =
LogMessaging::new(Arc::new(MemStorage::default()), kv.clone()).with_max_unflushed(8);
mq.set_topic_policy(
"s",
TopicPolicy {
max_unflushed: Some(0),
..Default::default()
},
)
.await
.unwrap();
for i in 0..6u32 {
mq.publish("s", format!("m{i}").as_bytes()).await.unwrap();
}
assert_eq!(
kv.relaxed.load(O::Relaxed),
0,
"the strong per-topic override NEVER took the relaxed path (durable, like budget 0)"
);
assert_eq!(
kv.batches.load(O::Relaxed),
6,
"every publish on the strong-override topic took a durable write_batch"
);
for i in 0..6u32 {
mq.publish("r", format!("m{i}").as_bytes()).await.unwrap();
}
assert!(
kv.relaxed.load(O::Relaxed) > 0,
"a topic with no override inherits the node's relaxed budget (fast-acked at least once)"
);
}
#[tokio::test]
async fn batch_min_over_topics_forces_durable_when_any_is_strong() {
use std::sync::atomic::Ordering as O;
let kv = Arc::new(CountingKv {
inner: MemoryKv::new(),
batches: std::sync::atomic::AtomicUsize::new(0),
relaxed: std::sync::atomic::AtomicUsize::new(0),
});
let mq =
LogMessaging::new(Arc::new(MemStorage::default()), kv.clone()).with_max_unflushed(100);
mq.set_topic_policy(
"strong",
TopicPolicy {
max_unflushed: Some(0),
..Default::default()
},
)
.await
.unwrap();
let msgs: Vec<(String, Vec<u8>)> = vec![
("relaxed".into(), b"a".to_vec()),
("strong".into(), b"b".to_vec()),
("relaxed".into(), b"c".to_vec()),
("relaxed".into(), b"d".to_vec()),
];
mq.publish_batch_ctx(&msgs, None).await.unwrap();
assert_eq!(
kv.relaxed.load(O::Relaxed),
0,
"the strong-override topic in the batch forced the whole coalesced batch durable"
);
assert_eq!(
kv.batches.load(O::Relaxed),
1,
"still ONE coalesced durable write_batch for the batch"
);
}
#[tokio::test]
async fn unsupported_backend_refuses_set_topic_policy() {
struct BareBackend;
#[async_trait]
impl Messaging for BareBackend {
async fn publish(&self, _: &str, _: &[u8]) -> Result<(), MessagingError> {
Ok(())
}
async fn claim(
&self,
_: &str,
_: Duration,
_: usize,
_: u32,
) -> Result<Vec<ClaimedMessage>, MessagingError> {
Ok(Vec::new())
}
async fn ack(&self, _: &ClaimedMessage) -> Result<(), MessagingError> {
Ok(())
}
async fn nack(&self, _: &ClaimedMessage) -> Result<(), MessagingError> {
Ok(())
}
}
let mq = BareBackend;
let err = mq
.set_topic_policy("t", TopicPolicy::default())
.await
.unwrap_err();
assert!(
matches!(err, MessagingError::Unsupported(_)),
"the default set_topic_policy fails closed, got {err:?}"
);
assert_eq!(
mq.topic_policy("t").await.unwrap(),
None,
"the default topic_policy reads None"
);
}
#[tokio::test]
async fn inlined_message_dead_letters_and_redrives_with_its_payload() {
let mq = mq();
mq.publish("t", b"poison").await.unwrap();
let m = mq.claim("t", Duration::ZERO, 10, 1).await.unwrap();
assert!(m[0].inline);
assert!(mq
.claim("t", Duration::ZERO, 10, 1)
.await
.unwrap()
.is_empty());
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 1);
assert_eq!(mq.redrive_dead_letters("t").await.unwrap(), 1);
let back = mq.claim("t", Duration::from_secs(60), 10, 1).await.unwrap();
assert_eq!(back[0].payload, b"poison");
assert!(back[0].inline);
}
#[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 selective_dlq_list_show_redrive_discard_by_filter() {
let mq = mq();
for p in [b"aaa".as_slice(), b"bbb", b"ccc"] {
mq.publish("t", p).await.unwrap();
}
let first = mq.claim("t", Duration::ZERO, 10, 1).await.unwrap();
let bbb = first.iter().find(|m| m.payload == b"bbb").unwrap().clone();
mq.set_last_error(&bbb, "guest-trap:\n injected\u{7} reason")
.await
.unwrap();
assert!(mq
.claim("t", Duration::ZERO, 10, 1)
.await
.unwrap()
.is_empty());
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 3);
let all = mq
.list_dead_letters("t", &DeadLetterFilter::default())
.await
.unwrap();
assert_eq!(all.len(), 3);
assert!(all.iter().all(|d| d.payload.is_none()));
let bbb_dl = all.iter().find(|d| d.id == bbb.id).unwrap();
let err = bbb_dl.last_error.as_deref().unwrap();
assert!(err.contains("guest-trap"));
assert!(
!err.contains('\n') && !err.contains('\u{7}'),
"control characters are sanitized out of last_error"
);
let matched = mq
.list_dead_letters(
"t",
&DeadLetterFilter {
match_last_error: Some("guest-trap".into()),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(matched.len(), 1);
assert_eq!(matched[0].id, bbb.id);
let limited = mq
.list_dead_letters(
"t",
&DeadLetterFilter {
limit: Some(2),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(limited.len(), 2);
let shown = mq
.show_dead_letter("t", "", &bbb.id)
.await
.unwrap()
.unwrap();
assert_eq!(shown.payload.as_deref(), Some(b"bbb".as_slice()));
assert!(shown.last_error.as_deref().unwrap().contains("guest-trap"));
let n = mq
.redrive_dead_letters_filtered(
"t",
&DeadLetterFilter {
id: Some(bbb.id.clone()),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(n, 1);
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 2);
let back = mq.claim("t", LEASE, 10, 5).await.unwrap();
let revived = back.iter().find(|m| m.payload == b"bbb").unwrap();
assert_eq!(revived.attempts, 1, "redrive resets attempts");
let d = mq
.discard_dead_letters("t", &DeadLetterFilter::default())
.await
.unwrap();
assert_eq!(d, 2);
assert_eq!(mq.dead_letter_count("t").await.unwrap(), 0);
}
#[tokio::test]
async fn selective_dlq_older_than_uses_id_age() {
let mq = mq();
mq.publish("t", b"recent").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 none = mq
.list_dead_letters(
"t",
&DeadLetterFilter {
older_than_ms: Some(3_600_000),
..Default::default()
},
)
.await
.unwrap();
assert!(
none.is_empty(),
"a just-published dead-letter isn't 'older than' 1h"
);
let any = mq
.list_dead_letters(
"t",
&DeadLetterFilter {
older_than_ms: Some(0),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(any.len(), 1);
}
#[tokio::test]
async fn group_and_pause_ops_are_confined_to_their_namespaced_topic() {
let mq = mq();
for t in ["siteA/ev", "siteB/ev"] {
assert!(mq
.claim_grouped(t, "g", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
mq.publish(t, b"m").await.unwrap();
}
mq.set_paused("siteA/ev", true).await.unwrap();
assert!(mq.is_paused("siteA/ev").await.unwrap());
assert!(
!mq.is_paused("siteB/ev").await.unwrap(),
"pausing one site's topic never pauses another's"
);
mq.delete_group("siteA/ev", "g").await.unwrap();
assert!(mq.list_groups("siteA/ev").await.unwrap().is_empty());
assert_eq!(
mq.list_groups("siteB/ev").await.unwrap().len(),
1,
"another site's group is untouched by a delete"
);
assert!(mq
.claim_grouped("siteA/ev", "g", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap()
.is_empty());
assert_eq!(
mq.claim_grouped("siteB/ev", "g", StartPosition::Earliest, LEASE, 10, 5)
.await
.unwrap()
.len(),
1,
"the other site's grouped delivery is unaffected"
);
}
#[test]
fn sanitize_reason_strips_control_chars_and_bounds_bytes() {
let s = sanitize_reason("guest-trap:\n\t\u{7}boom");
assert!(s.contains("guest-trap") && s.contains("boom"));
assert!(!s.chars().any(char::is_control), "no control chars survive");
let emoji = sanitize_reason(&"😀".repeat(100));
assert!(emoji.len() <= LAST_ERROR_MAX);
assert_eq!(
emoji.len() % 4,
0,
"bounded on a whole 4-byte char boundary"
);
assert!(sanitize_reason(&"€".repeat(200)).len() <= LAST_ERROR_MAX);
assert_eq!(sanitize_reason("\n\r\t\u{0}"), "");
}
#[tokio::test]
async fn dead_letter_ops_are_confined_to_their_namespaced_topic() {
let mq = mq();
for t in ["siteA/orders", "siteB/orders"] {
mq.publish(t, format!("{t}-poison").as_bytes())
.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());
}
let a = mq
.list_dead_letters("siteA/orders", &DeadLetterFilter::default())
.await
.unwrap();
assert_eq!(a.len(), 1);
assert!(!a[0].id.is_empty() && a.iter().all(|d| !d.id.contains("siteB")));
assert_eq!(
mq.discard_dead_letters("siteA/orders", &DeadLetterFilter::default())
.await
.unwrap(),
1
);
assert_eq!(mq.dead_letter_count("siteA/orders").await.unwrap(), 0);
assert_eq!(
mq.dead_letter_count("siteB/orders").await.unwrap(),
1,
"another site's DLQ is untouched"
);
}
#[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");
}
struct NonAtomicKv {
inner: MemoryKv,
}
#[async_trait]
impl KvStore for NonAtomicKv {
async fn get(&self, k: &str) -> Result<Option<Vec<u8>>, crate::kv::KvError> {
self.inner.get(k).await
}
async fn put(&self, k: &str, v: Vec<u8>) -> Result<(), crate::kv::KvError> {
self.inner.put(k, v).await
}
async fn delete(&self, k: &str) -> Result<(), crate::kv::KvError> {
self.inner.delete(k).await
}
async fn list_prefix(&self, p: &str) -> Result<Vec<String>, crate::kv::KvError> {
self.inner.list_prefix(p).await
}
}
#[tokio::test]
async fn ready_set_fast_path_refused_on_non_atomic_backend() {
let atomic = LogMessaging::new(Arc::new(MemStorage::default()), Arc::new(MemoryKv::new()));
assert!(
atomic.supports_ready_set(),
"an atomic write_batch backend supports the ready-set fast path"
);
let seq = LogMessaging::new(
Arc::new(MemStorage::default()),
Arc::new(NonAtomicKv {
inner: MemoryKv::new(),
}),
);
assert!(
!seq.supports_ready_set(),
"a sequential (non-atomic) write_batch backend MUST refuse the fast path (B2)"
);
}
#[tokio::test]
async fn publish_adds_ready_marker_and_claim_drain_prunes_it() {
let mq = mq();
assert!(
mq.ready_topics().await.unwrap().is_empty(),
"idle: no ready markers"
);
mq.publish("t", b"a").await.unwrap();
assert_eq!(
mq.ready_topics().await.unwrap(),
vec!["t".to_string()],
"a publish flags the topic ready (B1)"
);
let batch = mq.claim("t", LEASE, 16, 5).await.unwrap();
assert_eq!(batch.len(), 1);
assert_eq!(
mq.ready_topics().await.unwrap(),
vec!["t".to_string()],
"a leased-but-unacked message keeps the marker (still pending work)"
);
mq.ack(&batch[0]).await.unwrap();
let empty = mq.claim("t", LEASE, 16, 5).await.unwrap();
assert!(empty.is_empty(), "nothing left to claim");
assert!(
mq.ready_topics().await.unwrap().is_empty(),
"a claim over the now-empty work-queue prunes the ready marker (B4)"
);
}
#[tokio::test]
async fn rebuild_recovers_a_lost_ready_marker() {
let mq = mq();
mq.publish("t", b"a").await.unwrap();
mq.kv.delete(&ready_key("t")).await.unwrap();
assert!(
mq.ready_topics().await.unwrap().is_empty(),
"marker lost — the drainer would skip this topic on the fast path"
);
assert_eq!(mq.backlog("t").await.unwrap(), 1);
let n = mq.rebuild_ready_set().await.unwrap();
assert_eq!(n, 1, "rebuild finds the one topic with work");
assert_eq!(
mq.ready_topics().await.unwrap(),
vec!["t".to_string()],
"rebuild re-adds the lost marker (B6/B18)"
);
assert_eq!(mq.claim("t", LEASE, 16, 5).await.unwrap().len(), 1);
}
#[tokio::test]
async fn rebuild_is_add_only_never_prunes_a_stale_marker() {
let mq = mq();
mq.kv.put(&ready_key("stale"), Vec::new()).await.unwrap();
mq.publish("real", b"a").await.unwrap();
mq.kv.delete(&ready_key("real")).await.unwrap();
let _ = mq.rebuild_ready_set().await.unwrap();
let ready: std::collections::HashSet<String> =
mq.ready_topics().await.unwrap().into_iter().collect();
assert!(
ready.contains("real"),
"rebuild HEALS the missing marker for a topic with work"
);
assert!(
ready.contains("stale"),
"rebuild does NOT prune the stale marker (A-1: pruning is the gated claim path's job, \
never the racy rebuild)"
);
assert!(mq.claim("stale", LEASE, 16, 5).await.unwrap().is_empty());
assert!(
!mq.ready_topics()
.await
.unwrap()
.contains(&"stale".to_string()),
"a cheap empty claim over the stale-marked topic prunes it (B5)"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_publish_is_never_stranded_by_a_draining_prune() {
let mq = Arc::new(mq());
const PUBLISHERS: u32 = 8;
const PER: u32 = 60;
let total = PUBLISHERS * PER;
let mut pubs = Vec::new();
for p in 0..PUBLISHERS {
let mq = mq.clone();
pubs.push(tokio::spawn(async move {
for i in 0..PER {
mq.publish("t", format!("{p}-{i}").as_bytes())
.await
.unwrap();
tokio::task::yield_now().await;
}
}));
}
let mq_drain = mq.clone();
let drainer = tokio::spawn(async move {
let mut delivered = std::collections::HashSet::new();
let deadline = std::time::Instant::now() + Duration::from_secs(20);
while delivered.len() < total as usize {
if std::time::Instant::now() > deadline {
break;
}
let batch = mq_drain.claim("t", LEASE, 8, 5).await.unwrap();
if batch.is_empty() {
tokio::task::yield_now().await;
continue;
}
for m in batch {
delivered.insert(String::from_utf8(m.payload.clone()).unwrap());
mq_drain.ack(&m).await.unwrap();
}
}
delivered
});
for h in pubs {
h.await.unwrap();
}
let delivered = drainer.await.unwrap();
assert_eq!(
delivered.len(),
total as usize,
"gate 4: every published message was delivered with NO stranding (no rebuild used) — \
the prune-vs-publish race is closed directly"
);
}
#[tokio::test]
async fn ready_set_keys_are_per_topic_isolated() {
let mq = mq();
mq.publish("projA/orders", b"a").await.unwrap();
mq.publish("projB/orders", b"b").await.unwrap();
let mut topics = mq.ready_topics().await.unwrap();
topics.sort();
assert_eq!(
topics,
vec!["projA/orders".to_string(), "projB/orders".to_string()],
"same-named topics in two namespaces have distinct ready keys (gate 7)"
);
assert_eq!(ready_key("projA/orders"), "mqready/projA/orders");
assert_ne!(ready_key("projA/orders"), ready_key("projB/orders"));
let a = mq.claim("projA/orders", LEASE, 16, 5).await.unwrap();
assert_eq!(a.len(), 1);
mq.ack(&a[0]).await.unwrap();
assert!(mq
.claim("projA/orders", LEASE, 16, 5)
.await
.unwrap()
.is_empty());
assert_eq!(
mq.ready_topics().await.unwrap(),
vec!["projB/orders".to_string()],
"draining projA prunes only projA's key; projB's marker is untouched (gate 7)"
);
}
#[tokio::test]
async fn ready_set_removals_and_readds_are_never_relaxed() {
let counting = Arc::new(CountingKv {
inner: MemoryKv::new(),
batches: std::sync::atomic::AtomicUsize::new(0),
relaxed: std::sync::atomic::AtomicUsize::new(0),
});
let mq = LogMessaging::new(Arc::new(MemStorage::default()), counting.clone())
.with_max_unflushed(100); mq.publish("t", b"a").await.unwrap();
let relaxed_after_publish = counting.relaxed.load(std::sync::atomic::Ordering::Relaxed);
let batches_before = counting.batches.load(std::sync::atomic::Ordering::Relaxed);
let batch = mq.claim("t", LEASE, 16, 5).await.unwrap();
assert_eq!(batch.len(), 1);
mq.nack(&batch[0]).await.unwrap();
let relaxed_after = counting.relaxed.load(std::sync::atomic::Ordering::Relaxed);
assert_eq!(
relaxed_after, relaxed_after_publish,
"claim-drain prune + nack re-add must NEVER use write_batch_relaxed (B3)"
);
assert!(
counting.batches.load(std::sync::atomic::Ordering::Relaxed) > batches_before,
"the claim + nack committed via the DURABLE write_batch"
);
}
#[tokio::test]
async fn delivery_stats_report_ready_set_size() {
let mq = mq();
mq.publish("a", b"1").await.unwrap();
mq.publish("b", b"2").await.unwrap();
let stats = mq.delivery_stats("node-1").await.unwrap();
assert_eq!(stats.ready_set_size, 2, "two topics flagged ready");
assert_eq!(stats.this_node, "node-1");
mq.rebuild_ready_set().await.unwrap();
let stats = mq.delivery_stats("node-1").await.unwrap();
assert!(
stats.last_rebuild_age_ms.is_some(),
"a completed rebuild stamps the age"
);
}
#[tokio::test]
async fn due_topics_report_leased_deadlines_only() {
let mq = mq();
mq.publish("t", b"a").await.unwrap();
assert!(
mq.due_topics().await.unwrap().is_empty(),
"a claimable-now topic is not on the redelivery due-heap"
);
let batch = mq.claim("t", Duration::from_secs(60), 16, 5).await.unwrap();
assert_eq!(batch.len(), 1);
let due = mq.due_topics().await.unwrap();
assert_eq!(due.len(), 1, "the leased topic has a future deadline");
assert_eq!(due[0].0, "t");
assert!(due[0].1 > now_unix_ms(), "the deadline is in the future");
mq.ack(&batch[0]).await.unwrap();
}
#[tokio::test]
async fn grouped_topic_marker_not_pruned_by_workqueue_claim() {
let mq = mq();
mq.publish("g", b"x").await.unwrap();
let _ = mq
.claim_grouped("g", "workers", StartPosition::Earliest, LEASE, 16, 5)
.await
.unwrap();
mq.publish("g", b"y").await.unwrap();
assert_eq!(mq.ready_topics().await.unwrap(), vec!["g".to_string()]);
let _ = mq.claim("g", LEASE, 16, 5).await.unwrap();
assert_eq!(
mq.ready_topics().await.unwrap(),
vec!["g".to_string()],
"a work-queue claim never prunes a grouped topic's shared marker (B5)"
);
}
}