use std::fmt::{Debug, Formatter};
use std::future::{Future, ready};
use std::sync::{Arc, Mutex};
use bytes::BytesMut;
use fred::interfaces::{StreamsInterface, TransactionInterface};
use fred::types::Value;
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::Bindings;
use ruststream::{
DefaultPublish, HeaderMap, OutgoingMessage, OwnedTransactions, PairError, PublishPolicy,
Publisher, Take, Transaction, TransactionalPublisher,
};
use tracing::warn;
use crate::broker::{ConnectedRedisBroker, RedisCore};
use crate::envelope::frame;
use crate::list::push;
use crate::partition::{RedisPublishOptions, resolved_headers};
use crate::pubsub::PubSubMode;
use crate::pubsub::send;
use crate::route::Route;
use crate::{convert::fields_for_publish, error::RedisError};
type Buffered = (String, Vec<(String, Vec<u8>)>);
async fn flush_block(core: &RedisCore, buffered: Vec<Buffered>) -> Result<(), RedisError> {
if buffered.is_empty() {
return Ok(());
}
let pool = core.pool()?;
let txn = pool.next().multi();
for (key, fields) in buffered {
let _: () = txn
.xadd(key, false, None::<()>, "*", fields)
.await
.map_err(RedisError::publish)?;
}
let _: Value = txn.exec(true).await.map_err(RedisError::publish)?;
Ok(())
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[must_use]
pub struct RedisPublish;
impl PublishPolicy<ConnectedRedisBroker> for RedisPublish {
type Live = RedisPublisher;
fn pair(
self,
connected: &ConnectedRedisBroker,
) -> impl Future<Output = Result<Self::Live, PairError>> {
ready(Ok(connected.publisher()))
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self, channel: &str) -> Bindings {
crate::asyncapi::channel(&crate::asyncapi::Publish::stream(channel))
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[must_use]
pub struct RedisDefaultPublish;
impl PublishPolicy<ConnectedRedisBroker> for RedisDefaultPublish {
type Live = RedisDefaultPublisher;
fn pair(
self,
connected: &ConnectedRedisBroker,
) -> impl Future<Output = Result<Self::Live, PairError>> {
ready(Ok(connected.default_publisher()))
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self, channel: &str) -> Bindings {
crate::asyncapi::channel(&crate::asyncapi::Publish::stream(channel))
}
}
impl DefaultPublish for ConnectedRedisBroker {
type Policy = RedisDefaultPublish;
}
#[derive(Clone)]
pub struct RedisDefaultPublisher {
core: Arc<RedisCore>,
round_name: u64,
}
impl Debug for RedisDefaultPublisher {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisDefaultPublisher")
.field("core", &self.core)
.finish_non_exhaustive()
}
}
impl RedisDefaultPublisher {
pub(crate) const fn round_name(&self) -> u64 {
self.round_name
}
pub(crate) fn new(core: Arc<RedisCore>) -> Self {
Self {
round_name: core.rounds().publisher(),
core,
}
}
}
impl Publisher for RedisDefaultPublisher {
type Payload = Take;
type Error = RedisError;
type Options = RedisPublishOptions;
async fn publish(
&self,
msg: OutgoingMessage<'_, BytesMut>,
options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
let (key, payload, headers) = msg.into_parts();
let headers = resolved_headers(headers, options);
if let Some(round) = self
.core
.rounds()
.joined(self.round_name, joins_round(options))
{
let joined = match self.core.routes().route(key) {
Route::Stream => {
round
.xadd(key, fields_for_publish(Vec::from(payload), &headers))
.await
}
Route::List { envelope } => {
round
.lpush(key, frame(envelope.as_ref(), &payload, &headers), None)
.await
}
Route::Channel { mode, envelope } => {
round
.publish(
key,
frame(envelope.as_ref(), &payload, &headers),
mode == PubSubMode::Sharded,
)
.await
}
};
return joined.map_err(RedisError::publish);
}
match self.core.routes().route(key) {
Route::Stream => append(&self.core, key, Vec::from(payload), &headers).await,
Route::List { envelope } => {
push(&self.core, envelope.as_ref(), None, key, &payload, headers).await
}
Route::Channel { mode, envelope } => {
send(&self.core, mode, envelope.as_ref(), key, &payload, headers).await
}
}
}
}
pub(crate) fn joins_round(options: Option<&RedisPublishOptions>) -> bool {
options.is_some_and(|options| options.join_round)
}
async fn append(
core: &RedisCore,
key: &str,
payload: Vec<u8>,
headers: &HeaderMap,
) -> Result<(), RedisError> {
let pool = core.pool()?;
let _: String = pool
.xadd(
key,
false,
None::<()>,
"*",
fields_for_publish(payload, headers),
)
.await
.map_err(RedisError::publish)?;
Ok(())
}
#[derive(Clone)]
pub struct RedisPublisher {
core: Arc<RedisCore>,
txn: Arc<Mutex<Option<Vec<Buffered>>>>,
round_name: u64,
}
impl Debug for RedisPublisher {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisPublisher")
.field("core", &self.core)
.finish_non_exhaustive()
}
}
impl RedisPublisher {
pub(crate) const fn round_name(&self) -> u64 {
self.round_name
}
pub(crate) fn new(core: Arc<RedisCore>) -> Self {
Self {
round_name: core.rounds().publisher(),
core,
txn: Arc::new(Mutex::new(None)),
}
}
fn check_transactions_supported(&self) -> Result<(), RedisError> {
if self.core.transactions_supported() {
return Ok(());
}
Err(RedisError::InvalidOptions(
"transactions are only supported on standalone and sentinel topologies".to_owned(),
))
}
fn buffer_if_in_txn(&self, entry: &Buffered) -> bool {
let mut guard = self.txn.lock().expect("redis publisher mutex poisoned");
let buffered = guard.as_mut().is_some_and(|buffer| {
buffer.push(entry.clone());
true
});
drop(guard);
buffered
}
}
impl Publisher for RedisPublisher {
type Payload = Take;
type Error = RedisError;
type Options = RedisPublishOptions;
async fn publish(
&self,
msg: OutgoingMessage<'_, BytesMut>,
options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
if let Some(round) = self
.core
.rounds()
.joined(self.round_name, joins_round(options))
{
let (key, payload, headers) = msg.into_parts();
let fields =
fields_for_publish(Vec::from(payload), &resolved_headers(headers, options));
return round.xadd(key, fields).await.map_err(RedisError::publish);
}
let (key, payload, headers) = msg.into_parts();
let entry: Buffered = (
key.to_owned(),
fields_for_publish(Vec::from(payload), &resolved_headers(headers, options)),
);
if self.buffer_if_in_txn(&entry) {
return Ok(());
}
let pool = self.core.pool()?;
let (key, fields) = entry;
let _: String = pool
.xadd(key, false, None::<()>, "*", fields)
.await
.map_err(RedisError::publish)?;
Ok(())
}
}
impl TransactionalPublisher for RedisPublisher {
fn begin_transaction(&self) -> impl Future<Output = Result<(), Self::Error>> {
if let Err(err) = self.check_transactions_supported() {
return ready(Err(err));
}
{
let mut guard = self.txn.lock().expect("redis publisher mutex poisoned");
if guard.is_some() {
return ready(Err(RedisError::TransactionBusy));
}
*guard = Some(Vec::new());
}
ready(Ok(()))
}
async fn commit(&self) -> Result<(), Self::Error> {
let buffered = self
.txn
.lock()
.expect("redis publisher mutex poisoned")
.take()
.ok_or(RedisError::NoTransaction)?;
flush_block(&self.core, buffered).await
}
fn abort(&self) -> impl Future<Output = Result<(), Self::Error>> {
ready(
self.txn
.lock()
.expect("redis publisher mutex poisoned")
.take()
.ok_or(RedisError::NoTransaction)
.map(|_| ()),
)
}
}
impl OwnedTransactions for RedisPublisher {
type Transaction = RedisTransaction;
fn transaction(&self) -> impl Future<Output = Result<RedisTransaction, RedisError>> {
if let Err(err) = self.check_transactions_supported() {
return ready(Err(err));
}
ready(Ok(RedisTransaction {
core: Arc::clone(&self.core),
buffered: Vec::new(),
settled: false,
}))
}
}
#[must_use = "a transaction does nothing until settled with commit() or abort()"]
pub struct RedisTransaction {
core: Arc<RedisCore>,
buffered: Vec<Buffered>,
settled: bool,
}
impl Debug for RedisTransaction {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RedisTransaction")
.field("buffered", &self.buffered.len())
.field("settled", &self.settled)
.finish_non_exhaustive()
}
}
impl Drop for RedisTransaction {
fn drop(&mut self) {
if !self.settled {
warn!(
target: "ruststream_fred",
buffered = self.buffered.len(),
"owned transaction dropped without commit or abort; its buffered messages are \
discarded"
);
}
}
}
impl Transaction for RedisTransaction {
type Payload = Take;
type Error = RedisError;
type Options = RedisPublishOptions;
fn publish(
&mut self,
msg: OutgoingMessage<'_, BytesMut>,
options: Option<&Self::Options>,
) -> impl Future<Output = Result<(), Self::Error>> {
let (key, payload, headers) = msg.into_parts();
self.buffered.push((
key.to_owned(),
fields_for_publish(Vec::from(payload), &resolved_headers(headers, options)),
));
ready(Ok(()))
}
async fn commit(mut self) -> Result<(), Self::Error> {
self.settled = true;
flush_block(&self.core, std::mem::take(&mut self.buffered)).await
}
fn abort(mut self) -> impl Future<Output = Result<(), Self::Error>> {
self.settled = true;
ready(Ok(()))
}
}