use std::{fmt, io};
use crate::connection::ConnectionLike;
use crate::types::{
from_redis_value, ErrorKind, FromRedisValue, RedisResult, RedisWrite, ToRedisArgs, Value,
};
#[derive(Clone)]
pub enum Arg<D> {
Simple(D),
Cursor,
}
#[derive(Clone)]
pub struct Cmd {
data: Vec<u8>,
args: Vec<Arg<usize>>,
cursor: Option<u64>,
is_ignored: bool,
}
#[derive(Clone)]
pub struct Pipeline {
commands: Vec<Cmd>,
transaction_mode: bool,
}
pub struct Iter<'a, T: FromRedisValue> {
batch: std::vec::IntoIter<T>,
cursor: u64,
con: &'a mut (dyn ConnectionLike + 'a),
cmd: Cmd,
}
impl<'a, T: FromRedisValue> Iterator for Iter<'a, T> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<T> {
loop {
if let Some(v) = self.batch.next() {
return Some(v);
};
if self.cursor == 0 {
return None;
}
let pcmd = unwrap_or!(
self.cmd.get_packed_command_with_cursor(self.cursor),
return None
);
let rv = unwrap_or!(self.con.req_packed_command(&pcmd).ok(), return None);
let (cur, batch): (u64, Vec<T>) = unwrap_or!(from_redis_value(&rv).ok(), return None);
self.cursor = cur;
self.batch = batch.into_iter();
}
}
}
#[cfg(feature = "aio")]
use crate::aio::ConnectionLike as AsyncConnection;
#[cfg(feature = "aio")]
pub struct AsyncIter<'a, T: FromRedisValue + 'a> {
batch: std::vec::IntoIter<T>,
con: &'a mut (dyn AsyncConnection + Send + 'a),
cmd: Cmd,
}
#[cfg(feature = "aio")]
impl<'a, T: FromRedisValue + 'a> AsyncIter<'a, T> {
#[inline]
pub async fn next_item(&mut self) -> Option<T> {
loop {
if let Some(v) = self.batch.next() {
return Some(v);
};
if let Some(cursor) = self.cmd.cursor {
if cursor == 0 {
return None;
}
}
let rv = unwrap_or!(
self.con.req_packed_command(&self.cmd).await.ok(),
return None
);
let (cur, batch): (u64, Vec<T>) = unwrap_or!(from_redis_value(&rv).ok(), return None);
self.cmd.cursor = Some(cur);
self.batch = batch.into_iter();
}
}
}
fn countdigits(mut v: usize) -> usize {
let mut result = 1;
loop {
if v < 10 {
return result;
}
if v < 100 {
return result + 1;
}
if v < 1000 {
return result + 2;
}
if v < 10000 {
return result + 3;
}
v /= 10000;
result += 4;
}
}
#[inline]
fn bulklen(len: usize) -> usize {
1 + countdigits(len) + 2 + len + 2
}
fn args_len<'a, I>(args: I, cursor: u64) -> usize
where
I: IntoIterator<Item = Arg<&'a [u8]>> + ExactSizeIterator,
{
let mut totlen = 1 + countdigits(args.len()) + 2;
for item in args {
totlen += bulklen(match item {
Arg::Cursor => countdigits(cursor as usize),
Arg::Simple(val) => val.len(),
});
}
totlen
}
fn cmd_len(cmd: &Cmd) -> usize {
args_len(cmd.args_iter(), cmd.cursor.unwrap_or(0))
}
fn encode_command<'a, I>(args: I, cursor: u64) -> Vec<u8>
where
I: IntoIterator<Item = Arg<&'a [u8]>> + Clone + ExactSizeIterator,
{
let mut cmd = Vec::new();
write_command_to_vec(&mut cmd, args, cursor);
cmd
}
fn write_command_to_vec<'a, I>(cmd: &mut Vec<u8>, args: I, cursor: u64)
where
I: IntoIterator<Item = Arg<&'a [u8]>> + Clone + ExactSizeIterator,
{
let totlen = args_len(args.clone(), cursor);
cmd.reserve(totlen);
write_command(cmd, args, cursor).unwrap()
}
fn write_command<'a, I>(cmd: &mut (impl ?Sized + io::Write), args: I, cursor: u64) -> io::Result<()>
where
I: IntoIterator<Item = Arg<&'a [u8]>> + Clone + ExactSizeIterator,
{
cmd.write_all(b"*")?;
::itoa::write(&mut *cmd, args.len())?;
cmd.write_all(b"\r\n")?;
let mut cursor_bytes = itoa::Buffer::new();
for item in args {
let bytes = match item {
Arg::Cursor => cursor_bytes.format(cursor).as_bytes(),
Arg::Simple(val) => val,
};
cmd.write_all(b"$")?;
::itoa::write(&mut *cmd, bytes.len())?;
cmd.write_all(b"\r\n")?;
cmd.write_all(bytes)?;
cmd.write_all(b"\r\n")?;
}
Ok(())
}
fn encode_pipeline(cmds: &[Cmd], atomic: bool) -> Vec<u8> {
let mut rv = vec![];
write_pipeline(&mut rv, cmds, atomic);
rv
}
fn write_pipeline(rv: &mut Vec<u8>, cmds: &[Cmd], atomic: bool) {
let cmds_len = cmds.iter().map(cmd_len).sum();
if atomic {
let multi = cmd("MULTI");
let exec = cmd("EXEC");
rv.reserve(cmd_len(&multi) + cmd_len(&exec) + cmds_len);
multi.write_packed_command_preallocated(rv);
for cmd in cmds {
cmd.write_packed_command_preallocated(rv);
}
exec.write_packed_command_preallocated(rv);
} else {
rv.reserve(cmds_len);
for cmd in cmds {
cmd.write_packed_command_preallocated(rv);
}
}
}
impl RedisWrite for Cmd {
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 fmt::Display) {
use std::io::Write;
write!(self.data, "{}", arg).unwrap();
self.args.push(Arg::Simple(self.data.len()));
}
}
impl Default for Cmd {
fn default() -> Cmd {
Cmd::new()
}
}
impl Cmd {
pub fn new() -> Cmd {
Cmd {
data: vec![],
args: vec![],
cursor: None,
is_ignored: false,
}
}
#[inline]
pub fn arg<T: ToRedisArgs>(&mut self, arg: T) -> &mut Cmd {
arg.write_redis_args(self);
self
}
#[inline]
pub fn cursor_arg(&mut self, cursor: u64) -> &mut Cmd {
assert!(!self.in_scan_mode());
self.cursor = Some(cursor);
self.args.push(Arg::Cursor);
self
}
#[inline]
pub fn get_packed_command(&self) -> Vec<u8> {
let mut cmd = Vec::new();
self.write_packed_command(&mut cmd);
cmd
}
pub(crate) fn write_packed_command(&self, cmd: &mut Vec<u8>) {
write_command_to_vec(cmd, self.args_iter(), self.cursor.unwrap_or(0))
}
fn write_packed_command_preallocated(&self, cmd: &mut Vec<u8>) {
write_command(cmd, self.args_iter(), self.cursor.unwrap_or(0)).unwrap()
}
#[inline]
fn get_packed_command_with_cursor(&self, cursor: u64) -> Option<Vec<u8>> {
if !self.in_scan_mode() {
None
} else {
Some(encode_command(self.args_iter(), cursor))
}
}
#[inline]
pub fn in_scan_mode(&self) -> bool {
self.cursor.is_some()
}
#[inline]
pub fn query<T: FromRedisValue>(&self, con: &mut dyn ConnectionLike) -> RedisResult<T> {
let pcmd = self.get_packed_command();
match con.req_packed_command(&pcmd) {
Ok(val) => from_redis_value(&val),
Err(e) => Err(e),
}
}
#[inline]
#[cfg(feature = "aio")]
pub async fn query_async<C, T: FromRedisValue>(&self, con: &mut C) -> RedisResult<T>
where
C: crate::aio::ConnectionLike,
{
let val = con.req_packed_command(self).await?;
from_redis_value(&val)
}
#[inline]
pub fn iter<'a, T: FromRedisValue>(
self,
con: &'a mut dyn ConnectionLike,
) -> RedisResult<Iter<'a, T>> {
let pcmd = self.get_packed_command();
let rv = con.req_packed_command(&pcmd)?;
let (cursor, batch) = if rv.looks_like_cursor() {
from_redis_value::<(u64, Vec<T>)>(&rv)?
} else {
(0, from_redis_value(&rv)?)
};
Ok(Iter {
batch: batch.into_iter(),
cursor,
con,
cmd: self,
})
}
#[cfg(feature = "aio")]
#[inline]
pub async fn iter_async<'a, T: FromRedisValue + 'a>(
mut self,
con: &'a mut (dyn AsyncConnection + Send),
) -> RedisResult<AsyncIter<'a, T>> {
let rv = con.req_packed_command(&self).await?;
let (cursor, batch) = if rv.looks_like_cursor() {
from_redis_value::<(u64, Vec<T>)>(&rv)?
} else {
(0, from_redis_value(&rv)?)
};
if cursor == 0 {
self.cursor = None;
} else {
self.cursor = Some(cursor);
}
Ok(AsyncIter {
batch: batch.into_iter(),
con,
cmd: self,
})
}
#[inline]
pub fn execute(&self, con: &mut dyn ConnectionLike) {
self.query::<()>(con).unwrap();
}
pub fn args_iter(&self) -> impl Iterator<Item = Arg<&[u8]>> + Clone + ExactSizeIterator {
let mut prev = 0;
self.args.iter().map(move |arg| match *arg {
Arg::Simple(i) => {
let arg = Arg::Simple(&self.data[prev..i]);
prev = i;
arg
}
Arg::Cursor => Arg::Cursor,
})
}
}
impl Default for Pipeline {
fn default() -> Pipeline {
Pipeline::new()
}
}
impl Pipeline {
pub fn new() -> Pipeline {
Self::with_capacity(0)
}
pub fn with_capacity(capacity: usize) -> Pipeline {
Pipeline {
commands: Vec::with_capacity(capacity),
transaction_mode: false,
}
}
#[inline]
pub fn cmd(&mut self, name: &str) -> &mut Pipeline {
self.commands.push(cmd(name));
self
}
#[inline]
pub fn add_command(&mut self, cmd: Cmd) -> &mut Pipeline {
self.commands.push(cmd);
self
}
#[inline]
fn get_last_command(&mut self) -> &mut Cmd {
let idx = match self.commands.len() {
0 => panic!("No command on stack"),
x => x - 1,
};
&mut self.commands[idx]
}
#[inline]
pub fn arg<T: ToRedisArgs>(&mut self, arg: T) -> &mut Pipeline {
{
let cmd = self.get_last_command();
cmd.arg(arg);
}
self
}
#[inline]
pub fn ignore(&mut self) -> &mut Pipeline {
{
let cmd = self.get_last_command();
cmd.is_ignored = true;
}
self
}
#[inline]
pub fn atomic(&mut self) -> &mut Pipeline {
self.transaction_mode = true;
self
}
pub fn cmd_iter(&self) -> impl Iterator<Item = &Cmd> {
self.commands.iter()
}
fn make_pipeline_results(&self, resp: Vec<Value>) -> Value {
let mut rv = vec![];
for (idx, result) in resp.into_iter().enumerate() {
if !self.commands[idx].is_ignored {
rv.push(result);
}
}
Value::Bulk(rv)
}
pub fn get_packed_pipeline(&self) -> Vec<u8> {
encode_pipeline(&self.commands, self.transaction_mode)
}
#[cfg(feature = "aio")]
pub(crate) fn write_packed_pipeline(&self, out: &mut Vec<u8>) {
write_pipeline(out, &self.commands, self.transaction_mode)
}
fn execute_pipelined(&self, con: &mut dyn ConnectionLike) -> RedisResult<Value> {
Ok(self.make_pipeline_results(con.req_packed_commands(
&encode_pipeline(&self.commands, false),
0,
self.commands.len(),
)?))
}
fn execute_transaction(&self, con: &mut dyn ConnectionLike) -> RedisResult<Value> {
let mut resp = con.req_packed_commands(
&encode_pipeline(&self.commands, true),
self.commands.len() + 1,
1,
)?;
match resp.pop() {
Some(Value::Nil) => Ok(Value::Nil),
Some(Value::Bulk(items)) => Ok(self.make_pipeline_results(items)),
_ => fail!((
ErrorKind::ResponseError,
"Invalid response when parsing multi response"
)),
}
}
#[inline]
pub fn query<T: FromRedisValue>(&self, con: &mut dyn ConnectionLike) -> RedisResult<T> {
if !con.supports_pipelining() {
fail!((
ErrorKind::ResponseError,
"This connection does not support pipelining."
));
}
from_redis_value(
&(if self.commands.is_empty() {
Value::Bulk(vec![])
} else if self.transaction_mode {
self.execute_transaction(con)?
} else {
self.execute_pipelined(con)?
}),
)
}
#[inline]
pub fn clear(&mut self) {
self.commands.clear();
}
#[cfg(feature = "aio")]
async fn execute_pipelined_async<C>(&self, con: &mut C) -> RedisResult<Value>
where
C: crate::aio::ConnectionLike,
{
let value = con
.req_packed_commands(self, 0, self.commands.len())
.await?;
Ok(self.make_pipeline_results(value))
}
#[cfg(feature = "aio")]
async fn execute_transaction_async<C>(&self, con: &mut C) -> RedisResult<Value>
where
C: crate::aio::ConnectionLike,
{
let mut resp = con
.req_packed_commands(self, self.commands.len() + 1, 1)
.await?;
match resp.pop() {
Some(Value::Nil) => Ok(Value::Nil),
Some(Value::Bulk(items)) => Ok(self.make_pipeline_results(items)),
_ => Err((
ErrorKind::ResponseError,
"Invalid response when parsing multi response",
)
.into()),
}
}
#[inline]
#[cfg(feature = "aio")]
pub async fn query_async<C, T: FromRedisValue>(&self, con: &mut C) -> RedisResult<T>
where
C: crate::aio::ConnectionLike,
{
let v = if self.commands.is_empty() {
return from_redis_value(&Value::Bulk(vec![]));
} else if self.transaction_mode {
self.execute_transaction_async(con).await?
} else {
self.execute_pipelined_async(con).await?
};
from_redis_value(&v)
}
#[inline]
pub fn execute(&self, con: &mut dyn ConnectionLike) {
self.query::<()>(con).unwrap();
}
}
pub fn cmd(name: &str) -> Cmd {
let mut rv = Cmd::new();
rv.arg(name);
rv
}
pub fn pack_command(args: &[Vec<u8>]) -> Vec<u8> {
encode_command(args.iter().map(|x| Arg::Simple(&x[..])), 0)
}
pub fn pipe() -> Pipeline {
Pipeline::new()
}