use std::fmt::{Debug, Formatter};
use std::future::{Future, ready};
use std::num::NonZeroUsize;
use std::sync::Arc;
use bytes::Bytes;
use fred::clients::{Client, Pool};
use fred::interfaces::{ClientLike, PubsubInterface};
use fred::types::config::Server;
use fred::types::{Message, MessageKind};
use futures::Stream;
use futures::stream::unfold;
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::Bindings;
use ruststream::codec::Codec;
use ruststream::runtime::IntoSource;
use ruststream::{
AckError, AddressedCopies, BatchSubscriber, BufferedSubscriber, HeaderMap, IncomingMessage,
Lend, NamedCopies, OutgoingMessage, PairError, Partitioned, PublishPolicy, Publisher,
RedeliveryAddress, RedeliveryAddressed, RetryDeclaration, SubscriptionSource,
};
use tokio::runtime::Handle;
use tokio::sync::broadcast::Receiver;
use tokio::sync::broadcast::error::{RecvError, TryRecvError};
use tokio::task::JoinHandle;
use tracing::warn;
use crate::broker::{ConnectedRedisBroker, RedisCore};
use crate::envelope::{SharedEnvelope, frame, unframe};
use crate::loopback::{InFlight, Tap};
use crate::partition::{RedisPublishOptions, resolved_headers};
use crate::pipeline::{PubSubForm, RoundMessage, Window};
use crate::publisher::joins_round;
use crate::route::Route;
use crate::{error::RedisError, message::PARTITION_KEY_HEADER};
pub use crate::pubsub::RedisPubSubPublish as Publish;
pub mod prelude {
pub use ruststream::prelude::*;
pub use super::{PubSubMode, Publish, RedisPubSub, RedisPubSubPattern};
pub use crate::context::{PipelineContext, PoolContext, PubSubContext, keys};
pub use crate::pipeline::{Bindable, InRound, RedisPipelineSteps};
pub use crate::{PARTITION_KEY_HEADER, RedisBroker, RedisPublishOptions, RedisPublishSteps};
#[cfg(any(
feature = "tls-rustls",
feature = "tls-rustls-ring",
feature = "tls-native-tls"
))]
pub use crate::{TlsConfig, TlsConnector};
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum PubSubMode {
#[default]
Classic,
Sharded,
}
impl PubSubMode {
#[cfg(feature = "asyncapi")]
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::Classic => "classic",
Self::Sharded => "sharded",
}
}
}
#[derive(Clone)]
#[must_use]
pub struct RedisPubSub {
channel: String,
mode: PubSubMode,
codec: Option<SharedEnvelope>,
dead_letter: Option<String>,
buffer: Option<NonZeroUsize>,
}
impl Debug for RedisPubSub {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisPubSub")
.field("channel", &self.channel)
.field("mode", &self.mode)
.field("codec", &self.codec.is_some())
.field("dead_letter", &self.dead_letter)
.field("buffer", &self.buffer)
.finish()
}
}
impl RedisPubSub {
pub fn new(channel: impl Into<String>) -> Self {
Self {
channel: channel.into(),
mode: PubSubMode::default(),
codec: None,
dead_letter: None,
buffer: None,
}
}
pub const fn mode(mut self, mode: PubSubMode) -> Self {
self.mode = mode;
self
}
pub fn codec(mut self, codec: impl Codec + 'static) -> Self {
self.codec = Some(Arc::new(codec));
self
}
pub const fn buffer(mut self, messages: NonZeroUsize) -> Self {
self.buffer = Some(messages);
self
}
#[must_use]
pub fn channel(&self) -> &str {
&self.channel
}
pub(crate) const fn buffer_size(&self) -> Option<NonZeroUsize> {
self.buffer
}
pub(crate) const fn delivery_mode(&self) -> PubSubMode {
self.mode
}
pub(crate) fn codec_handle(&self) -> Option<SharedEnvelope> {
self.codec.clone()
}
pub(crate) fn route(&self) -> Route {
Route::Channel {
mode: self.mode,
envelope: self.codec.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
}
fn redelivery_channel(&self) -> RedeliveryAddress {
RedeliveryAddress::new(self.channel.clone())
}
#[cfg(feature = "asyncapi")]
fn describe(&self) -> Bindings {
crate::asyncapi::channel(&crate::asyncapi::PubSubSubscription::new(
self.mode.as_str(),
false,
crate::asyncapi::Envelope::of(self.codec.as_ref()),
))
}
}
#[derive(Clone)]
#[must_use]
pub struct RedisPubSubPattern {
pattern: String,
codec: Option<SharedEnvelope>,
buffer: Option<NonZeroUsize>,
}
impl Debug for RedisPubSubPattern {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisPubSubPattern")
.field("pattern", &self.pattern)
.field("codec", &self.codec.is_some())
.field("buffer", &self.buffer)
.finish()
}
}
impl RedisPubSubPattern {
pub fn new(pattern: impl Into<String>) -> Self {
Self {
pattern: pattern.into(),
codec: None,
buffer: None,
}
}
pub fn codec(mut self, codec: impl Codec + 'static) -> Self {
self.codec = Some(Arc::new(codec));
self
}
pub const fn buffer(mut self, messages: NonZeroUsize) -> Self {
self.buffer = Some(messages);
self
}
#[must_use]
pub fn pattern(&self) -> &str {
&self.pattern
}
pub(crate) const fn buffer_size(&self) -> Option<NonZeroUsize> {
self.buffer
}
pub(crate) fn codec_handle(&self) -> Option<SharedEnvelope> {
self.codec.clone()
}
#[cfg(feature = "asyncapi")]
fn describe(&self) -> Bindings {
crate::asyncapi::channel(&crate::asyncapi::PubSubSubscription::new(
PubSubMode::Classic.as_str(),
true,
crate::asyncapi::Envelope::of(self.codec.as_ref()),
))
}
}
impl IntoSource for RedisPubSub {
type Source = Self;
fn into_source(self) -> Self {
self
}
}
impl SubscriptionSource<ConnectedRedisBroker> for RedisPubSub {
type Subscriber = RedisPubSubSubscriber;
type Copies = AddressedCopies;
fn name(&self) -> &str {
self.channel()
}
fn declare_retry(self, declaration: &RetryDeclaration) -> Self {
self.with_declaration(declaration)
}
async fn subscribe(
self,
connected: &ConnectedRedisBroker,
) -> Result<Self::Subscriber, RedisError> {
connected.subscribe_pubsub(self).await
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self) -> Bindings {
self.describe()
}
}
impl RedeliveryAddressed<ConnectedRedisBroker> for RedisPubSub {
fn redelivery_address(
&self,
_connected: &ConnectedRedisBroker,
) -> impl Future<Output = Result<RedeliveryAddress, RedisError>> {
ready(Ok(self.redelivery_channel()))
}
}
impl IntoSource for RedisPubSubPattern {
type Source = Self;
fn into_source(self) -> Self {
self
}
}
impl SubscriptionSource<ConnectedRedisBroker> for RedisPubSubPattern {
type Subscriber = RedisPubSubSubscriber;
type Copies = NamedCopies;
fn name(&self) -> &str {
self.pattern()
}
async fn subscribe(
self,
connected: &ConnectedRedisBroker,
) -> Result<Self::Subscriber, RedisError> {
connected.subscribe_pubsub_pattern(self).await
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self) -> Bindings {
self.describe()
}
}
pub struct RedisPubSubSubscriber(BufferedSubscriber<PubSubWire>);
impl Debug for RedisPubSubSubscriber {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisPubSubSubscriber")
.finish_non_exhaustive()
}
}
impl RedisPubSubSubscriber {
pub(crate) fn new(wire: PubSubWire) -> Self {
Self(BufferedSubscriber::new(wire))
}
}
impl ruststream::Subscriber for RedisPubSubSubscriber {
type Message = RedisPubSubMessage;
type Error = RedisError;
fn stream(&mut self) -> impl Stream<Item = Result<Self::Message, Self::Error>> + Send + '_ {
self.0.stream()
}
}
impl BatchSubscriber for RedisPubSubSubscriber {
type Batch = Vec<RedisPubSubMessage>;
fn batches(
&mut self,
size: NonZeroUsize,
) -> impl Stream<Item = Result<Self::Batch, Self::Error>> + Send + '_ {
self.0.batches(size)
}
}
pub(crate) enum Resubscribe {
Channel(String),
ShardedChannel(String),
Pattern(String),
}
pub(crate) struct Resubscribing {
client: Client,
resubscribing: JoinHandle<()>,
}
impl Resubscribing {
pub(crate) fn new(
client: Client,
mut reconnects: Receiver<Server>,
target: Resubscribe,
runtime: &Handle,
) -> Self {
let resubscriber = client.clone();
let resubscribing = runtime.spawn(async move {
loop {
match reconnects.recv().await {
Ok(_) | Err(RecvError::Lagged(_)) => {}
Err(RecvError::Closed) => return,
}
let (name, outcome) = match &target {
Resubscribe::Channel(name) => (name, resubscriber.subscribe(name).await),
Resubscribe::ShardedChannel(name) => {
(name, resubscriber.ssubscribe(name).await)
}
Resubscribe::Pattern(name) => (name, resubscriber.psubscribe(name).await),
};
if let Err(error) = outcome {
warn!(
subscription = name.as_str(),
%error,
"a Redis Pub/Sub subscription reconnected but could not subscribe again; \
it receives nothing until its connection drops and reconnects"
);
}
}
});
Self {
client,
resubscribing,
}
}
}
pub(crate) struct PubSubWire {
name: String,
client: Resubscribing,
rx: Receiver<Message>,
codec: Option<SharedEnvelope>,
pool: Pool,
tap: Tap,
runtime: Handle,
}
impl Debug for PubSubWire {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PubSubWire").finish_non_exhaustive()
}
}
impl PubSubWire {
pub(crate) fn new(
name: String,
client: Resubscribing,
rx: Receiver<Message>,
codec: Option<SharedEnvelope>,
pool: Pool,
tap: Tap,
runtime: Handle,
) -> Self {
Self {
name,
client,
rx,
codec,
pool,
tap,
runtime,
}
}
pub(crate) fn round_client(&self) -> Client {
self.pool.next().clone()
}
pub(crate) fn round_stream<'a>(
&'a mut self,
window: &'a Arc<Window<PubSubForm>>,
) -> impl Stream<Item = Result<RoundMessage<PubSubForm>, RedisError>> + Send + 'a {
let codec = self.codec.clone();
let (pool, tap, name) = (&self.pool, &self.tap, self.name.as_str());
unfold((&mut self.rx, codec), move |(rx, codec)| async move {
loop {
if let Some(failure) = window.take_failure() {
return Some((Err(failure), (rx, codec)));
}
match rx.recv().await {
Ok(msg) => {
let delivery = to_message(&msg, codec.as_ref(), pool, tap);
window.yielded(rx.len());
let round = window.open();
return Some((
Ok(RoundMessage::new(delivery, Arc::clone(window), round)),
(rx, codec),
));
}
Err(RecvError::Lagged(skipped)) => lagged(name, tap, skipped),
Err(RecvError::Closed) => return None,
}
}
})
}
}
fn lagged(name: &str, tap: &Tap, skipped: u64) {
warn!(
subscription = name,
lost = skipped,
"a Redis Pub/Sub subscription fell behind its buffer and lost the oldest messages; raise \
its buffer or speed up its handler"
);
#[cfg(feature = "testing")]
tap.discard(usize::try_from(skipped).unwrap_or(usize::MAX));
#[cfg(not(feature = "testing"))]
let _ = tap;
}
impl PubSubWire {
pub(crate) fn round_batches<'a>(
&'a mut self,
window: &'a Arc<Window<PubSubForm>>,
size: NonZeroUsize,
) -> impl Stream<Item = Result<Vec<RoundMessage<PubSubForm>>, RedisError>> + Send + 'a {
window.fill_at(size.get());
let codec = self.codec.clone();
let (pool, tap, name) = (&self.pool, &self.tap, self.name.as_str());
unfold((&mut self.rx, codec), move |(rx, codec)| async move {
loop {
if let Some(failure) = window.take_failure() {
return Some((Err(failure), (rx, codec)));
}
let first = match rx.recv().await {
Ok(msg) => msg,
Err(RecvError::Lagged(skipped)) => {
lagged(name, tap, skipped);
continue;
}
Err(RecvError::Closed) => return None,
};
let mut messages = vec![first];
while messages.len() < size.get() {
match rx.try_recv() {
Ok(msg) => messages.push(msg),
Err(TryRecvError::Lagged(skipped)) => lagged(name, tap, skipped),
Err(_) => break,
}
}
let members = u32::try_from(messages.len()).unwrap_or(u32::MAX);
for _ in 0..members {
window.yielded(rx.len());
}
let round = window.open_batch(members);
let batch = messages
.iter()
.map(|msg| {
RoundMessage::new(
to_message(msg, codec.as_ref(), pool, tap),
Arc::clone(window),
round.clone(),
)
})
.collect();
return Some((Ok(batch), (rx, codec)));
}
})
}
}
impl Drop for PubSubWire {
fn drop(&mut self) {
#[cfg(feature = "testing")]
{
let mut unread = 0_usize;
loop {
match self.rx.try_recv() {
Ok(_) => unread += 1,
Err(TryRecvError::Lagged(skipped)) => {
unread =
unread.saturating_add(usize::try_from(skipped).unwrap_or(usize::MAX));
}
Err(_) => break,
}
}
self.tap.discard(unread);
}
self.client.resubscribing.abort();
let client = self.client.client.clone();
self.runtime.spawn(async move {
let _ = client.quit().await;
});
}
}
fn to_message(
msg: &Message,
codec: Option<&SharedEnvelope>,
pool: &Pool,
tap: &Tap,
) -> RedisPubSubMessage {
let raw = msg.value.as_bytes().unwrap_or(&[]);
let (payload, headers) = unframe(codec, raw);
RedisPubSubMessage {
channel: msg.channel.to_string(),
pattern: matches!(msg.kind, MessageKind::PMessage),
payload,
headers,
pool: pool.clone(),
flight: tap.delivered(),
}
}
impl ruststream::Subscriber for PubSubWire {
type Message = RedisPubSubMessage;
type Error = RedisError;
fn stream(&mut self) -> impl Stream<Item = Result<Self::Message, Self::Error>> + Send + '_ {
let codec = self.codec.clone();
let (pool, tap, name) = (&self.pool, &self.tap, self.name.as_str());
unfold((&mut self.rx, codec), move |(rx, codec)| async move {
loop {
match rx.recv().await {
Ok(msg) => {
let message = to_message(&msg, codec.as_ref(), pool, tap);
return Some((Ok(message), (rx, codec)));
}
Err(RecvError::Lagged(skipped)) => lagged(name, tap, skipped),
Err(RecvError::Closed) => return None,
}
}
})
}
}
pub struct RedisPubSubMessage {
channel: String,
pattern: bool,
payload: Bytes,
headers: HeaderMap,
pool: Pool,
#[cfg_attr(
not(feature = "testing"),
allow(
dead_code,
reason = "held for its release on drop, which only `testing` gives"
)
)]
flight: InFlight,
}
impl Debug for RedisPubSubMessage {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisPubSubMessage")
.field("channel", &self.channel)
.field("pattern", &self.pattern)
.field("payload_len", &self.payload.len())
.finish_non_exhaustive()
}
}
impl RedisPubSubMessage {
#[must_use]
pub fn channel(&self) -> &str {
&self.channel
}
#[must_use]
pub fn from_pattern(&self) -> bool {
self.pattern
}
pub(crate) const fn pool(&self) -> &Pool {
&self.pool
}
#[cfg(feature = "testing")]
pub(crate) fn take_flight(&mut self) -> InFlight {
std::mem::take(&mut self.flight)
}
}
impl IncomingMessage for RedisPubSubMessage {
fn payload(&self) -> &[u8] {
&self.payload
}
fn headers(&self) -> &HeaderMap {
&self.headers
}
fn partition_key(&self) -> Option<&[u8]> {
Partitioned::partition_key(self)
}
fn ack(self) -> impl Future<Output = Result<(), AckError>> {
ready(Err(AckError::Unsupported))
}
fn nack(self, _requeue: bool) -> impl Future<Output = Result<(), AckError>> {
ready(Err(AckError::Unsupported))
}
}
impl Partitioned for RedisPubSubMessage {
fn partition_key(&self) -> Option<&[u8]> {
self.headers().get(PARTITION_KEY_HEADER)
}
}
#[derive(Clone, Default)]
#[must_use]
pub struct RedisPubSubPublish {
mode: PubSubMode,
codec: Option<SharedEnvelope>,
}
impl Debug for RedisPubSubPublish {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisPubSubPublish")
.field("mode", &self.mode)
.field("codec", &self.codec.is_some())
.finish()
}
}
impl RedisPubSubPublish {
pub fn new() -> Self {
Self::default()
}
pub const fn mode(mut self, mode: PubSubMode) -> Self {
self.mode = mode;
self
}
pub fn codec(mut self, codec: impl Codec + 'static) -> Self {
self.codec = Some(Arc::new(codec));
self
}
}
impl PublishPolicy<ConnectedRedisBroker> for RedisPubSubPublish {
type Live = RedisPubSubPublisher;
fn pair(
self,
connected: &ConnectedRedisBroker,
) -> impl Future<Output = Result<Self::Live, PairError>> {
ready(Ok(connected.pubsub_publisher(self)))
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self, channel: &str) -> Bindings {
crate::asyncapi::channel(&crate::asyncapi::Publish::pubsub(
channel,
self.mode.as_str(),
crate::asyncapi::Envelope::of(self.codec.as_ref()),
))
}
}
#[derive(Clone)]
pub struct RedisPubSubPublisher {
core: Arc<RedisCore>,
mode: PubSubMode,
codec: Option<SharedEnvelope>,
round_name: u64,
}
impl Debug for RedisPubSubPublisher {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisPubSubPublisher")
.field("mode", &self.mode)
.field("codec", &self.codec.is_some())
.finish_non_exhaustive()
}
}
impl RedisPubSubPublisher {
pub(crate) const fn round_name(&self) -> u64 {
self.round_name
}
pub(crate) fn new(core: Arc<RedisCore>, publish: RedisPubSubPublish) -> Self {
Self {
round_name: core.rounds().publisher(),
core,
mode: publish.mode,
codec: publish.codec,
}
}
}
impl Publisher for RedisPubSubPublisher {
type Payload = Lend;
type Error = RedisError;
type Options = RedisPublishOptions;
async fn publish(
&self,
msg: OutgoingMessage<'_>,
options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
let (channel, payload, headers) = msg.into_parts();
if let Some(round) = self
.core
.rounds()
.joined(self.round_name, joins_round(options))
{
let body = frame(
self.codec.as_ref(),
payload,
&resolved_headers(headers, options),
);
return round
.publish(channel, body, self.mode == PubSubMode::Sharded)
.await
.map_err(RedisError::publish);
}
send(
&self.core,
self.mode,
self.codec.as_ref(),
channel,
payload,
resolved_headers(headers, options),
)
.await
}
}
pub(crate) async fn send(
core: &RedisCore,
mode: PubSubMode,
codec: Option<&SharedEnvelope>,
channel: &str,
payload: &[u8],
headers: HeaderMap,
) -> Result<(), RedisError> {
let pool = core.pool()?;
let client = pool.next();
let body = frame(codec, payload, &headers);
let channel = channel.to_owned();
let _: i64 = match mode {
PubSubMode::Classic => client.publish(channel, body).await,
PubSubMode::Sharded => client.spublish(channel, body).await,
}
.map_err(RedisError::publish)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context::PubSubContext;
use fred::types::config::Config;
use ruststream::BuildContext;
fn offline_pool() -> Pool {
Pool::new(Config::default(), None, None, None, 1).expect("offline pool")
}
#[test]
fn build_context_reads_channel_and_pattern_flag() {
let exact = RedisPubSubMessage {
channel: "events".to_owned(),
pattern: false,
payload: Bytes::from_static(b"{}"),
headers: HeaderMap::new(),
pool: offline_pool(),
flight: InFlight::default(),
};
let cx = PubSubContext::build(&exact);
assert_eq!(cx.channel(), "events");
assert!(!cx.from_pattern());
let matched = RedisPubSubMessage {
channel: "events.user".to_owned(),
pattern: true,
payload: Bytes::from_static(b"{}"),
headers: HeaderMap::new(),
pool: offline_pool(),
flight: InFlight::default(),
};
assert!(PubSubContext::build(&matched).from_pattern());
}
#[test]
fn a_channel_addresses_its_copies_to_itself() {
assert_eq!(
RedisPubSub::new("events").redelivery_channel().as_str(),
"events"
);
assert_eq!(
RedisPubSub::new("events")
.mode(PubSubMode::Sharded)
.redelivery_channel()
.as_str(),
"events"
);
}
}