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
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102

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(2).peekable();

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

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

            // consume the next item (10)
            iter.next();

            // then, if there are more args, there should be a 36
            // if there's nothing, we're done
            match iter.next() {
                Some(_) => (),
                None => break 'outer,
            };
            // then the length of the args (which in theory can be any length)
            // then another 13, 10
            'inner: loop {
                let b = *iter.next().unwrap();
                if b == 13 && iter.peek().unwrap() == &&10 {
                    iter.next();
                    break 'inner;
                }
            }
        } else {
            arg.push(b);
        }
    }
    args
}