use std::{io, pin::Pin};
use crate::connection::ConnectionLike;
use crate::types::{
from_redis_value, ErrorKind, FromRedisValue, RedisResult, RedisWrite, ToRedisArgs, Value,
};
use tokio::io::{AsyncWrite, AsyncWriteExt};
#[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: Vec<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.pop() {
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, mut batch): (u64, Vec<T>) =
unwrap_or!(from_redis_value(&rv).ok(), return None);
batch.reverse();
self.cursor = cur;
self.batch = batch;
}
}
}
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(())
}
async fn write_command_async<'a, I>(
mut cmd: Pin<&mut (impl ?Sized + AsyncWrite)>,
args: I,
cursor: u64,
) -> io::Result<()>
where
I: IntoIterator<Item = Arg<&'a [u8]>> + Clone + ExactSizeIterator,
{
cmd.write_all(b"*").await?;
cmd.write_all(itoa::Buffer::new().format(args.len()).as_bytes())
.await?;
cmd.write_all(b"\r\n").await?;
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"$").await?;
cmd.write_all(itoa::Buffer::new().format(bytes.len()).as_bytes())
.await?;
cmd.write_all(b"\r\n").await?;
cmd.write_all(bytes).await?;
cmd.write_all(b"\r\n").await?;
}
Ok(())
}
fn encode_pipeline(cmds: &[Cmd], atomic: bool) -> Vec<u8> {
let mut rv = vec![];
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(&mut rv);
for cmd in cmds {
cmd.write_packed_command_preallocated(&mut rv);
}
exec.write_packed_command_preallocated(&mut rv);
} else {
rv.reserve(cmds_len);
for cmd in cmds {
cmd.write_packed_command_preallocated(&mut rv);
}
}
rv
}
async fn write_pipeline_async(
mut out: Pin<&mut (impl ?Sized + AsyncWrite)>,
cmds: &[Cmd],
atomic: bool,
) -> io::Result<()> {
if atomic {
cmd("MULTI").write_command_async(out.as_mut()).await?;
for cmd in cmds {
cmd.write_command_async(out.as_mut()).await?;
}
cmd("EXEC").write_command_async(out.as_mut()).await?;
} else {
for cmd in cmds {
cmd.write_command_async(out.as_mut()).await?;
}
}
Ok(())
}
impl RedisWrite for Cmd {
fn write_arg(&mut self, arg: &[u8]) {
let prev = self.data.len();
self.args.push(Arg::Simple(prev + arg.len()));
self.data.extend_from_slice(arg);
}
}
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) async fn write_command_async(
&self,
out: Pin<&mut (impl ?Sized + AsyncWrite)>,
) -> io::Result<()> {
write_command_async(out, self.args_iter(), self.cursor.unwrap_or(0)).await
}
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]
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, mut batch) = if rv.looks_like_cursor() {
from_redis_value::<(u64, Vec<T>)>(&rv)?
} else {
(0, from_redis_value(&rv)?)
};
batch.reverse();
Ok(Iter {
batch,
cursor,
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)
}
pub(crate) async fn write_pipeline_async(
&self,
out: Pin<&mut (impl ?Sized + AsyncWrite)>,
) -> io::Result<()> {
write_pipeline_async(out, &self.commands, self.transaction_mode).await
}
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();
}
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))
}
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]
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()
}