golimit

README

golimit

A simple redis backed rate limiter for distributed systems based on token bucket algorithm.

Install

go get "github.com/akhiljames/golimit"

Usage

package main

import (
	"fmt"
	"strings"
	"time"

	"github.com/akhiljames/golimit"
	"github.com/go-redis/redis"
)

// Redis client. You can use your own redis client
type Redis struct {
	client *redis.Client
}

// Eval exec script
func (r *Redis) Eval(script string, keys []string, args ...interface{}) (interface{}, error) {
	return r.client.Eval(script, keys, args...).Result()
}

// EvalSha Optimize Lua script execution
func (r *Redis) EvalSha(sha1 string, keys []string, args ...interface{}) (interface{}, error, bool) {
	result, err := r.client.EvalSha(sha1, keys, args...).Result()
	noScript := err != nil && strings.HasPrefix(err.Error(), "NOSCRIPT ")
	return result, err, noScript
}

func main() {
	// create a new rate limiter
	tb := golimit.NewRateLimiter(
		&Redis{redis.NewClient(&redis.Options{
			Addr: "localhost:6379",
		})},
		"golimit:example", // this is a unique key for redis
		&golimit.Config{
			Interval: 1 * time.Second, // interval for rate limit
			Capacity: 5,               // capacity to be processed in the interval
			// this will rate limt as per - capacity amount in interval time
		},
	)
	if ok, err := tb.Take(1); ok { // usage of operation can be taken. for eg. heavy db ops
		// serve the user request
	} else {
		if err != nil {
			fmt.Println(err.Error())
		}
		// reject the user request
	}
}

Documentation

Index

• type Config
• type RateLimiter
  • func NewRateLimiter(redis Redis, key string, config *Config) *RateLimiter
  • func (b *RateLimiter) Config() Config
  • func (b *RateLimiter) SetConfig(config *Config)
  • func (b *RateLimiter) Take(amount int64) (bool, error)
• type Redis
• type Script
  • func NewScript(redis Redis, src string) *Script
  • func (s *Script) Run(keys []string, args ...interface{}) (interface{}, error)

Examples

• RateLimiter.Take

Types

type Config

type Config struct {
	Interval time.Duration // the interval between each addition of one token

	// the capacity of the bucket
	Capacity int64
}

Config is the bucket configuration.

type RateLimiter

type RateLimiter struct {
	// contains filtered or unexported fields
}

RateLimiter implements the Token Bucket Algorithm. See https://en.wikipedia.org/wiki/Token_bucket.

func NewRateLimiter

func NewRateLimiter(redis Redis, key string, config *Config) *RateLimiter

NewRateLimiter returns a new token-bucket rate limiter special for key in redis with the specified bucket configuration.

func (*RateLimiter) Config

func (b *RateLimiter) Config() Config

Config returns the bucket configuration in a concurrency-safe way.

func (*RateLimiter) SetConfig

func (b *RateLimiter) SetConfig(config *Config)

SetConfig updates the bucket configuration in a concurrency-safe way.

func (*RateLimiter) Take

func (b *RateLimiter) Take(amount int64) (bool, error)

Take takes amount tokens from the bucket.

Example

package main

import (
	"fmt"
	"strings"
	"time"

	"github.com/akhiljames/golimit"
	"github.com/go-redis/redis"
)

type Redis struct {
	client *redis.Client
}

func (r *Redis) Eval(script string, keys []string, args ...interface{}) (interface{}, error) {
	return r.client.Eval(script, keys, args...).Result()
}

// EvalSha Optimize Lua script execution
func (r *Redis) EvalSha(sha1 string, keys []string, args ...interface{}) (interface{}, error, bool) {
	result, err := r.client.EvalSha(sha1, keys, args...).Result()
	noScript := err != nil && strings.HasPrefix(err.Error(), "NOSCRIPT ")
	return result, err, noScript
}

func main() {
	tb := golimit.NewRateLimiter(
		&Redis{redis.NewClient(&redis.Options{
			Addr: "localhost:6379",
		})},
		"golimit:example",
		&golimit.Config{
			Interval: 1 * time.Second / 2,
			Capacity: 5,
		},
	)
	if ok, err := tb.Take(1); ok {
		fmt.Println("PASS")
	} else {
		if err != nil {
			fmt.Println(err.Error())
		}
		fmt.Println("DROP")
	}
}

Output:

PASS

type Redis

type Redis interface {
	Eval(script string, keys []string, args ...interface{}) (interface{}, error)
	EvalSha(sha1 string, keys []string, args ...interface{}) (interface{}, error, bool)
}

Redis interface

type Script

type Script struct {
	// contains filtered or unexported fields
}

Script struct

func NewScript

func NewScript(redis Redis, src string) *Script

NewScript returns new script config

func (*Script) Run

func (s *Script) Run(keys []string, args ...interface{}) (interface{}, error)

Run script execution by using the EvalSha command