Full Code of yangwenmai/ratelimit for AI

Repository: yangwenmai/ratelimit
Branch: master
Commit: 44221c2292e1
Files: 11
Total size: 12.7 KB

Directory structure:
gitextract_1wr6y951/

├── .gitignore
├── .travis.yml
├── LICENSE
├── Makefile
├── README.md
├── example/
│   └── main.go
├── leakybucket/
│   ├── bucket.go
│   ├── leaky_bucket_test.go
│   └── leaky_memory.go
└── simpleratelimit/
    ├── ratelimit.go
    └── ratelimit_test.go

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
bin
pkg
*.swp
.DS_Store
.idea
*.iml
.vscode/
*.overprofile
*/*.overprofile


================================================
FILE: .travis.yml
================================================
language: go
sudo: false

notifications:
  email:
    recipients:
      - yangwen.yw@gmail.com
  on_success: change
  on_failure: always

go:
  - 1.8.3

install:
  - go get github.com/go-playground/overalls
  - go get github.com/mattn/goveralls
  - go get github.com/smartystreets/goconvey
  - mkdir -p $GOPATH/src/github.com/yangwenmai
  - cd $GOPATH/src/github.com/yangwenmai/ratelimit

script:
    - overalls -project=github.com/yangwenmai/ratelimit -covermode=count -ignore='.git,_vendor'
    - goveralls -coverprofile=overalls.coverprofile -service=travis-ci -repotoken $COVERALLS_TOKEN
    - go test ./...

env:
  global:
    secure: "bopJ2qeaADkDDOTDOPfYiO3mHigtyUBM3Ec8vhc9jkQWQ9Spmd4rXnKjE0ry7ZWTBdLze5Fkgga83/5YuRaecLKYPLtYLXIF1KioJvJiuTlAkMVrlvuuOMR7YN1qq4hHybeWZI/YEvj+fPjm0VRFU7syrLzsvIX0wuhDOrW5fvDOBW/Xn0isZ0UZSLx/AmU5TJp8lYG/WMPDUSyyxMNDOT4Kv0ZExMIXlIEtzFJXSb8yuOAZVIVJ4IiFXjDeMIrc7krzQa/sqQb2OujjjJqrDqdQTcbaPpGL4q+eewXdVhG6+JQlK6pa8K2LfQTERD0fVXcxrF3AXsmOiqh4/aLpez69pqYebiwwm2SK+hcwm2AZiUg2Phrvk0gqyoDRWSEV2AgOeMxLdSWyAWjreA2H3YtbvLpe2RCXiv05WS0BYvUwmFxYUDgq679AzZ1YHZRQTZTtMmCciunQekBLb+P+qbmuyv7h60mTYGrfZLRjoPUreQ0MDqwEjB8f0mFMrq66LPlMThbsvwzUFqtfBbclWP+HVLiRnd0kKFkV0HnBt15dvYQxr80Oad+4TYYJ0WWiuy/LV9D44BDgL59CKOTIu7jd0yjP5gmBO7wJhx0UhZ8f/SZPNX8GXXx7tfmjpM2H9aA84Zf+Bl+kV2/Ra2NiTWeroblsEutv795gLFwF0i4="


================================================
FILE: LICENSE
================================================
MIT License

Copyright (c) 2017 maiyang

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.


================================================
FILE: Makefile
================================================
default: test

testdeps:
	@go get github.com/smartystreets/goconvey

test: testdeps
	@go test ./...


testrace: testdeps
	@go test ./... -race

testall: test testrace

bench:
	@go test ./... -run=NONE -bench=.


================================================
FILE: README.md
================================================
# ratelimit [![Build Status](https://travis-ci.org/yangwenmai/ratelimit.png?branch=master)](https://travis-ci.org/yangwenmai/ratelimit) [![Documentation](https://godoc.org/github.com/yangwenmai/ratelimit?status.svg)](http://godoc.org/github.com/yangwenmai/ratelimit) [![Go Report Card](https://goreportcard.com/badge/github.com/yangwenmai/ratelimit)](https://goreportcard.com/report/github.com/yangwenmai/ratelimit) [![Coverage Status](https://coveralls.io/repos/github/yangwenmai/ratelimit/badge.svg?branch=master)](https://coveralls.io/github/yangwenmai/ratelimit?branch=master)

基于令牌桶算法和漏桶算法来实现的限速限流，Golang实现。

## Stargazers over time

[![Stargazers over time](https://starcharts.herokuapp.com/yangwenmai/ratelimit.svg)](https://starcharts.herokuapp.com/yangwenmai/ratelimit)

# 算法介绍

## 漏桶算法

漏桶算法思路很简单，水（请求）先进入到漏桶里，漏桶以一定的速度出水，当水流入速度过大会直接溢出，可以看出漏桶算法能强行限制数据的传输速率。

![leaky-bucket](https://github.com/yangwenmai/ratelimit/blob/master/doc/leaky-bucket.png)

漏桶算法示意图

漏桶算法可以很好地限制容量池的大小，从而防止流量暴增。

## 令牌桶算法

令牌桶算法的原理是系统会以一个恒定的速度往桶里放入令牌，而如果请求需要被处理，则需要先从桶里获取一个令牌，当桶里没有令牌可取时，则拒绝服务。

![token-bucket](https://github.com/yangwenmai/ratelimit/blob/master/doc/token-bucket.jpg)

令牌桶算法示意图

令牌桶算法通过发放令牌，根据令牌的rate频率做请求频率限制，容量限制等。

### 示例

```go
package main

import (
	"log"
	"time"

	"github.com/yangwenmai/ratelimit/leakybucket"
	"github.com/yangwenmai/ratelimit/simpleratelimit"
)

func main() {
    // rate limit: simple
	rl := simpleratelimit.New(10, time.Second)

	for i := 0; i < 100; i++ {
		log.Printf("limit result: %v\n", rl.Limit())
	}
	log.Printf("limit result: %v\n", rl.Limit())

    // rate limit: leaky-bucket
	lb := leakybucket.New()
	b, err := lb.Create("leaky_bucket", 10, time.Second)
	if err != nil {
		log.Println(err)
	}
	log.Printf("bucket capacity:%v", b.Capacity())

    // rate limit: token-bucket
}
```

# 参考资料

1. [限流:漏桶算法和令牌桶算法](http://maiyang.github.io/技术/算法/2017/05/28/rate-limit-algorithm)
2. [维基百科:Token_bucket](https://en.wikipedia.org/wiki/Token_bucket)
3. [维基百科:Leaky_bucket](https://en.wikipedia.org/wiki/Leaky_bucket)
4. [接口限流实践](http://www.cnblogs.com/LBSer/p/4083131.html)
5. [流量调整和限流技术](http://colobu.com/2014/11/13/rate-limiting/)


================================================
FILE: example/main.go
================================================
package main

import (
	"log"
	"time"

	"github.com/yangwenmai/ratelimit/leakybucket"
	"github.com/yangwenmai/ratelimit/simpleratelimit"
)

func main() {
	// rate limit: simple
	rl := simpleratelimit.New(10, time.Second)

	for i := 0; i < 100; i++ {
		log.Printf("limit result: %v\n", rl.Limit())
	}
	log.Printf("limit result: %v\n", rl.Limit())

	// rate limit: leaky-bucket
	lb := leakybucket.New()
	b, err := lb.Create("leaky_bucket", 10, time.Second)
	if err != nil {
		log.Println(err)
	}
	log.Printf("bucket capacity:%v", b.Capacity())

	// rate limit: token-bucket
}


================================================
FILE: leakybucket/bucket.go
================================================
package leakybucket

import (
	"errors"
	"time"
)

var (
	// ErrorFull is returned when the amount requested to add exceeds the remaining space in the bucket.
	ErrorFull = errors.New("add exceeds free capacity")
)

// BucketI interface for interacting with leaky buckets: https://en.wikipedia.org/wiki/Leaky_bucket
type BucketI interface {
	// Capacity of the bucket.
	Capacity() uint
	// Remaining space in the bucket.
	Remaining() uint
	// Reset returns when the bucket will be drained.
	Reset() time.Time
	// Add to the bucket. Returns bucket state after adding.
	Add(uint) (BucketState, error)
}

// BucketState is a snapshot of a bucket's properties.
type BucketState struct {
	Capacity  uint
	Remaining uint
	Reset     time.Time
}

// StorageI interface for generating buckets keyed by a string.
type StorageI interface {
	// Create a bucket with a name, capacity, and rate.
	// rate is how long it takes for full capacity to drain.
	Create(name string, capacity uint, rate time.Duration) (BucketI, error)
}


================================================
FILE: leakybucket/leaky_bucket_test.go
================================================
package leakybucket

import (
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestLimit(t *testing.T) {
	Convey("start testing", t, func() {
		Convey("create capacity 100", func() {
			s := New()
			bucket, err := s.Create("testCreateBucket", 100, time.Second)
			if err != nil {
				t.Error(err)
			}
			if capacity := bucket.Capacity(); capacity != 100 {
				t.Fatalf("expected capacity of %d, got %d", 100, capacity)
			}
		})
		Convey("same name", func() {
			s := New()
			bucket, err := s.Create("testSame", 10, time.Second)
			if err != nil {
				t.Error(err)
			}
			bucket2, err := s.Create("testSame", 100, time.Second)
			if err != nil {
				t.Error(err)
			}
			if bucket != bucket2 {
				t.Fatalf("not same.")
			}
		})
		Convey("add()", func() {
			s := New()
			bucket, err := s.Create("testAdd", 100, time.Second)
			if err != nil {
				t.Fatal(err)
			}

			addAndTestRemaining := func(add, remaining uint) {
				if state, err := bucket.Add(add); err != nil {
					t.Fatal(err)
				} else if bucket.Remaining() != state.Remaining {
					t.Fatalf("expected bucket and state remaining to match, bucket is %d, state is %d",
						bucket.Remaining(), state.Remaining)
				} else if state.Remaining != remaining {
					t.Fatalf("expected %d remaining, got %d", remaining, state.Remaining)
				}
			}
			addAndTestRemaining(1, 99)
			addAndTestRemaining(3, 96)
			addAndTestRemaining(90, 6)
			if _, err := bucket.Add(7); err == nil {
				t.Fatalf("expected ErrorFull, received no error")
			} else if err != ErrorFull {
				t.Fatalf("expected ErrorFull, received %v", err)
			}
		})
		Convey("Reset()", func() {
			s := New()
			bucket, err := s.Create("testReset", 1, time.Millisecond)
			if err != nil {
				t.Fatal(err)
			}
			if _, err := bucket.Add(1); err != nil {
				t.Fatal(err)
			}
			time.Sleep(time.Millisecond * 2)
			bucket.Reset()
			if bucket.Remaining() != 1 {
				t.Fatalf("expected 1, got %d", bucket.Remaining())
			}
			state, err := bucket.Add(1)
			if err != nil {
				t.Fatal(err)
			}
			if state.Remaining != 0 {
				t.Fatalf("expected full bucket, got %d", state.Remaining)
			}
			if state.Reset.Unix() < time.Now().Unix() {
				t.Fatalf("reset time is in the past")
			}
			if bucket.Reset().Unix() < time.Now().Unix() {
				t.Fatalf("reset time is in the past")
			}
		})
	})
}

func BenchmarkLimit(b *testing.B) {
	Convey("start testing benchmark", b, func() {
	})
}


================================================
FILE: leakybucket/leaky_memory.go
================================================
package leakybucket

import (
	"sync"
	"time"
)

type bucket struct {
	capacity  uint
	remaining uint
	reset     time.Time
	rate      time.Duration
	mutex     sync.Mutex
}

// Capacity return
func (b *bucket) Capacity() uint {
	return b.capacity
}

// Remaining space in the bucket.
func (b *bucket) Remaining() uint {
	return b.remaining
}

// Reset returns when the bucket will be drained.
func (b *bucket) Reset() time.Time {
	b.remaining = b.capacity
	return b.reset
}

// Add to the bucket.
func (b *bucket) Add(amount uint) (BucketState, error) {
	b.mutex.Lock()
	defer b.mutex.Unlock()
	if time.Now().After(b.reset) {
		b.reset = time.Now().Add(b.rate)
		b.remaining = b.capacity
	}
	if amount > b.remaining {
		return BucketState{Capacity: b.capacity, Remaining: b.remaining, Reset: b.reset}, ErrorFull
	}
	b.remaining -= amount
	return BucketState{Capacity: b.capacity, Remaining: b.remaining, Reset: b.reset}, nil
}

// Storage is a non thread-safe in-memory leaky bucket factory.
type Storage struct {
	buckets map[string]*bucket
}

// New initializes the in-memory bucket store.
func New() *Storage {
	return &Storage{
		buckets: make(map[string]*bucket),
	}
}

// Create a bucket.
func (s *Storage) Create(name string, capacity uint, rate time.Duration) (BucketI, error) {
	b, ok := s.buckets[name]
	if ok {
		return b, nil
	}
	b = &bucket{
		capacity:  capacity,
		remaining: capacity,
		reset:     time.Now().Add(rate),
		rate:      rate,
	}
	s.buckets[name] = b
	return b, nil
}


================================================
FILE: simpleratelimit/ratelimit.go
================================================
package simpleratelimit

import (
	"sync/atomic"
	"time"
)

// RateLimiter 限速器
type RateLimiter struct {
	rate      uint64
	allowance uint64
	max       uint64
	unit      uint64
	lastCheck uint64
}

// New 创建RateLimiter实例
func New(rate int, per time.Duration) *RateLimiter {
	nano := uint64(per)
	if nano < 1 {
		nano = uint64(time.Second)
	}
	if rate < 1 {
		rate = 1
	}

	return &RateLimiter{
		rate:      uint64(rate),
		allowance: uint64(rate) * nano,
		max:       uint64(rate) * nano,
		unit:      nano,

		lastCheck: unixNano(),
	}
}

// Limit 判断是否超过限制
func (rl *RateLimiter) Limit() bool {
	now := unixNano()
	// 计算上一次调用到现在过了多少纳秒
	passed := now - atomic.SwapUint64(&rl.lastCheck, now)

	rate := atomic.LoadUint64(&rl.rate)
	current := atomic.AddUint64(&rl.allowance, passed*rate)

	if max := atomic.LoadUint64(&rl.max); current > max {
		atomic.AddUint64(&rl.allowance, max-current)
		current = max
	}

	if current < rl.unit {
		return true
	}

	// 没有超过限额
	atomic.AddUint64(&rl.allowance, -rl.unit)
	return false
}

// UpdateRate 更新速率值
func (rl *RateLimiter) UpdateRate(rate int) {
	atomic.StoreUint64(&rl.rate, uint64(rate))
	atomic.StoreUint64(&rl.max, uint64(rate)*rl.unit)
}

// Undo 重置上一次调用Limit()，返回没有使用过的限额
func (rl *RateLimiter) Undo() {
	current := atomic.AddUint64(&rl.allowance, rl.unit)

	if max := atomic.LoadUint64(&rl.max); current > max {
		atomic.AddUint64(&rl.allowance, max-current)
	}
}

// unixNano 当前时间（纳秒）
func unixNano() uint64 {
	return uint64(time.Now().UnixNano())
}


================================================
FILE: simpleratelimit/ratelimit_test.go
================================================
package simpleratelimit

import (
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

func TestLimit(t *testing.T) {
	Convey("start testing New(0, 0)", t, func() {
		rl := New(0, 0*time.Second)
		So(rl.Limit(), ShouldEqual, false)
	})
	Convey("start testing limit", t, func() {
		rl := New(1, time.Second)
		for i := 0; i < 10; i++ {
			if i == 0 {
				So(rl.Limit(), ShouldEqual, false)
			} else {
				So(rl.Limit(), ShouldEqual, true)
			}
		}
	})
	Convey("start testing updateRate", t, func() {
		rl := New(1, time.Second)
		for i := 0; i < 2; i++ {
			if i == 0 {
				So(rl.Limit(), ShouldEqual, false)
			} else {
				rl.UpdateRate(2)
				So(rl.Limit(), ShouldEqual, true)
			}
		}
	})
	Convey("start testing undo", t, func() {
		rl := New(1, time.Second)
		for i := 0; i < 10; i++ {
			if i == 0 {
				So(rl.Limit(), ShouldEqual, false)
			} else {
				rl.Undo()
				So(rl.Limit(), ShouldEqual, false)
			}
		}
		for i := 0; i < 10; i++ {
			rl.Limit()
		}
		rl.Undo()
	})
}

func BenchmarkLimit(b *testing.B) {
	Convey("start testing benchmark", b, func() {
		rl := New(1, time.Second)
		for i := 0; i < b.N; i++ {
			rl.Limit()
		}
	})
}