├── .gitignore ├── COPYING ├── LICENSE ├── README ├── README.md ├── tokenbucket.go └── tokenbucket_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | *.swp 2 | *.swo 3 | *.swn 4 | conf.sh 5 | -------------------------------------------------------------------------------- /COPYING: -------------------------------------------------------------------------------- 1 | GNU LESSER GENERAL PUBLIC LICENSE 2 | Version 3, 29 June 2007 3 | 4 | Copyright (C) 2007 Free Software Foundation, Inc. 5 | Everyone is permitted to copy and distribute verbatim copies 6 | of this license document, but changing it is not allowed. 7 | 8 | 9 | This version of the GNU Lesser General Public License incorporates 10 | the terms and conditions of version 3 of the GNU General Public 11 | License, supplemented by the additional permissions listed below. 12 | 13 | 0. Additional Definitions. 14 | 15 | As used herein, "this License" refers to version 3 of the GNU Lesser 16 | General Public License, and the "GNU GPL" refers to version 3 of the GNU 17 | General Public License. 18 | 19 | "The Library" refers to a covered work governed by this License, 20 | other than an Application or a Combined Work as defined below. 21 | 22 | An "Application" is any work that makes use of an interface provided 23 | by the Library, but which is not otherwise based on the Library. 24 | Defining a subclass of a class defined by the Library is deemed a mode 25 | of using an interface provided by the Library. 26 | 27 | A "Combined Work" is a work produced by combining or linking an 28 | Application with the Library. The particular version of the Library 29 | with which the Combined Work was made is also called the "Linked 30 | Version". 31 | 32 | The "Minimal Corresponding Source" for a Combined Work means the 33 | Corresponding Source for the Combined Work, excluding any source code 34 | for portions of the Combined Work that, considered in isolation, are 35 | based on the Application, and not on the Linked Version. 36 | 37 | The "Corresponding Application Code" for a Combined Work means the 38 | object code and/or source code for the Application, including any data 39 | and utility programs needed for reproducing the Combined Work from the 40 | Application, but excluding the System Libraries of the Combined Work. 41 | 42 | 1. Exception to Section 3 of the GNU GPL. 43 | 44 | You may convey a covered work under sections 3 and 4 of this License 45 | without being bound by section 3 of the GNU GPL. 46 | 47 | 2. Conveying Modified Versions. 48 | 49 | If you modify a copy of the Library, and, in your modifications, a 50 | facility refers to a function or data to be supplied by an Application 51 | that uses the facility (other than as an argument passed when the 52 | facility is invoked), then you may convey a copy of the modified 53 | version: 54 | 55 | a) under this License, provided that you make a good faith effort to 56 | ensure that, in the event an Application does not supply the 57 | function or data, the facility still operates, and performs 58 | whatever part of its purpose remains meaningful, or 59 | 60 | b) under the GNU GPL, with none of the additional permissions of 61 | this License applicable to that copy. 62 | 63 | 3. Object Code Incorporating Material from Library Header Files. 64 | 65 | The object code form of an Application may incorporate material from 66 | a header file that is part of the Library. You may convey such object 67 | code under terms of your choice, provided that, if the incorporated 68 | material is not limited to numerical parameters, data structure 69 | layouts and accessors, or small macros, inline functions and templates 70 | (ten or fewer lines in length), you do both of the following: 71 | 72 | a) Give prominent notice with each copy of the object code that the 73 | Library is used in it and that the Library and its use are 74 | covered by this License. 75 | 76 | b) Accompany the object code with a copy of the GNU GPL and this license 77 | document. 78 | 79 | 4. Combined Works. 80 | 81 | You may convey a Combined Work under terms of your choice that, 82 | taken together, effectively do not restrict modification of the 83 | portions of the Library contained in the Combined Work and reverse 84 | engineering for debugging such modifications, if you also do each of 85 | the following: 86 | 87 | a) Give prominent notice with each copy of the Combined Work that 88 | the Library is used in it and that the Library and its use are 89 | covered by this License. 90 | 91 | b) Accompany the Combined Work with a copy of the GNU GPL and this license 92 | document. 93 | 94 | c) For a Combined Work that displays copyright notices during 95 | execution, include the copyright notice for the Library among 96 | these notices, as well as a reference directing the user to the 97 | copies of the GNU GPL and this license document. 98 | 99 | d) Do one of the following: 100 | 101 | 0) Convey the Minimal Corresponding Source under the terms of this 102 | License, and the Corresponding Application Code in a form 103 | suitable for, and under terms that permit, the user to 104 | recombine or relink the Application with a modified version of 105 | the Linked Version to produce a modified Combined Work, in the 106 | manner specified by section 6 of the GNU GPL for conveying 107 | Corresponding Source. 108 | 109 | 1) Use a suitable shared library mechanism for linking with the 110 | Library. A suitable mechanism is one that (a) uses at run time 111 | a copy of the Library already present on the user's computer 112 | system, and (b) will operate properly with a modified version 113 | of the Library that is interface-compatible with the Linked 114 | Version. 115 | 116 | e) Provide Installation Information, but only if you would otherwise 117 | be required to provide such information under section 6 of the 118 | GNU GPL, and only to the extent that such information is 119 | necessary to install and execute a modified version of the 120 | Combined Work produced by recombining or relinking the 121 | Application with a modified version of the Linked Version. (If 122 | you use option 4d0, the Installation Information must accompany 123 | the Minimal Corresponding Source and Corresponding Application 124 | Code. If you use option 4d1, you must provide the Installation 125 | Information in the manner specified by section 6 of the GNU GPL 126 | for conveying Corresponding Source.) 127 | 128 | 5. Combined Libraries. 129 | 130 | You may place library facilities that are a work based on the 131 | Library side by side in a single library together with other library 132 | facilities that are not Applications and are not covered by this 133 | License, and convey such a combined library under terms of your 134 | choice, if you do both of the following: 135 | 136 | a) Accompany the combined library with a copy of the same work based 137 | on the Library, uncombined with any other library facilities, 138 | conveyed under the terms of this License. 139 | 140 | b) Give prominent notice with the combined library that part of it 141 | is a work based on the Library, and explaining where to find the 142 | accompanying uncombined form of the same work. 143 | 144 | 6. Revised Versions of the GNU Lesser General Public License. 145 | 146 | The Free Software Foundation may publish revised and/or new versions 147 | of the GNU Lesser General Public License from time to time. Such new 148 | versions will be similar in spirit to the present version, but may 149 | differ in detail to address new problems or concerns. 150 | 151 | Each version is given a distinguishing version number. If the 152 | Library as you received it specifies that a certain numbered version 153 | of the GNU Lesser General Public License "or any later version" 154 | applies to it, you have the option of following the terms and 155 | conditions either of that published version or of any later version 156 | published by the Free Software Foundation. If the Library as you 157 | received it does not specify a version number of the GNU Lesser 158 | General Public License, you may choose any version of the GNU Lesser 159 | General Public License ever published by the Free Software Foundation. 160 | 161 | If the Library as you received it specifies that a proxy can decide 162 | whether future versions of the GNU Lesser General Public License shall 163 | apply, that proxy's public statement of acceptance of any version is 164 | permanent authorization for you to choose that version for the 165 | Library. 166 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | COPYING -------------------------------------------------------------------------------- /README: -------------------------------------------------------------------------------- 1 | [![GoDoc](http://godoc.org/github.com/ChimeraCoder/tokenbucket?status.png)](http://godoc.org/github.com/ChimeraCoder/tokenbucket) 2 | 3 | tokenbucket 4 | ==================== 5 | 6 | This package provides an implementation of [Token bucket](https://en.wikipedia.org/wiki/Token_bucket) scheduling in Go. It is useful for implementing rate-limiting, traffic shaping, or other sorts of scheduling that depend on bandwidth constraints. 7 | 8 | 9 | Example 10 | ------------ 11 | 12 | 13 | To create a new bucket, specify a capacity (how many tokens can be stored "in the bank"), and a rate (how often a new token is added). 14 | 15 | ````go 16 | 17 | // Create a new bucket 18 | // Allow a new action every 5 seconds, with a maximum of 3 "in the bank" 19 | bucket := tokenbucket.NewBucket(3, 5 * time.Second) 20 | ```` 21 | 22 | This bucket should be shared between any functions that share the same constraints. (These functions may or may not run in separate goroutines). 23 | 24 | 25 | Anytime a regulated action is performed, spend a token. 26 | 27 | ````go 28 | // To perform a regulated action, we must spend a token 29 | // RegulatedAction will not be performed until the bucket contains enough tokens 30 | <-bucket.SpendToken(1) 31 | RegulatedAction() 32 | ```` 33 | 34 | `SpendToken` returns immediately. Reading from the channel that it returns will block until the action has "permission" to continue (ie, until there are enough tokens in the bucket). 35 | 36 | 37 | (The channel that `SpendToken` returns is of type `error`. For now, the value will always be `nil`, so it can be ignored.) 38 | 39 | 40 | 41 | ####License 42 | 43 | `tokenbucket` is free software provided under version 3 of the LGPL license. 44 | 45 | 46 | Software that uses `tokenbucket` may be released under *any* license, as long as the source code for `tokenbucket` (including any modifications) are made available under the LGPLv3 license. 47 | 48 | You do not need to release the rest of the software under the LGPL, or any free/open-source license, for that matter (though we would encourage you to do so!). 49 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | README -------------------------------------------------------------------------------- /tokenbucket.go: -------------------------------------------------------------------------------- 1 | package tokenbucket 2 | 3 | import ( 4 | "sync" 5 | "time" 6 | ) 7 | 8 | type Bucket struct { 9 | capacity int64 10 | tokens chan struct{} 11 | rate time.Duration // Add a token to the bucket every 1/r units of time 12 | rateMutex sync.Mutex 13 | } 14 | 15 | func NewBucket(rate time.Duration, capacity int64) *Bucket { 16 | 17 | //A bucket is simply a channel with a buffer representing the maximum size 18 | tokens := make(chan struct{}, capacity) 19 | 20 | b := &Bucket{capacity, tokens, rate, sync.Mutex{}} 21 | 22 | //Set off a function that will continuously add tokens to the bucket 23 | go func(b *Bucket) { 24 | ticker := time.NewTicker(rate) 25 | for _ = range ticker.C { 26 | b.tokens <- struct{}{} 27 | } 28 | }(b) 29 | 30 | return b 31 | } 32 | 33 | func (b *Bucket) GetRate() time.Duration { 34 | b.rateMutex.Lock() 35 | tmp := b.rate 36 | b.rateMutex.Unlock() 37 | return tmp 38 | } 39 | 40 | func (b *Bucket) SetRate(rate time.Duration) { 41 | b.rateMutex.Lock() 42 | b.rate = rate 43 | b.rateMutex.Unlock() 44 | } 45 | 46 | //AddTokens manually adds n tokens to the bucket 47 | func (b *Bucket) AddToken(n int64) { 48 | } 49 | 50 | func (b *Bucket) withdrawTokens(n int64) error { 51 | for i := int64(0); i < n; i++ { 52 | <-b.tokens 53 | } 54 | return nil 55 | } 56 | 57 | func (b *Bucket) SpendToken(n int64) <-chan error { 58 | // Default to spending a single token 59 | if n < 0 { 60 | n = 1 61 | } 62 | 63 | c := make(chan error) 64 | go func(b *Bucket, n int64, c chan error) { 65 | c <- b.withdrawTokens(n) 66 | close(c) 67 | return 68 | }(b, n, c) 69 | 70 | return c 71 | } 72 | 73 | // Drain will empty all tokens in the bucket 74 | // If the tokens are being added too quickly (if the rate is too fast) 75 | // this will never drain 76 | func (b *Bucket) Drain() error{ 77 | // TODO replace this with a more solid approach (such as replacing the channel altogether) 78 | for { 79 | select { 80 | case _ = <-b.tokens: 81 | continue 82 | default: 83 | return nil 84 | } 85 | } 86 | } 87 | -------------------------------------------------------------------------------- /tokenbucket_test.go: -------------------------------------------------------------------------------- 1 | package tokenbucket_test 2 | 3 | import ( 4 | "github.com/ChimeraCoder/tokenbucket" 5 | "testing" 6 | "time" 7 | ) 8 | 9 | func Example_BucketUse() { 10 | // Allow a new action every 5 seconds, with a maximum of 3 "in the bank" 11 | bucket := tokenbucket.NewBucket(5*time.Second, 3) 12 | 13 | // To perform a regulated action, we must spend a token 14 | // RegulatedAction will not be performed until the bucket contains enough tokens 15 | <-bucket.SpendToken(1) 16 | RegulatedAction() 17 | } 18 | 19 | // RegulatedAction represents some function that is rate-limited, monitored, 20 | // or otherwise regulated 21 | func RegulatedAction() { 22 | // Some expensive action goes on here 23 | } 24 | 25 | // Test that a bucket that is full does not block execution 26 | func Test_BucketBuffering(t *testing.T) { 27 | // Create a bucket with capacity 3, that adds tokens every 4 seconds 28 | const RATE = 4 * time.Second 29 | const CAPACITY = 3 30 | const ERROR = 500 * time.Millisecond 31 | b := tokenbucket.NewBucket(RATE, CAPACITY) 32 | 33 | // Allow the bucket enough time to fill to capacity 34 | time.Sleep(CAPACITY * RATE) 35 | 36 | // Check that we can empty the bucket without wasting any time 37 | before := time.Now() 38 | <-b.SpendToken(1) 39 | <-b.SpendToken(1) 40 | <-b.SpendToken(1) 41 | after := time.Now() 42 | 43 | if diff := after.Sub(before); diff > RATE { 44 | t.Errorf("Waited %d seconds, though this should have been nearly instantaneous", diff) 45 | } 46 | } 47 | 48 | // Test that a bucket that is empty blocks execution for the correct amount of time 49 | func Test_BucketCreation(t *testing.T) { 50 | // Create a bucket with capacity 3, that adds tokens every 4 seconds 51 | const RATE = 4 * time.Second 52 | const CAPACITY = 3 53 | const ERROR = 500 * time.Millisecond 54 | const EXPECTED_DURATION = RATE * CAPACITY 55 | 56 | b := tokenbucket.NewBucket(RATE, CAPACITY) 57 | 58 | // Ensure that the bucket is empty 59 | <-b.SpendToken(1) 60 | <-b.SpendToken(1) 61 | <-b.SpendToken(1) 62 | <-b.SpendToken(1) 63 | 64 | // Spending three times on an empty bucket should take 12 seconds 65 | // (Take the average across three, due to imprecision/scheduling) 66 | before := time.Now() 67 | <-b.SpendToken(1) 68 | <-b.SpendToken(1) 69 | <-b.SpendToken(1) 70 | after := time.Now() 71 | 72 | lower := EXPECTED_DURATION - ERROR 73 | upper := EXPECTED_DURATION + ERROR 74 | if diff := after.Sub(before); diff < lower || diff > upper { 75 | t.Errorf("Waited %s seconds, though really should have waited between %s and %s", diff.String(), lower.String(), upper.String()) 76 | } 77 | } 78 | --------------------------------------------------------------------------------