redis_cell/
lib.rs

1#[macro_use]
2extern crate bitflags;
3extern crate libc;
4extern crate time;
5
6#[macro_use]
7mod macros;
8
9pub mod cell;
10pub mod error;
11mod redis;
12
13use cell::store;
14use error::CellError;
15use libc::c_int;
16use redis::Command;
17use redis::raw;
18
19const MODULE_NAME: &str = "redis-cell";
20const MODULE_VERSION: c_int = 1;
21
22// ThrottleCommand provides GCRA rate limiting as a command in Redis.
23struct ThrottleCommand {}
24
25impl Command for ThrottleCommand {
26    // Should return the name of the command to be registered.
27    fn name(&self) -> &'static str {
28        "cl.throttle"
29    }
30
31    // Run the command.
32    fn run(&self, r: redis::Redis, args: &[&str]) -> Result<(), CellError> {
33        if args.len() != 5 && args.len() != 6 {
34            return Err(error!(
35                "Usage: {} <key> <max_burst> <count per period> \
36                 <period> [<quantity>]",
37                self.name()
38            ));
39        }
40
41        // the first argument is command name "cl.throttle" (ignore it)
42        let key = args[1];
43        let max_burst = parse_i64(args[2])?;
44        let count = parse_i64(args[3])?;
45        let period = parse_i64(args[4])?;
46        let quantity = match args.get(5) {
47            Some(n) => parse_i64(n)?,
48            None => 1,
49        };
50
51        // We reinitialize a new store and rate limiter every time this command
52        // is run, but these structures don't have a huge overhead to them so
53        // it's not that big of a problem.
54        let mut store = store::InternalRedisStore::new(&r);
55        let rate = cell::Rate::per_period(count, time::Duration::seconds(period));
56        let mut limiter = cell::RateLimiter::new(
57            &mut store,
58            &cell::RateQuota {
59                max_burst,
60                max_rate: rate,
61            },
62        );
63
64        let (throttled, rate_limit_result) = limiter.rate_limit(key, quantity)?;
65
66        // Reply with an array containing rate limiting results. Note that
67        // Redis' support for interesting data types is quite weak, so we have
68        // to jam a few square pegs into round holes. It's a little messy, but
69        // the interface comes out as pretty workable.
70        r.reply_array(5)?;
71        r.reply_integer(if throttled { 1 } else { 0 })?;
72        r.reply_integer(rate_limit_result.limit)?;
73        r.reply_integer(rate_limit_result.remaining)?;
74        r.reply_integer(rate_limit_result.retry_after.num_seconds())?;
75        r.reply_integer(rate_limit_result.reset_after.num_seconds())?;
76
77        Ok(())
78    }
79
80    // Should return any flags to be registered with the name as a string
81    // separated list. See the Redis module API documentation for a complete
82    // list of the ones that are available.
83    fn str_flags(&self) -> &'static str {
84        "write"
85    }
86}
87
88#[allow(non_snake_case)]
89#[allow(unused_variables)]
90#[no_mangle]
91pub extern "C" fn Throttle_RedisCommand(
92    ctx: *mut raw::RedisModuleCtx,
93    argv: *mut *mut raw::RedisModuleString,
94    argc: c_int,
95) -> raw::Status {
96    Command::harness(&ThrottleCommand {}, ctx, argv, argc)
97}
98
99#[allow(non_snake_case)]
100#[allow(unused_variables)]
101#[no_mangle]
102pub extern "C" fn RedisModule_OnLoad(
103    ctx: *mut raw::RedisModuleCtx,
104    argv: *mut *mut raw::RedisModuleString,
105    argc: c_int,
106) -> raw::Status {
107    if raw::init(
108        ctx,
109        format!("{}\0", MODULE_NAME).as_ptr(),
110        MODULE_VERSION,
111        raw::REDISMODULE_APIVER_1,
112    ) == raw::Status::Err
113    {
114        return raw::Status::Err;
115    }
116
117    let command = ThrottleCommand {};
118    if raw::create_command(
119        ctx,
120        format!("{}\0", command.name()).as_ptr(),
121        Some(Throttle_RedisCommand),
122        format!("{}\0", command.str_flags()).as_ptr(),
123        0,
124        0,
125        0,
126    ) == raw::Status::Err
127    {
128        return raw::Status::Err;
129    }
130
131    raw::Status::Ok
132}
133
134fn parse_i64(arg: &str) -> Result<i64, CellError> {
135    arg.parse::<i64>()
136        .map_err(|_| error!("Couldn't parse as integer: {}", arg))
137}
redis_cell/lib.rs

redis_cell/
lib.rs
