1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

use crc16::key_hash_slot;
use redis::{Cmd, Connection, ToRedisArgs, FromRedisValue, RedisResult};

/// Redis::Cmd's `args` field is private,
/// but we need it to determine a slot from the command.
/// So this is a simple wrapper around Redis::Cmd that keeps
/// track of the args.
pub struct ClusterCmd {
    cmd: Cmd,
    args: Vec<Vec<u8>>,
}

impl ClusterCmd {
    pub fn new() -> ClusterCmd {
        ClusterCmd {
            cmd: Cmd::new(),
            args: Vec::new(),
        }
    }

    /// Add an arg to the command.
    pub fn arg<T: ToRedisArgs>(&mut self, arg: T) -> &mut ClusterCmd {
        for item in arg.to_redis_args().into_iter() {
            self.args.push(item);
        }
        self.cmd.arg(arg);
        self
    }

    /// Execute the command, returning the result.
    pub fn query<T: FromRedisValue>(&self, conn: &Connection) -> RedisResult<T> {
        self.cmd.query(conn)
    }

    /// Get the slot for this command.
    pub fn slot(&self) -> Option<u16> {
        slot_for_command(&self.args)
    }
}

fn slot_for_command(args: &Vec<Vec<u8>>) -> Option<u16> {
    if args.len() > 1 {
        Some(key_hash_slot(args[1].as_slice()))
    } else {
        None
    }
}

pub fn slot_for_packed_command(cmd: &[u8]) -> Option<u16> {
    let args = unpack_command(cmd);
    slot_for_command(&args)
}

/// `redis-rs` passes packed commands (as a u8 slice)
/// to the methods of the Commands trait
/// we need to "unpack" the command into the
/// original arguments to properly compute
/// the command's slot.
/// This is pretty messy/can probably be better
fn unpack_command(cmd: &[u8]) -> Vec<Vec<u8>> {
    let mut arg: Vec<u8> = Vec::new();
    let mut args: Vec<Vec<u8>> = Vec::new();

    // first 4 are some leading info ('*', len of args, '\r', '\n')
    // the next 4 precede the first arg
    // see: <https://github.com/mitsuhiko/redis-rs/blob/master/src/cmd.rs#L85>
    let mut iter = cmd.iter().skip(4).skip(4).peekable();

    'outer: loop {
        let b = *iter.next().unwrap();

        // args are separated by [13, 10]
        if b == 13 && iter.peek().unwrap() == &&10 {
            args.push(arg.clone());
            arg.clear();

            // consume next 4 which precede the next arg
            for _ in 0..5 {
                match iter.next() {
                    Some(_) => (),
                    None => break 'outer,
                };
            }
        } else {
            arg.push(b);
        }
    }
    args
}