golang redis 令牌桶

const (

    tpsBucketToken = "TpsTokenBucket"

)

 

type tpsToken struct{}

 

// NewTpsToken 控制mysql tps

func NewTpsToken() *tpsToken {

    return &tpsToken{}

}

 

func (t *tpsToken) InitTokenBucket(key string, cnt int) error {

    globalRedis.Del(key)

    for i := 1; i <= cnt; i++ {

        _, err := globalRedis.SAdd(key, i).Result()

        if err != nil {

            return err

        }

    }

    return nil

}

 

func (t *tpsToken) GetToken(key string) string {

    token, err := globalRedis.SPop(key).Result()

    if err == nil && token != "" {

        return token

    }

    return ""

}

func (t *tpsToken) ReleaseToken(key, token string) (err error) {

    _, err = globalRedis.SAdd(key, token).Result()

    return

}