use std::future::{Future, ready};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::Bindings;
use ruststream::{
AddressedCopies, RedeliveryAddress, RedeliveryAddressed, RetryDeclaration, SubscriptionSource,
};
#[cfg(feature = "asyncapi")]
use crate::asyncapi;
use crate::broker::ConnectedRedisBroker;
use crate::delay::{DelayConfig, DelayedRetry};
use crate::{error::RedisError, subscriber::RedisSubscriber};
pub use crate::publisher::RedisPublish as Publish;
pub use crate::publisher::RedisPublish as TransactionalPublish;
pub mod prelude {
pub use ruststream::prelude::*;
pub use ruststream::{
OwnedTransactions, Positioned, Seeker, Transaction, TransactionalPublisher,
};
pub use super::{Publish, RedisStream, StreamStart, TransactionalPublish};
pub use crate::context::{
PipelineContext, PoolContext, StreamBatchContext, StreamContext, keys,
};
pub use crate::pipeline::{Bindable, InRound, RedisPipelineSteps};
pub use crate::{
DelayedRetry, PARTITION_KEY_HEADER, RedisBroker, RedisGroupPosition, RedisGroupSeeker,
RedisPublishOptions, RedisPublishSteps, RedisSubscribeExt,
};
#[cfg(any(
feature = "tls-rustls",
feature = "tls-rustls-ring",
feature = "tls-native-tls"
))]
pub use crate::{TlsConfig, TlsConnector};
}
const DEFAULT_BLOCK: Duration = Duration::from_secs(5);
fn auto_consumer() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(0);
let n = COUNTER.fetch_add(1, Ordering::Relaxed);
format!("ruststream-{n}")
}
#[derive(Debug, Clone, Default)]
pub enum StreamStart {
#[default]
New,
Beginning,
Id(String),
}
impl StreamStart {
pub(crate) fn as_id(&self) -> &str {
match self {
Self::New => "$",
Self::Beginning => "0",
Self::Id(id) => id,
}
}
}
#[derive(Debug, Clone)]
pub(crate) enum ReadMode {
Fresh,
Reclaim { min_idle: Duration },
Claiming { min_idle: Duration },
}
impl ReadMode {
pub(crate) const fn requeue(&self) -> RequeueMode {
match self {
Self::Fresh | Self::Reclaim { .. } => RequeueMode::Republish,
Self::Claiming { min_idle } => RequeueMode::LeavePending {
min_idle: *min_idle,
},
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequeueMode {
Republish,
LeavePending { min_idle: Duration },
}
#[derive(Debug, Clone)]
#[must_use]
pub struct RedisStream {
key: String,
group: Option<String>,
consumer: Option<String>,
block: Option<Duration>,
start: StreamStart,
mode: ReadMode,
delayed_retry: Option<DelayedRetry>,
dead_letter: Option<String>,
}
impl RedisStream {
pub fn new(key: impl Into<String>) -> Self {
Self {
key: key.into(),
group: None,
consumer: None,
block: None,
start: StreamStart::New,
mode: ReadMode::Fresh,
delayed_retry: None,
dead_letter: None,
}
}
pub fn reclaim(key: impl Into<String>, min_idle: Duration) -> Self {
Self {
key: key.into(),
group: None,
consumer: None,
block: None,
start: StreamStart::New,
mode: ReadMode::Reclaim { min_idle },
delayed_retry: None,
dead_letter: None,
}
}
pub fn claiming(key: impl Into<String>, min_idle: Duration) -> Self {
Self {
key: key.into(),
group: None,
consumer: None,
block: None,
start: StreamStart::New,
mode: ReadMode::Claiming { min_idle },
delayed_retry: None,
dead_letter: None,
}
}
pub fn group(mut self, group: impl Into<String>) -> Self {
self.group = Some(group.into());
self
}
pub fn consumer(mut self, consumer: impl Into<String>) -> Self {
self.consumer = Some(consumer.into());
self
}
pub const fn block(mut self, block: Duration) -> Self {
self.block = Some(block);
self
}
pub fn start_id(mut self, start: StreamStart) -> Self {
self.start = start;
self
}
pub fn delayed_retry(mut self, retry: DelayedRetry) -> Self {
self.delayed_retry = Some(retry);
self
}
#[must_use]
pub fn key(&self) -> &str {
&self.key
}
pub(crate) fn group_or_err(&self) -> Result<&str, RedisError> {
self.group.as_deref().ok_or_else(|| {
RedisError::InvalidOptions(format!(
"stream subscription on `{}` requires a consumer group: call .group(name)",
self.key
))
})
}
pub(crate) fn consumer_or_auto(&self) -> String {
self.consumer.clone().unwrap_or_else(auto_consumer)
}
pub(crate) fn block_or_default(&self) -> Duration {
self.block.unwrap_or(DEFAULT_BLOCK)
}
pub(crate) const fn start(&self) -> &StreamStart {
&self.start
}
pub(crate) fn mode(&self) -> ReadMode {
self.mode.clone()
}
pub(crate) fn dead_letter(&self) -> Option<&str> {
self.dead_letter.as_deref()
}
fn with_declaration(mut self, declaration: &RetryDeclaration) -> Self {
self.dead_letter = declaration.dead_letter().map(str::to_owned);
self
}
pub(crate) fn delay_config(&self) -> Option<DelayConfig> {
self.delayed_retry.as_ref().map(DelayConfig::from_retry)
}
fn redelivery_key(&self) -> RedeliveryAddress {
RedeliveryAddress::new(self.key.clone())
}
#[cfg(feature = "asyncapi")]
fn describe(&self) -> Bindings {
let (mode, min_idle) = match self.mode {
ReadMode::Fresh => (asyncapi::ReadMode::Fresh, None),
ReadMode::Reclaim { min_idle } => (asyncapi::ReadMode::Reclaim, Some(min_idle)),
ReadMode::Claiming { min_idle } => (asyncapi::ReadMode::Claiming, Some(min_idle)),
};
asyncapi::channel(&asyncapi::StreamSubscription::new(
self.group.as_deref(),
self.consumer.as_deref(),
mode,
min_idle.map(|idle| u64::try_from(idle.as_millis()).unwrap_or(u64::MAX)),
))
}
}
impl SubscriptionSource<ConnectedRedisBroker> for RedisStream {
type Subscriber = RedisSubscriber;
type Copies = AddressedCopies;
fn name(&self) -> &str {
self.key()
}
fn declare_retry(self, declaration: &RetryDeclaration) -> Self {
self.with_declaration(declaration)
}
async fn subscribe(
self,
connected: &ConnectedRedisBroker,
) -> Result<Self::Subscriber, RedisError> {
connected.subscribe(self).await
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self) -> Bindings {
self.describe()
}
}
impl RedeliveryAddressed<ConnectedRedisBroker> for RedisStream {
fn redelivery_address(
&self,
_connected: &ConnectedRedisBroker,
) -> impl Future<Output = Result<RedeliveryAddress, RedisError>> {
ready(Ok(self.redelivery_key()))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn group_is_required() {
let err = RedisStream::new("orders").group_or_err().unwrap_err();
assert!(matches!(err, RedisError::InvalidOptions(msg) if msg.contains("consumer group")));
}
#[test]
fn group_set_resolves() {
let s = RedisStream::new("orders").group("workers");
assert_eq!(s.group_or_err().expect("group set"), "workers");
}
#[test]
fn start_maps_to_redis_ids() {
assert_eq!(StreamStart::New.as_id(), "$");
assert_eq!(StreamStart::Beginning.as_id(), "0");
assert_eq!(StreamStart::Id("5-0".into()).as_id(), "5-0");
}
#[test]
fn reclaim_carries_min_idle() {
let s = RedisStream::reclaim("orders", Duration::from_secs(30)).group("g");
assert!(matches!(s.mode(), ReadMode::Reclaim { min_idle } if min_idle.as_secs() == 30));
}
#[test]
fn claiming_carries_min_idle() {
let s = RedisStream::claiming("orders", Duration::from_secs(30)).group("g");
assert!(matches!(s.mode(), ReadMode::Claiming { min_idle } if min_idle.as_secs() == 30));
}
#[test]
fn every_read_mode_addresses_its_copies_to_the_stream_key() {
let idle = Duration::from_secs(30);
for source in [
RedisStream::new("orders"),
RedisStream::claiming("orders", idle),
RedisStream::reclaim("orders", idle),
] {
assert_eq!(source.redelivery_key().as_str(), "orders");
}
}
#[test]
fn the_claiming_mode_retries_by_leaving_the_entry_pending() {
let idle = Duration::from_secs(30);
assert_eq!(
RedisStream::claiming("orders", idle).mode().requeue(),
RequeueMode::LeavePending { min_idle: idle }
);
assert_eq!(
RedisStream::new("orders").mode().requeue(),
RequeueMode::Republish
);
assert_eq!(
RedisStream::reclaim("orders", idle).mode().requeue(),
RequeueMode::Republish
);
}
}