[
  {
    "path": ".gitignore",
    "content": "bin\npkg\n*.swp\n.DS_Store\n.idea\n*.iml\n.vscode/\n*.overprofile\n*/*.overprofile\n"
  },
  {
    "path": ".travis.yml",
    "content": "language: go\nsudo: false\n\nnotifications:\n  email:\n    recipients:\n      - yangwen.yw@gmail.com\n  on_success: change\n  on_failure: always\n\ngo:\n  - 1.8.3\n\ninstall:\n  - go get github.com/go-playground/overalls\n  - go get github.com/mattn/goveralls\n  - go get github.com/smartystreets/goconvey\n  - mkdir -p $GOPATH/src/github.com/yangwenmai\n  - cd $GOPATH/src/github.com/yangwenmai/ratelimit\n\nscript:\n    - overalls -project=github.com/yangwenmai/ratelimit -covermode=count -ignore='.git,_vendor'\n    - goveralls -coverprofile=overalls.coverprofile -service=travis-ci -repotoken $COVERALLS_TOKEN\n    - go test ./...\n\nenv:\n  global:\n    secure: \"bopJ2qeaADkDDOTDOPfYiO3mHigtyUBM3Ec8vhc9jkQWQ9Spmd4rXnKjE0ry7ZWTBdLze5Fkgga83/5YuRaecLKYPLtYLXIF1KioJvJiuTlAkMVrlvuuOMR7YN1qq4hHybeWZI/YEvj+fPjm0VRFU7syrLzsvIX0wuhDOrW5fvDOBW/Xn0isZ0UZSLx/AmU5TJp8lYG/WMPDUSyyxMNDOT4Kv0ZExMIXlIEtzFJXSb8yuOAZVIVJ4IiFXjDeMIrc7krzQa/sqQb2OujjjJqrDqdQTcbaPpGL4q+eewXdVhG6+JQlK6pa8K2LfQTERD0fVXcxrF3AXsmOiqh4/aLpez69pqYebiwwm2SK+hcwm2AZiUg2Phrvk0gqyoDRWSEV2AgOeMxLdSWyAWjreA2H3YtbvLpe2RCXiv05WS0BYvUwmFxYUDgq679AzZ1YHZRQTZTtMmCciunQekBLb+P+qbmuyv7h60mTYGrfZLRjoPUreQ0MDqwEjB8f0mFMrq66LPlMThbsvwzUFqtfBbclWP+HVLiRnd0kKFkV0HnBt15dvYQxr80Oad+4TYYJ0WWiuy/LV9D44BDgL59CKOTIu7jd0yjP5gmBO7wJhx0UhZ8f/SZPNX8GXXx7tfmjpM2H9aA84Zf+Bl+kV2/Ra2NiTWeroblsEutv795gLFwF0i4=\"\n"
  },
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2017 maiyang\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "Makefile",
    "content": "default: test\n\ntestdeps:\n\t@go get github.com/smartystreets/goconvey\n\ntest: testdeps\n\t@go test ./...\n\n\ntestrace: testdeps\n\t@go test ./... -race\n\ntestall: test testrace\n\nbench:\n\t@go test ./... -run=NONE -bench=.\n"
  },
  {
    "path": "README.md",
    "content": "# 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)\n\n基于令牌桶算法和漏桶算法来实现的限速限流，Golang实现。\n\n## Stargazers over time\n\n[![Stargazers over time](https://starcharts.herokuapp.com/yangwenmai/ratelimit.svg)](https://starcharts.herokuapp.com/yangwenmai/ratelimit)\n\n# 算法介绍\n\n## 漏桶算法\n\n漏桶算法思路很简单，水（请求）先进入到漏桶里，漏桶以一定的速度出水，当水流入速度过大会直接溢出，可以看出漏桶算法能强行限制数据的传输速率。\n\n![leaky-bucket](https://github.com/yangwenmai/ratelimit/blob/master/doc/leaky-bucket.png)\n\n漏桶算法示意图\n\n漏桶算法可以很好地限制容量池的大小，从而防止流量暴增。\n\n## 令牌桶算法\n\n令牌桶算法的原理是系统会以一个恒定的速度往桶里放入令牌，而如果请求需要被处理，则需要先从桶里获取一个令牌，当桶里没有令牌可取时，则拒绝服务。\n\n![token-bucket](https://github.com/yangwenmai/ratelimit/blob/master/doc/token-bucket.jpg)\n\n令牌桶算法示意图\n\n令牌桶算法通过发放令牌，根据令牌的rate频率做请求频率限制，容量限制等。\n\n### 示例\n\n```go\npackage main\n\nimport (\n\t\"log\"\n\t\"time\"\n\n\t\"github.com/yangwenmai/ratelimit/leakybucket\"\n\t\"github.com/yangwenmai/ratelimit/simpleratelimit\"\n)\n\nfunc main() {\n    // rate limit: simple\n\trl := simpleratelimit.New(10, time.Second)\n\n\tfor i := 0; i < 100; i++ {\n\t\tlog.Printf(\"limit result: %v\\n\", rl.Limit())\n\t}\n\tlog.Printf(\"limit result: %v\\n\", rl.Limit())\n\n    // rate limit: leaky-bucket\n\tlb := leakybucket.New()\n\tb, err := lb.Create(\"leaky_bucket\", 10, time.Second)\n\tif err != nil {\n\t\tlog.Println(err)\n\t}\n\tlog.Printf(\"bucket capacity:%v\", b.Capacity())\n\n    // rate limit: token-bucket\n}\n```\n\n# 参考资料\n\n1. [限流:漏桶算法和令牌桶算法](http://maiyang.github.io/技术/算法/2017/05/28/rate-limit-algorithm)\n2. [维基百科:Token_bucket](https://en.wikipedia.org/wiki/Token_bucket)\n3. [维基百科:Leaky_bucket](https://en.wikipedia.org/wiki/Leaky_bucket)\n4. [接口限流实践](http://www.cnblogs.com/LBSer/p/4083131.html)\n5. [流量调整和限流技术](http://colobu.com/2014/11/13/rate-limiting/)\n"
  },
  {
    "path": "example/main.go",
    "content": "package main\n\nimport (\n\t\"log\"\n\t\"time\"\n\n\t\"github.com/yangwenmai/ratelimit/leakybucket\"\n\t\"github.com/yangwenmai/ratelimit/simpleratelimit\"\n)\n\nfunc main() {\n\t// rate limit: simple\n\trl := simpleratelimit.New(10, time.Second)\n\n\tfor i := 0; i < 100; i++ {\n\t\tlog.Printf(\"limit result: %v\\n\", rl.Limit())\n\t}\n\tlog.Printf(\"limit result: %v\\n\", rl.Limit())\n\n\t// rate limit: leaky-bucket\n\tlb := leakybucket.New()\n\tb, err := lb.Create(\"leaky_bucket\", 10, time.Second)\n\tif err != nil {\n\t\tlog.Println(err)\n\t}\n\tlog.Printf(\"bucket capacity:%v\", b.Capacity())\n\n\t// rate limit: token-bucket\n}\n"
  },
  {
    "path": "leakybucket/bucket.go",
    "content": "package leakybucket\n\nimport (\n\t\"errors\"\n\t\"time\"\n)\n\nvar (\n\t// ErrorFull is returned when the amount requested to add exceeds the remaining space in the bucket.\n\tErrorFull = errors.New(\"add exceeds free capacity\")\n)\n\n// BucketI interface for interacting with leaky buckets: https://en.wikipedia.org/wiki/Leaky_bucket\ntype BucketI interface {\n\t// Capacity of the bucket.\n\tCapacity() uint\n\t// Remaining space in the bucket.\n\tRemaining() uint\n\t// Reset returns when the bucket will be drained.\n\tReset() time.Time\n\t// Add to the bucket. Returns bucket state after adding.\n\tAdd(uint) (BucketState, error)\n}\n\n// BucketState is a snapshot of a bucket's properties.\ntype BucketState struct {\n\tCapacity  uint\n\tRemaining uint\n\tReset     time.Time\n}\n\n// StorageI interface for generating buckets keyed by a string.\ntype StorageI interface {\n\t// Create a bucket with a name, capacity, and rate.\n\t// rate is how long it takes for full capacity to drain.\n\tCreate(name string, capacity uint, rate time.Duration) (BucketI, error)\n}\n"
  },
  {
    "path": "leakybucket/leaky_bucket_test.go",
    "content": "package leakybucket\n\nimport (\n\t\"testing\"\n\t\"time\"\n\n\t. \"github.com/smartystreets/goconvey/convey\"\n)\n\nfunc TestLimit(t *testing.T) {\n\tConvey(\"start testing\", t, func() {\n\t\tConvey(\"create capacity 100\", func() {\n\t\t\ts := New()\n\t\t\tbucket, err := s.Create(\"testCreateBucket\", 100, time.Second)\n\t\t\tif err != nil {\n\t\t\t\tt.Error(err)\n\t\t\t}\n\t\t\tif capacity := bucket.Capacity(); capacity != 100 {\n\t\t\t\tt.Fatalf(\"expected capacity of %d, got %d\", 100, capacity)\n\t\t\t}\n\t\t})\n\t\tConvey(\"same name\", func() {\n\t\t\ts := New()\n\t\t\tbucket, err := s.Create(\"testSame\", 10, time.Second)\n\t\t\tif err != nil {\n\t\t\t\tt.Error(err)\n\t\t\t}\n\t\t\tbucket2, err := s.Create(\"testSame\", 100, time.Second)\n\t\t\tif err != nil {\n\t\t\t\tt.Error(err)\n\t\t\t}\n\t\t\tif bucket != bucket2 {\n\t\t\t\tt.Fatalf(\"not same.\")\n\t\t\t}\n\t\t})\n\t\tConvey(\"add()\", func() {\n\t\t\ts := New()\n\t\t\tbucket, err := s.Create(\"testAdd\", 100, time.Second)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\n\t\t\taddAndTestRemaining := func(add, remaining uint) {\n\t\t\t\tif state, err := bucket.Add(add); err != nil {\n\t\t\t\t\tt.Fatal(err)\n\t\t\t\t} else if bucket.Remaining() != state.Remaining {\n\t\t\t\t\tt.Fatalf(\"expected bucket and state remaining to match, bucket is %d, state is %d\",\n\t\t\t\t\t\tbucket.Remaining(), state.Remaining)\n\t\t\t\t} else if state.Remaining != remaining {\n\t\t\t\t\tt.Fatalf(\"expected %d remaining, got %d\", remaining, state.Remaining)\n\t\t\t\t}\n\t\t\t}\n\t\t\taddAndTestRemaining(1, 99)\n\t\t\taddAndTestRemaining(3, 96)\n\t\t\taddAndTestRemaining(90, 6)\n\t\t\tif _, err := bucket.Add(7); err == nil {\n\t\t\t\tt.Fatalf(\"expected ErrorFull, received no error\")\n\t\t\t} else if err != ErrorFull {\n\t\t\t\tt.Fatalf(\"expected ErrorFull, received %v\", err)\n\t\t\t}\n\t\t})\n\t\tConvey(\"Reset()\", func() {\n\t\t\ts := New()\n\t\t\tbucket, err := s.Create(\"testReset\", 1, time.Millisecond)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t\tif _, err := bucket.Add(1); err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t\ttime.Sleep(time.Millisecond * 2)\n\t\t\tbucket.Reset()\n\t\t\tif bucket.Remaining() != 1 {\n\t\t\t\tt.Fatalf(\"expected 1, got %d\", bucket.Remaining())\n\t\t\t}\n\t\t\tstate, err := bucket.Add(1)\n\t\t\tif err != nil {\n\t\t\t\tt.Fatal(err)\n\t\t\t}\n\t\t\tif state.Remaining != 0 {\n\t\t\t\tt.Fatalf(\"expected full bucket, got %d\", state.Remaining)\n\t\t\t}\n\t\t\tif state.Reset.Unix() < time.Now().Unix() {\n\t\t\t\tt.Fatalf(\"reset time is in the past\")\n\t\t\t}\n\t\t\tif bucket.Reset().Unix() < time.Now().Unix() {\n\t\t\t\tt.Fatalf(\"reset time is in the past\")\n\t\t\t}\n\t\t})\n\t})\n}\n\nfunc BenchmarkLimit(b *testing.B) {\n\tConvey(\"start testing benchmark\", b, func() {\n\t})\n}\n"
  },
  {
    "path": "leakybucket/leaky_memory.go",
    "content": "package leakybucket\n\nimport (\n\t\"sync\"\n\t\"time\"\n)\n\ntype bucket struct {\n\tcapacity  uint\n\tremaining uint\n\treset     time.Time\n\trate      time.Duration\n\tmutex     sync.Mutex\n}\n\n// Capacity return\nfunc (b *bucket) Capacity() uint {\n\treturn b.capacity\n}\n\n// Remaining space in the bucket.\nfunc (b *bucket) Remaining() uint {\n\treturn b.remaining\n}\n\n// Reset returns when the bucket will be drained.\nfunc (b *bucket) Reset() time.Time {\n\tb.remaining = b.capacity\n\treturn b.reset\n}\n\n// Add to the bucket.\nfunc (b *bucket) Add(amount uint) (BucketState, error) {\n\tb.mutex.Lock()\n\tdefer b.mutex.Unlock()\n\tif time.Now().After(b.reset) {\n\t\tb.reset = time.Now().Add(b.rate)\n\t\tb.remaining = b.capacity\n\t}\n\tif amount > b.remaining {\n\t\treturn BucketState{Capacity: b.capacity, Remaining: b.remaining, Reset: b.reset}, ErrorFull\n\t}\n\tb.remaining -= amount\n\treturn BucketState{Capacity: b.capacity, Remaining: b.remaining, Reset: b.reset}, nil\n}\n\n// Storage is a non thread-safe in-memory leaky bucket factory.\ntype Storage struct {\n\tbuckets map[string]*bucket\n}\n\n// New initializes the in-memory bucket store.\nfunc New() *Storage {\n\treturn &Storage{\n\t\tbuckets: make(map[string]*bucket),\n\t}\n}\n\n// Create a bucket.\nfunc (s *Storage) Create(name string, capacity uint, rate time.Duration) (BucketI, error) {\n\tb, ok := s.buckets[name]\n\tif ok {\n\t\treturn b, nil\n\t}\n\tb = &bucket{\n\t\tcapacity:  capacity,\n\t\tremaining: capacity,\n\t\treset:     time.Now().Add(rate),\n\t\trate:      rate,\n\t}\n\ts.buckets[name] = b\n\treturn b, nil\n}\n"
  },
  {
    "path": "simpleratelimit/ratelimit.go",
    "content": "package simpleratelimit\n\nimport (\n\t\"sync/atomic\"\n\t\"time\"\n)\n\n// RateLimiter 限速器\ntype RateLimiter struct {\n\trate      uint64\n\tallowance uint64\n\tmax       uint64\n\tunit      uint64\n\tlastCheck uint64\n}\n\n// New 创建RateLimiter实例\nfunc New(rate int, per time.Duration) *RateLimiter {\n\tnano := uint64(per)\n\tif nano < 1 {\n\t\tnano = uint64(time.Second)\n\t}\n\tif rate < 1 {\n\t\trate = 1\n\t}\n\n\treturn &RateLimiter{\n\t\trate:      uint64(rate),\n\t\tallowance: uint64(rate) * nano,\n\t\tmax:       uint64(rate) * nano,\n\t\tunit:      nano,\n\n\t\tlastCheck: unixNano(),\n\t}\n}\n\n// Limit 判断是否超过限制\nfunc (rl *RateLimiter) Limit() bool {\n\tnow := unixNano()\n\t// 计算上一次调用到现在过了多少纳秒\n\tpassed := now - atomic.SwapUint64(&rl.lastCheck, now)\n\n\trate := atomic.LoadUint64(&rl.rate)\n\tcurrent := atomic.AddUint64(&rl.allowance, passed*rate)\n\n\tif max := atomic.LoadUint64(&rl.max); current > max {\n\t\tatomic.AddUint64(&rl.allowance, max-current)\n\t\tcurrent = max\n\t}\n\n\tif current < rl.unit {\n\t\treturn true\n\t}\n\n\t// 没有超过限额\n\tatomic.AddUint64(&rl.allowance, -rl.unit)\n\treturn false\n}\n\n// UpdateRate 更新速率值\nfunc (rl *RateLimiter) UpdateRate(rate int) {\n\tatomic.StoreUint64(&rl.rate, uint64(rate))\n\tatomic.StoreUint64(&rl.max, uint64(rate)*rl.unit)\n}\n\n// Undo 重置上一次调用Limit()，返回没有使用过的限额\nfunc (rl *RateLimiter) Undo() {\n\tcurrent := atomic.AddUint64(&rl.allowance, rl.unit)\n\n\tif max := atomic.LoadUint64(&rl.max); current > max {\n\t\tatomic.AddUint64(&rl.allowance, max-current)\n\t}\n}\n\n// unixNano 当前时间（纳秒）\nfunc unixNano() uint64 {\n\treturn uint64(time.Now().UnixNano())\n}\n"
  },
  {
    "path": "simpleratelimit/ratelimit_test.go",
    "content": "package simpleratelimit\n\nimport (\n\t\"testing\"\n\t\"time\"\n\n\t. \"github.com/smartystreets/goconvey/convey\"\n)\n\nfunc TestLimit(t *testing.T) {\n\tConvey(\"start testing New(0, 0)\", t, func() {\n\t\trl := New(0, 0*time.Second)\n\t\tSo(rl.Limit(), ShouldEqual, false)\n\t})\n\tConvey(\"start testing limit\", t, func() {\n\t\trl := New(1, time.Second)\n\t\tfor i := 0; i < 10; i++ {\n\t\t\tif i == 0 {\n\t\t\t\tSo(rl.Limit(), ShouldEqual, false)\n\t\t\t} else {\n\t\t\t\tSo(rl.Limit(), ShouldEqual, true)\n\t\t\t}\n\t\t}\n\t})\n\tConvey(\"start testing updateRate\", t, func() {\n\t\trl := New(1, time.Second)\n\t\tfor i := 0; i < 2; i++ {\n\t\t\tif i == 0 {\n\t\t\t\tSo(rl.Limit(), ShouldEqual, false)\n\t\t\t} else {\n\t\t\t\trl.UpdateRate(2)\n\t\t\t\tSo(rl.Limit(), ShouldEqual, true)\n\t\t\t}\n\t\t}\n\t})\n\tConvey(\"start testing undo\", t, func() {\n\t\trl := New(1, time.Second)\n\t\tfor i := 0; i < 10; i++ {\n\t\t\tif i == 0 {\n\t\t\t\tSo(rl.Limit(), ShouldEqual, false)\n\t\t\t} else {\n\t\t\t\trl.Undo()\n\t\t\t\tSo(rl.Limit(), ShouldEqual, false)\n\t\t\t}\n\t\t}\n\t\tfor i := 0; i < 10; i++ {\n\t\t\trl.Limit()\n\t\t}\n\t\trl.Undo()\n\t})\n}\n\nfunc BenchmarkLimit(b *testing.B) {\n\tConvey(\"start testing benchmark\", b, func() {\n\t\trl := New(1, time.Second)\n\t\tfor i := 0; i < b.N; i++ {\n\t\t\trl.Limit()\n\t\t}\n\t})\n}\n"
  }
]