use std::future::{Future, ready};
use std::sync::Arc;
use std::time::Duration;
use fred::clients::{Client, Pipeline, Pool};
use fred::error::Error;
use fred::interfaces::{KeysInterface, ListInterface, SortedSetsInterface, StreamsInterface};
use fred::types::Value;
use ruststream::AckError;
use super::window::{Form, Segment, Sent};
use crate::convert::fields_for_publish;
use crate::delay::DelayConfig;
use crate::list::RedisListMessage;
#[cfg(feature = "testing")]
use crate::loopback::InFlight;
use crate::message::RedisMessage;
use crate::pubsub::RedisPubSubMessage;
use crate::stream::RequeueMode;
#[derive(Debug)]
pub enum StreamOp {
Ack(String),
Republish(Vec<(String, Vec<u8>)>),
Schedule { score: f64, member: Vec<u8> },
}
#[derive(Debug)]
pub struct StreamForm {
key: Arc<str>,
group: Arc<str>,
delay: Option<DelayConfig>,
requeue: RequeueMode,
}
impl StreamForm {
pub(crate) fn new(
key: impl Into<Arc<str>>,
group: impl Into<Arc<str>>,
delay: Option<DelayConfig>,
requeue: RequeueMode,
) -> Self {
Self {
key: key.into(),
group: group.into(),
delay,
requeue,
}
}
async fn queue<Sink>(&self, sink: &Sink, op: StreamOp) -> Result<(), Error>
where
Sink: StreamsInterface + SortedSetsInterface + KeysInterface + Sync,
{
self.queue_counted(sink, op).await.1
}
async fn queue_counted<Sink>(&self, sink: &Sink, op: StreamOp) -> (usize, Result<(), Error>)
where
Sink: StreamsInterface + SortedSetsInterface + KeysInterface + Sync,
{
let one = |queued: Result<(), Error>| (usize::from(queued.is_ok()), queued);
match op {
StreamOp::Ack(id) => one(sink.xack::<(), _, _, _>(&*self.key, &*self.group, id).await),
StreamOp::Republish(fields) => one(sink
.xadd::<(), _, _, _, _>(&*self.key, false, None::<()>, "*", fields)
.await),
StreamOp::Schedule { score, member } => {
let cfg = self
.delay
.as_ref()
.expect("a schedule is owed only where a delay queue is named");
let scheduled = sink
.zadd::<(), _, _>(cfg.zset_key(), None, None, false, false, (score, member))
.await;
match (scheduled, cfg.ttl_millis()) {
(Err(err), _) => (0, Err(err)),
(Ok(()), None) => (1, Ok(())),
(Ok(()), Some(ttl)) => {
match sink.pexpire::<(), _>(cfg.zset_key(), ttl, None).await {
Ok(()) => (2, Ok(())),
Err(err) => (1, Err(err)),
}
}
}
}
}
}
}
impl Form for StreamForm {
type Message = RedisMessage;
type Op = StreamOp;
fn pool(msg: &RedisMessage) -> &Pool {
msg.seeker().pool()
}
#[cfg(feature = "testing")]
fn take_flight(msg: &mut RedisMessage) -> InFlight {
msg.take_flight()
}
fn ack(&self, msg: RedisMessage, ops: &mut Vec<StreamOp>) -> Result<(), AckError> {
let (id, _, _) = msg.into_settle();
ops.push(StreamOp::Ack(id));
Ok(())
}
fn nack(
&self,
msg: RedisMessage,
requeue: bool,
ops: &mut Vec<StreamOp>,
) -> Result<(), AckError> {
if requeue && let RequeueMode::LeavePending { .. } = self.requeue {
drop(msg.into_settle());
return Ok(());
}
let (id, payload, headers) = msg.into_settle();
if requeue {
ops.push(StreamOp::Republish(fields_for_publish(
payload.to_vec(),
&headers,
)));
}
ops.push(StreamOp::Ack(id));
Ok(())
}
fn nack_after(
&self,
msg: RedisMessage,
delay_by: Duration,
ops: &mut Vec<StreamOp>,
) -> Result<(), AckError> {
let Some(cfg) = self.delay.as_ref() else {
if let RequeueMode::LeavePending { .. } = self.requeue {
drop(msg.into_settle());
return Ok(());
}
return Err(AckError::Unsupported);
};
let (id, payload, headers) = msg.into_settle();
let (score, member) = cfg.scheduled(&id, &payload, &headers, delay_by);
ops.push(StreamOp::Schedule { score, member });
ops.push(StreamOp::Ack(id));
Ok(())
}
async fn close_atomic(&self, segment: &Segment, ops: &mut Vec<StreamOp>) -> Result<(), Error> {
for op in ops.drain(..) {
match segment {
Segment::Plain(pipeline) => self.queue(pipeline, op).await?,
Segment::Atomic(pinned) => self.queue(pinned, op).await?,
}
}
Ok(())
}
async fn send_ops(&self, client: &Client, ops: &mut Vec<StreamOp>) -> Sent {
let mut sent = Sent::default();
if ops.is_empty() {
return sent;
}
let mut writes = Writes::default();
let mut acked = Vec::new();
for op in ops.drain(..) {
match op {
StreamOp::Ack(id) => {
if let Some(id) = writes.close(id) {
acked.push(id);
}
}
op => {
let (commands, queued) = self.queue_counted(writes.pipeline(client), op).await;
writes.queued(commands, queued, &mut sent);
}
}
}
writes.send(&mut acked, &mut sent).await;
let settles = client.pipeline();
if !acked.is_empty()
&& let Err(err) = settles
.xack::<(), _, _, _>(&*self.key, &*self.group, acked)
.await
{
sent.failed += 1;
sent.first_error.get_or_insert(err);
}
sent.record(settles.try_all::<Value>().await);
sent
}
}
#[derive(Debug)]
pub struct ListOp {
value: Vec<u8>,
member: Option<Vec<u8>>,
requeue: bool,
}
#[derive(Debug)]
pub struct ListForm {
main: Arc<str>,
processing: Arc<str>,
recovery_zset: Option<Arc<str>>,
}
impl ListForm {
pub(crate) fn new(
main: impl Into<Arc<str>>,
processing: impl Into<Arc<str>>,
recovery_zset: Option<&str>,
) -> Self {
Self {
main: main.into(),
processing: processing.into(),
recovery_zset: recovery_zset.map(Arc::from),
}
}
fn settle(msg: RedisListMessage, requeue: bool, ops: &mut Vec<ListOp>) -> Result<(), AckError> {
let Some((value, member)) = msg.into_settle() else {
return Err(AckError::Unsupported);
};
ops.push(ListOp {
value,
member,
requeue,
});
Ok(())
}
async fn queue<Sink>(&self, sink: &Sink, op: ListOp) -> Result<(), Error>
where
Sink: ListInterface + SortedSetsInterface + Sync,
{
if op.requeue {
self.queue_requeue(sink, &op).await?;
}
self.queue_release(sink, op).await
}
async fn queue_requeue<Sink>(&self, sink: &Sink, op: &ListOp) -> Result<(), Error>
where
Sink: ListInterface + Sync,
{
sink.lpush::<(), _, _>(&*self.main, op.value.clone()).await
}
async fn queue_release<Sink>(&self, sink: &Sink, op: ListOp) -> Result<(), Error>
where
Sink: ListInterface + SortedSetsInterface + Sync,
{
sink.lrem::<(), _, _>(&*self.processing, 1, op.value)
.await?;
if let (Some(zset), Some(member)) = (&self.recovery_zset, op.member) {
sink.zrem::<(), _, _>(&**zset, member).await?;
}
Ok(())
}
}
impl Form for ListForm {
type Message = RedisListMessage;
type Op = ListOp;
fn pool(msg: &RedisListMessage) -> &Pool {
msg.pool()
}
#[cfg(feature = "testing")]
fn take_flight(msg: &mut RedisListMessage) -> InFlight {
msg.take_flight()
}
fn ack(&self, msg: RedisListMessage, ops: &mut Vec<ListOp>) -> Result<(), AckError> {
Self::settle(msg, false, ops)
}
fn nack(
&self,
msg: RedisListMessage,
requeue: bool,
ops: &mut Vec<ListOp>,
) -> Result<(), AckError> {
Self::settle(msg, requeue, ops)
}
fn nack_after(
&self,
msg: RedisListMessage,
_delay: Duration,
_ops: &mut Vec<ListOp>,
) -> Result<(), AckError> {
drop(msg);
Err(AckError::Unsupported)
}
async fn close_atomic(&self, segment: &Segment, ops: &mut Vec<ListOp>) -> Result<(), Error> {
for op in ops.drain(..) {
match segment {
Segment::Plain(pipeline) => self.queue(pipeline, op).await?,
Segment::Atomic(pinned) => self.queue(pinned, op).await?,
}
}
Ok(())
}
async fn send_ops(&self, client: &Client, ops: &mut Vec<ListOp>) -> Sent {
let mut sent = Sent::default();
if ops.is_empty() {
return sent;
}
let mut writes = Writes::default();
let mut released = Vec::new();
for op in ops.drain(..) {
if op.requeue {
let queued = self.queue_requeue(writes.pipeline(client), &op).await;
writes.queued(usize::from(queued.is_ok()), queued, &mut sent);
}
if let Some(op) = writes.close(op) {
released.push(op);
}
}
writes.send(&mut released, &mut sent).await;
let settles = client.pipeline();
for op in released {
if let Err(err) = self.queue_release(&settles, op).await {
sent.failed += 1;
sent.first_error.get_or_insert(err);
}
}
sent.record(settles.try_all::<Value>().await);
sent
}
}
struct Writes<Settle> {
pipeline: Option<Pipeline<Client>>,
commands: usize,
wrote: bool,
refused: bool,
owed: Vec<(usize, bool, Settle)>,
}
impl<Settle> Default for Writes<Settle> {
fn default() -> Self {
Self {
pipeline: None,
commands: 0,
wrote: false,
refused: false,
owed: Vec::new(),
}
}
}
impl<Settle> Writes<Settle> {
fn pipeline(&mut self, client: &Client) -> &Pipeline<Client> {
self.pipeline.get_or_insert_with(|| client.pipeline())
}
fn queued(&mut self, commands: usize, queued: Result<(), Error>, sent: &mut Sent) {
self.wrote = true;
self.commands += commands;
match queued {
Ok(()) => {}
Err(err) => {
self.refused = true;
sent.failed += 1;
sent.first_error.get_or_insert(err);
}
}
}
fn close(&mut self, settle: Settle) -> Option<Settle> {
if !std::mem::take(&mut self.wrote) {
return Some(settle);
}
let commands = std::mem::take(&mut self.commands);
let refused = std::mem::take(&mut self.refused);
self.owed.push((commands, refused, settle));
None
}
async fn send(self, due: &mut Vec<Settle>, sent: &mut Sent) {
let Some(pipeline) = self.pipeline else {
return;
};
let results = pipeline.try_all::<Value>().await;
let expected: usize = self.owed.iter().map(|(commands, _, _)| commands).sum();
if results.len() != expected {
sent.record(results);
return;
}
sent.commands += results.len();
let mut results = results.into_iter();
for (commands, refused, settle) in self.owed {
let mut done = !refused;
for result in results.by_ref().take(commands) {
if let Err(err) = result {
done = false;
sent.failed += 1;
sent.first_error.get_or_insert(err);
}
}
if done {
due.push(settle);
}
}
}
}
#[derive(Debug, Default)]
pub struct PubSubForm;
impl Form for PubSubForm {
type Message = RedisPubSubMessage;
type Op = std::convert::Infallible;
fn pool(msg: &RedisPubSubMessage) -> &Pool {
msg.pool()
}
#[cfg(feature = "testing")]
fn take_flight(msg: &mut RedisPubSubMessage) -> InFlight {
msg.take_flight()
}
fn ack(&self, _msg: RedisPubSubMessage, _ops: &mut Vec<Self::Op>) -> Result<(), AckError> {
Err(AckError::Unsupported)
}
fn nack(
&self,
_msg: RedisPubSubMessage,
_requeue: bool,
_ops: &mut Vec<Self::Op>,
) -> Result<(), AckError> {
Err(AckError::Unsupported)
}
fn nack_after(
&self,
_msg: RedisPubSubMessage,
_delay: Duration,
_ops: &mut Vec<Self::Op>,
) -> Result<(), AckError> {
Err(AckError::Unsupported)
}
fn close_atomic(
&self,
_segment: &Segment,
_ops: &mut Vec<Self::Op>,
) -> impl Future<Output = Result<(), Error>> + Send {
ready(Ok(()))
}
fn send_ops(
&self,
_client: &Client,
_ops: &mut Vec<Self::Op>,
) -> impl Future<Output = Sent> + Send {
ready(Sent::default())
}
}