redis_cluster/
cmd.rs

1use crc16::key_hash_slot;
2use redis::{Cmd, Connection, ToRedisArgs, FromRedisValue, RedisResult};
3
4/// Redis::Cmd's `args` field is private,
5/// but we need it to determine a slot from the command.
6/// So this is a simple wrapper around Redis::Cmd that keeps
7/// track of the args.
8pub struct ClusterCmd {
9    cmd: Cmd,
10    args: Vec<Vec<u8>>,
11}
12
13impl ClusterCmd {
14    pub fn new() -> ClusterCmd {
15        ClusterCmd {
16            cmd: Cmd::new(),
17            args: Vec::new(),
18        }
19    }
20
21    /// Add an arg to the command.
22    pub fn arg<T: ToRedisArgs>(&mut self, arg: T) -> &mut ClusterCmd {
23        for item in arg.to_redis_args().into_iter() {
24            self.args.push(item);
25        }
26        self.cmd.arg(arg);
27        self
28    }
29
30    /// Execute the command, returning the result.
31    pub fn query<T: FromRedisValue>(&self, conn: &Connection) -> RedisResult<T> {
32        self.cmd.query(conn)
33    }
34
35    /// Get the slot for this command.
36    pub fn slot(&self) -> Option<u16> {
37        slot_for_command(&self.args)
38    }
39}
40
41fn slot_for_command(args: &Vec<Vec<u8>>) -> Option<u16> {
42    if args.len() > 1 {
43        Some(key_hash_slot(args[1].as_slice()))
44    } else {
45        None
46    }
47}
48
49pub fn slot_for_packed_command(cmd: &[u8]) -> Option<u16> {
50    let args = unpack_command(cmd);
51    slot_for_command(&args)
52}
53
54/// `redis-rs` passes packed commands (as a u8 slice)
55/// to the methods of the Commands trait
56/// we need to "unpack" the command into the
57/// original arguments to properly compute
58/// the command's slot.
59/// This is pretty messy/can probably be better
60fn unpack_command(cmd: &[u8]) -> Vec<Vec<u8>> {
61    let mut arg: Vec<u8> = Vec::new();
62    let mut args: Vec<Vec<u8>> = Vec::new();
63
64    // first 4 are some leading info ('*', len of args, '\r', '\n')
65    // the next 4 precede the first arg
66    // see: <https://github.com/mitsuhiko/redis-rs/blob/master/src/cmd.rs#L85>
67    let mut iter = cmd.iter().skip(2).peekable();
68
69    'outer: loop {
70        let b = *iter.next().unwrap();
71
72        // args are separated by 13, 10
73        if b == 13 && iter.peek().unwrap() == &&10 {
74            if arg.len() > 0 {
75                args.push(arg.clone());
76                arg.clear();
77            }
78
79            // consume the next item (10)
80            iter.next();
81
82            // then, if there are more args, there should be a 36
83            // if there's nothing, we're done
84            match iter.next() {
85                Some(_) => (),
86                None => break 'outer,
87            };
88            // then the length of the args (which in theory can be any length)
89            // then another 13, 10
90            'inner: loop {
91                let b = *iter.next().unwrap();
92                if b == 13 && iter.peek().unwrap() == &&10 {
93                    iter.next();
94                    break 'inner;
95                }
96            }
97        } else {
98            arg.push(b);
99        }
100    }
101    args
102}
redis_cluster/cmd.rs

redis_cluster/
cmd.rs
