Redis 限流器

之前一直在思考如何实现限流器，最近看redis命令学习到了可以用redis来实现限流器的功能，简单方便。

用redis来设置限流器，20秒钟不超过10次，根据key取出value，如果value不存在则设置value自动加一(incr)，然后设置超时时间(20)；如果value存在并且小于10，则自增1(incr)；如果value大于10则返回错误。

package main

import (
    "errors"
    "fmt"
    "net/http"
    "strconv"
    "strings"

    "github.com/garyburd/redigo/redis"
)

var c redis.Conn

func limiter(id string) (int64, error) {

    var sum int64
    var err error
    val, _ := redis.String(c.Do("GET", id))
    if val == "" {
        sum, err = redis.Int64(c.Do("INCR", id))
        if err != nil {
            fmt.Println("---incr is failed, err: ", err)
            return 0, err
        }
        _, err = c.Do("EXPIRE", id, 20)
        if err != nil {
            fmt.Println("---err: ", err)
            return 0, err
        }

    } else {
        sum, err = strconv.ParseInt(val, 10, 0)
        if err != nil {
            fmt.Println("--atoi is failed, err: ", err)
            return 0, err
        }
        if sum > 10 {
            return 0, errors.New("sum is max.")
        } else {
            sum, err = redis.Int64(c.Do("INCR", id))
            if err != nil {
                fmt.Println("---incr2 is failed, err: ", err)
            }
        }
    }

    return sum, nil
}

func restrictor(resp http.ResponseWriter, req *http.Request) {
    urls := strings.Split(req.URL.Path, "/")

    num, err := limiter(urls[2])
    if err != nil {
        resp.Write([]byte(err.Error()))
    }
    resp.Write([]byte("num is " + strconv.Itoa(int(num))))
}

func main() {
    var err error
    c, err = redis.Dial("tcp", "10.10.30.54:6379")
    if err != nil {
        fmt.Println("connect redis server is failed")
        return
    }
    http.HandleFunc("/restrictor/", restrictor)
    http.ListenAndServe(":9978", nil)
}

 

链接上 Redis，在浏览器中输入 http://localhost:9978/restrictor/{id}即可。根据不同用户不同的id，来进行区分。

例：http://localhost:9978/restrictor/10102

会收到返回 num is 7