#![macro_use]
#[cfg(feature = "cache-aio")]
use crate::cmd::CommandCacheConfig;
use crate::cmd::{Arg, Cmd, CmdRef, args_len, write_command};
use crate::connection::ConnectionLike;
use crate::errors::ErrorKind;
use crate::types::{FromRedisValue, RedisResult, ToRedisArgs, Value, from_redis_value};
#[derive(Clone, Debug)]
pub(crate) struct CommandRecord {
pub(crate) args_start: usize,
pub(crate) data_start: usize,
pub(crate) cursor: Option<u64>,
pub(crate) ignored: bool,
pub(crate) no_response: bool,
#[cfg(feature = "cache-aio")]
pub(crate) cache: Option<CommandCacheConfig>,
}
#[derive(Clone, Debug)]
pub struct Pipeline {
pub(crate) data: Vec<u8>,
pub(crate) args: Vec<Arg<usize>>,
pub(crate) commands: Vec<CommandRecord>,
pub(crate) transaction_mode: bool,
pub(crate) ignore_errors: bool,
#[cfg(feature = "aio")]
pub(crate) subscription_confirmations: Option<ExpectedConfirmations>,
}
#[cfg(feature = "aio")]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ExpectedConfirmations {
Count(usize),
UnsubscribeAllChannels,
UnsubscribeAllPatterns,
#[cfg(feature = "cluster-async")]
UnsubscribeAllShardChannels,
}
impl Pipeline {
pub fn new() -> Self {
Self {
data: Vec::new(),
args: Vec::new(),
commands: Vec::new(),
transaction_mode: false,
ignore_errors: false,
#[cfg(feature = "aio")]
subscription_confirmations: None,
}
}
#[inline]
pub fn atomic(&mut self) -> &mut Self {
self.transaction_mode = true;
self
}
#[inline]
pub fn ignore_errors(&mut self) -> &mut Self {
self.ignore_errors = true;
self
}
pub fn is_transaction(&self) -> bool {
self.transaction_mode
}
pub fn get_packed_pipeline(&self) -> Vec<u8> {
encode_pipeline(self, self.transaction_mode)
}
fn command_parts(&self, idx: usize) -> (&[Arg<usize>], usize, u64) {
let record = &self.commands[idx];
let args_end = self
.commands
.get(idx + 1)
.map_or(self.args.len(), |r| r.args_start);
(
&self.args[record.args_start..args_end],
record.data_start,
record.cursor.unwrap_or(0),
)
}
pub fn len(&self) -> usize {
self.commands.len()
}
pub fn is_empty(&self) -> bool {
self.commands.is_empty()
}
#[inline]
pub fn query<T: FromRedisValue>(&self, con: &mut dyn ConnectionLike) -> RedisResult<T> {
if !con.supports_pipelining() {
fail!((
ErrorKind::Client,
"This connection does not support pipelining."
));
}
let response = if self.transaction_mode {
con.req_packed_commands(&encode_pipeline(self, true), self.len() + 1, 1)?
} else {
con.req_packed_commands(&encode_pipeline(self, false), 0, self.len())?
};
self.complete_request(response)
}
#[inline]
#[cfg(feature = "aio")]
pub async fn query_async<T: FromRedisValue>(
&self,
con: &mut impl crate::aio::ConnectionLike,
) -> RedisResult<T> {
let response = if self.transaction_mode {
con.req_packed_commands(self, self.len() + 1, 1).await?
} else {
con.req_packed_commands(self, 0, self.len()).await?
};
self.complete_request(response)
}
#[inline]
pub fn exec(&self, con: &mut dyn ConnectionLike) -> RedisResult<()> {
self.query::<()>(con)
}
#[cfg(feature = "aio")]
pub async fn exec_async(&self, con: &mut impl crate::aio::ConnectionLike) -> RedisResult<()> {
self.query_async::<()>(con).await
}
fn complete_request<T: FromRedisValue>(&self, mut response: Vec<Value>) -> RedisResult<T> {
let response = if self.is_transaction() {
match response.pop() {
Some(Value::Nil) => {
return Ok(from_redis_value(Value::Nil)?);
}
Some(Value::Array(items)) => items,
_ => {
return Err((
ErrorKind::UnexpectedReturnType,
"Invalid response when parsing multi response",
)
.into());
}
}
} else {
response
};
self.compose_response(response)
}
}
fn encode_pipeline(pipe: &Pipeline, atomic: bool) -> Vec<u8> {
let mut rv = vec![];
write_pipeline(&mut rv, pipe, atomic);
rv
}
const PACKED_MULTI: &[u8] = b"*1\r\n$5\r\nMULTI\r\n";
const PACKED_EXEC: &[u8] = b"*1\r\n$4\r\nEXEC\r\n";
fn command_args_iter<'a>(
data: &'a [u8],
args: &'a [Arg<usize>],
data_start: usize,
) -> impl Clone + ExactSizeIterator<Item = Arg<&'a [u8]>> {
let mut prev = data_start;
args.iter().map(move |arg| match *arg {
Arg::Simple(end) => {
let arg = Arg::Simple(&data[prev..end]);
prev = end;
arg
}
Arg::Cursor => Arg::Cursor,
})
}
fn write_pipeline(rv: &mut Vec<u8>, pipe: &Pipeline, atomic: bool) {
let cmds_len: usize = (0..pipe.commands.len())
.map(|idx| {
let (args, data_start, cursor) = pipe.command_parts(idx);
args_len(command_args_iter(&pipe.data, args, data_start), cursor)
})
.sum();
let write_commands = |rv: &mut Vec<u8>| {
for idx in 0..pipe.commands.len() {
let (args, data_start, cursor) = pipe.command_parts(idx);
write_command(rv, command_args_iter(&pipe.data, args, data_start), cursor);
}
};
if atomic {
rv.reserve(PACKED_MULTI.len() + PACKED_EXEC.len() + cmds_len);
rv.extend_from_slice(PACKED_MULTI);
write_commands(rv);
rv.extend_from_slice(PACKED_EXEC);
} else {
rv.reserve(cmds_len);
write_commands(rv);
}
}
macro_rules! implement_pipeline_commands {
($struct_name:ident) => {
impl $struct_name {
#[inline]
pub(crate) fn start_command(&mut self) {
self.commands.push($crate::pipeline::CommandRecord {
args_start: self.args.len(),
data_start: self.data.len(),
cursor: None,
ignored: false,
no_response: false,
#[cfg(feature = "cache-aio")]
cache: None,
});
}
#[inline]
pub fn add_command(&mut self, cmd: Cmd) -> &mut Self {
self.start_command();
for arg in cmd.args_iter() {
match arg {
Arg::Simple(bytes) => $crate::types::RedisWrite::write_arg(self, bytes),
Arg::Cursor => self.args.push(Arg::Cursor),
}
}
if let Some(record) = self.commands.last_mut() {
record.cursor = cmd.cursor_value();
record.no_response = cmd.is_no_response();
#[cfg(feature = "cache-aio")]
{
record.cache = cmd.get_cache_config().clone();
}
}
self
}
#[inline]
pub fn cmd(&mut self, name: &str) -> &mut Self {
self.start_command();
$crate::types::RedisWrite::write_arg(self, name.as_bytes());
self
}
pub fn cmd_iter(&self) -> impl Iterator<Item = CmdRef<'_>> {
(0..self.commands.len()).map(move |idx| self.command_ref(idx))
}
#[inline]
pub fn reserve_for_commands(&mut self, additional: usize) -> &mut Self {
self.commands.reserve(additional);
self
}
#[inline]
pub fn reserve_for_args(&mut self, additional: usize) -> &mut Self {
self.args.reserve(additional);
self
}
#[inline]
pub fn reserve_for_data(&mut self, additional: usize) -> &mut Self {
self.data.reserve(additional);
self
}
#[inline]
fn command_ref(&self, idx: usize) -> CmdRef<'_> {
let record = &self.commands[idx];
let next = self.commands.get(idx + 1);
let args_end = next.map_or(self.args.len(), |r| r.args_start);
let data_end = next.map_or(self.data.len(), |r| r.data_start);
CmdRef::new(
&self.data[record.data_start..data_end],
&self.args[record.args_start..args_end],
record.data_start,
record.cursor,
record.no_response,
record.ignored,
#[cfg(feature = "cache-aio")]
&record.cache,
)
}
#[inline]
pub fn ignore(&mut self) -> &mut Self {
if let Some(record) = self.commands.last_mut() {
record.ignored = true;
}
self
}
#[inline]
pub fn arg<T: ToRedisArgs>(&mut self, arg: T) -> &mut Self {
assert!(!self.commands.is_empty(), "No command on stack");
arg.write_redis_args(self);
self
}
#[inline]
pub fn clear(&mut self) {
self.data.clear();
self.args.clear();
self.commands.clear();
}
fn filter_ignored_results(&self, resp: Vec<Value>) -> Vec<Value> {
if !self.commands.iter().any(|command| command.ignored) {
return resp;
}
resp.into_iter()
.enumerate()
.filter_map(|(index, result)| {
let ignored = self
.commands
.get(index)
.is_some_and(|record| record.ignored);
(!ignored).then_some(result)
})
.collect()
}
fn compose_response<T: FromRedisValue>(&self, response: Vec<Value>) -> RedisResult<T> {
if self.ignore_errors {
return Ok(from_redis_value(Value::Array(
self.filter_ignored_results(response),
))?);
}
let server_errors: Vec<_> = response
.iter()
.enumerate()
.filter_map(|(index, value)| match value {
Value::ServerError(error) => Some((index, error.clone())),
_ => None,
})
.collect();
if server_errors.is_empty() {
Ok(from_redis_value(
Value::Array(self.filter_ignored_results(response)).extract_error()?,
)?)
} else {
Err(crate::RedisError::pipeline(server_errors))
}
}
}
impl $crate::types::RedisWrite for $struct_name {
fn write_arg(&mut self, arg: &[u8]) {
self.data.extend_from_slice(arg);
self.args.push(Arg::Simple(self.data.len()));
}
fn write_arg_fmt(&mut self, arg: impl std::fmt::Display) {
use std::io::Write as _;
write!(self.data, "{arg}").unwrap();
self.args.push(Arg::Simple(self.data.len()));
}
fn writer_for_next_arg(&mut self) -> impl std::io::Write + '_ {
struct PipelineArgGuard<'a>(&'a mut $struct_name);
impl Drop for PipelineArgGuard<'_> {
fn drop(&mut self) {
self.0.args.push(Arg::Simple(self.0.data.len()));
}
}
impl std::io::Write for PipelineArgGuard<'_> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0.data.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
PipelineArgGuard(self)
}
fn reserve_space_for_args(&mut self, additional: impl IntoIterator<Item = usize>) {
let mut capacity = 0;
let mut args = 0;
for add in additional {
capacity += add;
args += 1;
}
self.data.reserve(capacity);
self.args.reserve(args);
}
}
impl Default for $struct_name {
fn default() -> Self {
Self::new()
}
}
};
}
implement_pipeline_commands!(Pipeline);
impl Pipeline {
#[cfg(feature = "cache-aio")]
#[cfg_attr(docsrs, doc(cfg(feature = "cache-aio")))]
pub fn set_cache_config(&mut self, command_cache_config: CommandCacheConfig) -> &mut Self {
if let Some(record) = self.commands.last_mut() {
record.cache = Some(command_cache_config);
}
self
}
#[cfg(feature = "cache-aio")]
pub(crate) fn add_command_ref(&mut self, cmd: CmdRef<'_>) -> &mut Self {
self.start_command();
for arg in cmd.args_iter() {
match arg {
Arg::Simple(bytes) => crate::types::RedisWrite::write_arg(self, bytes),
Arg::Cursor => self.args.push(Arg::Cursor),
}
}
if let Some(record) = self.commands.last_mut() {
record.cursor = cmd.cursor();
record.no_response = cmd.is_no_response();
record.cache = cmd.get_cache_config().clone();
}
self
}
#[cfg(feature = "cluster-async")]
pub(crate) fn into_cmd_iter(self) -> impl Iterator<Item = Cmd> {
self.cmd_iter()
.map(|cmd| cmd.to_cmd())
.collect::<Vec<_>>()
.into_iter()
}
}
#[cfg(test)]
mod tests {
use crate::{
ServerErrorKind,
errors::{RedisError, Repr, ServerError},
pipe,
};
use super::*;
fn test_pipe() -> Pipeline {
let mut pipeline = pipe();
pipeline
.cmd("FOO")
.cmd("BAR")
.ignore()
.cmd("baz")
.ignore()
.cmd("barvaz");
pipeline
}
fn server_error() -> Value {
Value::ServerError(ServerError(Repr::Known {
kind: ServerErrorKind::CrossSlot,
detail: None,
}))
}
#[test]
fn test_pipeline_passes_values_only_from_non_ignored_commands() {
let pipeline = test_pipe();
let inputs = vec![Value::Int(1), Value::Int(2), Value::Int(3), Value::Okay];
let result = pipeline.complete_request(inputs);
let expected = vec!["1".to_string(), "OK".to_string()];
assert_eq!(result, Ok(expected));
}
#[test]
fn test_pipeline_passes_errors_from_ignored_commands() {
let pipeline = test_pipe();
let inputs = vec![Value::Okay, server_error(), Value::Okay, server_error()];
let error = pipeline.compose_response::<Vec<Value>>(inputs).unwrap_err();
let error_message = error.to_string();
assert!(error_message.contains("Index 1"), "{error_message}");
assert!(error_message.contains("Index 3"), "{error_message}");
}
#[test]
fn test_pipeline_passes_response_from_each_commands_despite_errors() {
let mut pipeline = test_pipe();
let inputs = vec![Value::Okay, server_error(), Value::Okay, server_error()];
let result = pipeline
.ignore_errors()
.compose_response::<Vec<RedisResult<Value>>>(inputs)
.unwrap();
assert_eq!(result[0], Ok::<Value, RedisError>(Value::Okay));
assert!(
result[1]
.clone()
.is_err_and(|e| e.to_string().contains("CrossSlot"))
);
}
}