├── .gitignore ├── reader.go ├── writer.go ├── LICENSE ├── ioshape.go ├── README.md ├── ioshape_test.go └── bucket.go /.gitignore: -------------------------------------------------------------------------------- 1 | *.bak 2 | *.old 3 | *~ 4 | *.cache 5 | *.swap 6 | *.swp 7 | *.temp 8 | *.tmp 9 | 10 | *.o 11 | *.a 12 | *.so 13 | *.exe 14 | *.dll 15 | *.py[co] 16 | 17 | *.DS_Store 18 | 19 | *.project 20 | *.includepath 21 | *.settings 22 | *.sublime-project 23 | *.sublime-workspace 24 | *.idea 25 | *.vscode 26 | -------------------------------------------------------------------------------- /reader.go: -------------------------------------------------------------------------------- 1 | package ioshape 2 | 3 | import ( 4 | "io" 5 | "time" 6 | ) 7 | 8 | // Reader is a traffic shaper struct that implements io.Reader interface. A 9 | // Reader reads from R by B. 10 | // Priority changes between 0(highest) and 15(lowest). 11 | type Reader struct { 12 | R io.Reader // underlying reader 13 | B *Bucket // bucket 14 | Pr int // priority 15 | } 16 | 17 | // Read reads from R by b. 18 | func (rr *Reader) Read(p []byte) (n int, err error) { 19 | if rr.B == nil { 20 | n, err = rr.R.Read(p) 21 | return 22 | } 23 | 24 | l := len(p) 25 | m := l 26 | for n < l && err == nil { 27 | k := int(rr.B.getTokens(int64(m), rr.Pr)) 28 | if k <= 0 { 29 | time.Sleep(time.Second / (freq * freqMul)) 30 | continue 31 | } 32 | var nn int 33 | nn, err = rr.R.Read(p[n : n+k]) 34 | if nn < 0 || nn > k { 35 | rr.B.giveTokens(int64(k)) 36 | err = ErrOutOfRange 37 | continue 38 | } 39 | if nn != k { 40 | rr.B.giveTokens(int64(k - nn)) 41 | } 42 | n += nn 43 | m -= nn 44 | } 45 | return 46 | } 47 | -------------------------------------------------------------------------------- /writer.go: -------------------------------------------------------------------------------- 1 | package ioshape 2 | 3 | import ( 4 | "io" 5 | "time" 6 | ) 7 | 8 | // Writer is a traffic shaper struct that implements io.Writer interface. A 9 | // Writer writes to W by B. 10 | // Priority changes between 0(highest) and 15(lowest). 11 | type Writer struct { 12 | W io.Writer // underlying reader 13 | B *Bucket // bucket 14 | Pr int // priority 15 | } 16 | 17 | // Write writes to W by b. 18 | func (wr *Writer) Write(p []byte) (n int, err error) { 19 | if wr.B == nil { 20 | n, err = wr.W.Write(p) 21 | return 22 | } 23 | 24 | l := len(p) 25 | m := l 26 | for n < l && err == nil { 27 | k := int(wr.B.getTokens(int64(m), wr.Pr)) 28 | if k <= 0 { 29 | time.Sleep(time.Second / (freq * freqMul)) 30 | continue 31 | } 32 | var nn int 33 | nn, err = wr.W.Write(p[n : n+k]) 34 | if nn < 0 || nn > k { 35 | wr.B.giveTokens(int64(k)) 36 | err = ErrOutOfRange 37 | continue 38 | } 39 | if nn != k { 40 | wr.B.giveTokens(int64(k - nn)) 41 | err = io.ErrShortWrite 42 | } 43 | n += nn 44 | m -= nn 45 | } 46 | return 47 | } 48 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2018, Orkun Karaduman. 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are met: 6 | * Redistributions of source code must retain the above copyright 7 | notice, this list of conditions and the following disclaimer. 8 | * Redistributions in binary form must reproduce the above copyright 9 | notice, this list of conditions and the following disclaimer in the 10 | documentation and/or other materials provided with the distribution. 11 | * Neither the name of Orkun Karaduman nor the 12 | names of its contributors may be used to endorse or promote products 13 | derived from this software without specific prior written permission. 14 | 15 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 16 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 17 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 18 | DISCLAIMED. IN NO EVENT SHALL BE LIABLE FOR ANY 19 | DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 20 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 21 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND 22 | ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 23 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 24 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 25 | -------------------------------------------------------------------------------- /ioshape.go: -------------------------------------------------------------------------------- 1 | /* 2 | Package ioshape provides I/O structures and functions for Traffic Shaping using 3 | token-bucket algorithm. 4 | */ 5 | package ioshape 6 | 7 | import ( 8 | "errors" 9 | "io" 10 | ) 11 | 12 | const ( 13 | freq = 16 14 | freqMul = 4 15 | priorityScale = 16 16 | chunkDiv = 1 17 | chunkSize = 32 * 1024 18 | ) 19 | 20 | // ErrOutOfRange is the error used for the result of r/w is out of range. 21 | var ErrOutOfRange = errors.New("out of range") 22 | 23 | // CopyB is identical to io.Copy except that it shapes traffic by b *Bucket. 24 | func CopyB(dst io.Writer, src io.Reader, b *Bucket) (written int64, err error) { 25 | return io.Copy(dst, &Reader{R: src, B: b}) 26 | } 27 | 28 | // CopyBN is identical to io.CopyN except that it shapes traffic by b *Bucket. 29 | func CopyBN(dst io.Writer, src io.Reader, b *Bucket, n int64) (written int64, err error) { 30 | return io.CopyN(dst, &Reader{R: src, B: b}, n) 31 | } 32 | 33 | // CopyRate is identical to io.Copy except that it shapes traffic with rate 34 | // in bytes per second. 35 | func CopyRate(dst io.Writer, src io.Reader, rate int64) (written int64, err error) { 36 | b := NewBucketRate(rate) 37 | written, err = io.Copy(dst, &Reader{R: src, B: b}) 38 | b.Stop() 39 | return 40 | } 41 | 42 | // CopyRateN is identical to io.CopyN except that it shapes traffic with rate 43 | // in bytes per second. 44 | func CopyRateN(dst io.Writer, src io.Reader, rate int64, n int64) (written int64, err error) { 45 | b := NewBucketRate(rate) 46 | written, err = io.CopyN(dst, &Reader{R: src, B: b}, n) 47 | b.Stop() 48 | return 49 | } 50 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Go Traffic Shaper 2 | 3 | [![GoDoc](https://godoc.org/github.com/orkunkaraduman/go-ioshape?status.svg)](https://godoc.org/github.com/orkunkaraduman/go-ioshape) 4 | 5 | The repository provides `ioshape` package shapes I/O traffic using 6 | token-bucket algorithm. It is used for creating bandwidth limiting applications, 7 | needing traffic limiting or throttling or prioritization. 8 | 9 | ## Examples 10 | 11 | ### Limit copy operation simply 12 | 13 | It limits copy operation to 2 MBps. 14 | 15 | ```go 16 | n, err := ioshape.CopyRate(dst, src, 2*1024*1024) 17 | ``` 18 | 19 | ### Limit multiple operations with bucket 20 | 21 | It limits two copy operation to 3MBps totally. Traffic will be balanced equally. 22 | 23 | ```go 24 | bucket := ioshape.NewBucketRate(3*1024*1024) 25 | var wg sync.WaitGroup 26 | wg.Add(1) 27 | go func() { 28 | ioshape.CopyB(dst1, src1, bucket) 29 | wg.Done() 30 | }() 31 | wg.Add(1) 32 | go func() { 33 | ioshape.CopyB(dst2, src2, bucket) 34 | wg.Done() 35 | }() 36 | wg.Wait() 37 | bucket.Stop() // its necessary to free resources 38 | ``` 39 | 40 | ### Limit multiple operations with burst and priority 41 | 42 | It limits three copy operation to 5MBps totally. Traffic will be balanced with 43 | given priorities. 44 | 45 | ```go 46 | bucket := ioshape.NewBucket() 47 | rate := 5*1024*1024 // the rate is 5MBps 48 | burst := rate*10 // the burst is ten times of the rate 49 | bucket.Set(rate, burst) 50 | var wg sync.WaitGroup 51 | wg.Add(1) 52 | go func() { 53 | rr1 := &ioshape.Reader{R: src1, B: bucket, Pr: 0} // highest priority 54 | io.Copy(dst1, rr1) 55 | wg.Done() 56 | } 57 | wg.Add(1) 58 | go func() { 59 | rr2 := &ioshape.Reader{R: src2, B: bucket, Pr: 15} // lowest priority 60 | io.Copy(dst2, rr2) 61 | wg.Done() 62 | } 63 | wg.Add(1) 64 | go func() { 65 | rr3 := &ioshape.Reader{R: src3, B: bucket, Pr: 2} // higher priority 66 | io.Copy(dst3, rr3) 67 | wg.Done() 68 | } 69 | wg.Wait() 70 | bucket.Stop() 71 | ``` 72 | -------------------------------------------------------------------------------- /ioshape_test.go: -------------------------------------------------------------------------------- 1 | package ioshape 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "io/ioutil" 7 | "net/http" 8 | "sync" 9 | "testing" 10 | "time" 11 | ) 12 | 13 | const ( 14 | testURL = "http://ipv4.download.thinkbroadband.com/1GB.zip" 15 | ) 16 | 17 | func TestReader(t *testing.T) { 18 | bu := NewBucket() 19 | bu.Set(64*1024, 0) 20 | size := 4 * 64 * 1024 21 | 22 | var wg sync.WaitGroup 23 | f := func(r io.Reader) { 24 | defer wg.Done() 25 | defer func() { 26 | if rr, ok := r.(io.Closer); ok { 27 | rr.Close() 28 | } 29 | }() 30 | start := time.Now() 31 | rr := &Reader{R: r, B: bu} 32 | _, err := io.CopyN(ioutil.Discard, rr, int64(size)) 33 | if err != nil { 34 | panic(err) 35 | } 36 | fmt.Println(time.Now().Sub(start)) 37 | } 38 | 39 | j := 4 40 | rrs := make([]io.Reader, j) 41 | for i := 0; i < j; i++ { 42 | resp, err := http.Get(testURL) 43 | if err != nil { 44 | panic(err) 45 | } 46 | rrs[i] = resp.Body 47 | wg.Add(1) 48 | } 49 | for i := 0; i < j; i++ { 50 | fmt.Println(time.Now()) 51 | go f(rrs[i]) 52 | } 53 | 54 | wg.Wait() 55 | bu.Stop() 56 | } 57 | 58 | func TestWriter(t *testing.T) { 59 | bu := NewBucket() 60 | bu.Set(64*1024, 0) 61 | size := 4 * 64 * 1024 62 | 63 | var wg sync.WaitGroup 64 | f := func(r io.Reader) { 65 | defer wg.Done() 66 | defer func() { 67 | if wr, ok := r.(io.Closer); ok { 68 | wr.Close() 69 | } 70 | }() 71 | start := time.Now() 72 | wr := &Writer{W: ioutil.Discard, B: bu} 73 | _, err := io.CopyN(wr, r, int64(size)) 74 | if err != nil { 75 | panic(err) 76 | } 77 | fmt.Println(time.Now().Sub(start)) 78 | } 79 | 80 | j := 4 81 | rrs := make([]io.Reader, j) 82 | for i := 0; i < j; i++ { 83 | resp, err := http.Get(testURL) 84 | if err != nil { 85 | panic(err) 86 | } 87 | rrs[i] = resp.Body 88 | wg.Add(1) 89 | } 90 | for i := 0; i < j; i++ { 91 | fmt.Println(time.Now()) 92 | go f(rrs[i]) 93 | } 94 | 95 | wg.Wait() 96 | bu.Stop() 97 | } 98 | 99 | func TestStopping(t *testing.T) { 100 | bu := NewBucket() 101 | bu.Set(128*1024, 0) 102 | size := 4 * 128 * 1024 103 | 104 | var wg sync.WaitGroup 105 | f := func(r io.Reader) { 106 | defer wg.Done() 107 | defer func() { 108 | if rr, ok := r.(io.Closer); ok { 109 | rr.Close() 110 | } 111 | }() 112 | start := time.Now() 113 | rr := &Reader{R: r, B: bu} 114 | _, err := io.CopyN(ioutil.Discard, rr, int64(size)) 115 | if err != nil { 116 | panic(err) 117 | } 118 | fmt.Println(time.Now().Sub(start)) 119 | } 120 | 121 | j := 4 122 | rrs := make([]io.Reader, j) 123 | for i := 0; i < j; i++ { 124 | resp, err := http.Get(testURL) 125 | if err != nil { 126 | panic(err) 127 | } 128 | rrs[i] = resp.Body 129 | wg.Add(1) 130 | } 131 | for i := 0; i < j; i++ { 132 | fmt.Println(time.Now()) 133 | go f(rrs[i]) 134 | } 135 | time.Sleep(8 * time.Second) 136 | bu.Stop() 137 | 138 | wg.Wait() 139 | } 140 | -------------------------------------------------------------------------------- /bucket.go: -------------------------------------------------------------------------------- 1 | package ioshape 2 | 3 | import ( 4 | "sync" 5 | "sync/atomic" 6 | "time" 7 | ) 8 | 9 | type bucketTokenRequest struct { 10 | count int64 11 | callback chan int64 12 | priority int 13 | } 14 | 15 | type bucketTokenReturn struct { 16 | count int64 17 | } 18 | 19 | // Bucket shapes traffic by given rate, burst and Reader/Writer priorities. 20 | type Bucket struct { 21 | tokens int64 22 | n int64 23 | k int64 24 | b int64 25 | m int64 26 | setMu sync.RWMutex 27 | ticker *time.Ticker 28 | ticks int64 29 | stopCh chan struct{} 30 | stopped int32 31 | tokenRequests chan *bucketTokenRequest 32 | tokenReturns chan *bucketTokenReturn 33 | } 34 | 35 | // NewBucket returns a new Bucket. 36 | func NewBucket() (bu *Bucket) { 37 | bu = &Bucket{} 38 | bu.ticker = time.NewTicker(time.Second / freq) 39 | bu.stopCh = make(chan struct{}, 1) 40 | bu.tokenRequests = make(chan *bucketTokenRequest) 41 | bu.tokenReturns = make(chan *bucketTokenReturn) 42 | go bu.timer() 43 | return 44 | } 45 | 46 | // NewBucketRate returns a new Bucket and sets rate. 47 | func NewBucketRate(rate int64) (bu *Bucket) { 48 | bu = NewBucket() 49 | bu.SetRate(rate) 50 | return 51 | } 52 | 53 | func (bu *Bucket) timer() { 54 | var n, k, b, m int64 55 | for { 56 | select { 57 | case <-bu.stopCh: 58 | atomic.StoreInt32(&bu.stopped, 1) 59 | time.Sleep(10 * time.Millisecond) 60 | for ok := true; ok; { 61 | select { 62 | case tokenRequest := <-bu.tokenRequests: 63 | tokenRequest.callback <- tokenRequest.count 64 | default: 65 | ok = false 66 | } 67 | } 68 | return 69 | case <-bu.ticker.C: 70 | bu.setMu.RLock() 71 | n = bu.n 72 | k = bu.k 73 | b = bu.b 74 | m = bu.m 75 | bu.setMu.RUnlock() 76 | bu.tokens += n 77 | if bu.ticks%freq < k { 78 | bu.tokens++ 79 | } 80 | if bu.tokens > b { 81 | bu.tokens = b 82 | } 83 | bu.ticks++ 84 | if bu.ticks > freq { 85 | bu.ticks = 0 86 | } 87 | case tokenRequest := <-bu.tokenRequests: 88 | count := tokenRequest.count 89 | if count > bu.tokens { 90 | count = bu.tokens 91 | } 92 | if count > m { 93 | count = m 94 | } 95 | if tokenRequest.priority > int((priorityScale*bu.ticks/freq)%priorityScale) { 96 | count = 0 97 | } 98 | tokenRequest.callback <- count 99 | bu.tokens -= count 100 | case tokenReturn := <-bu.tokenReturns: 101 | count := tokenReturn.count 102 | bu.tokens += count 103 | if bu.tokens > b { 104 | bu.tokens = b 105 | } 106 | } 107 | } 108 | } 109 | 110 | // Stop turns off a bucket. After Stop, bucket won't shape traffic. Stop 111 | // must be call to free resources, after the bucket doesn't be needing. 112 | func (bu *Bucket) Stop() { 113 | bu.ticker.Stop() 114 | select { 115 | case bu.stopCh <- struct{}{}: 116 | default: 117 | } 118 | } 119 | 120 | // Set sets buckets rate and burst in bytes per second. The burst should be 121 | // greater or equal than the rate. Otherwise burst will be equal rate. 122 | func (bu *Bucket) Set(rate, burst int64) { 123 | if rate < 0 { 124 | return 125 | } 126 | bu.setMu.Lock() 127 | defer bu.setMu.Unlock() 128 | if rate > burst { 129 | burst = rate 130 | } 131 | bu.n = rate / freq 132 | bu.k = rate % freq 133 | bu.b = burst + bu.n 134 | bu.m = bu.n / chunkDiv 135 | if bu.m == 0 { 136 | bu.m = 1 137 | } 138 | } 139 | 140 | // SetRate sets rate and burst to the rate in bytes per second. 141 | func (bu *Bucket) SetRate(rate int64) { 142 | bu.Set(rate, 0) 143 | } 144 | 145 | func (bu *Bucket) getTokens(count int64, priority int) int64 { 146 | callback := make(chan int64) 147 | if count > 0 && bu.stopped == 0 { 148 | bu.tokenRequests <- &bucketTokenRequest{ 149 | count: count, 150 | callback: callback, 151 | priority: priority} 152 | return <-callback 153 | } 154 | return count 155 | } 156 | 157 | func (bu *Bucket) giveTokens(count int64) { 158 | bu.tokenReturns <- &bucketTokenReturn{ 159 | count: count} 160 | } 161 | --------------------------------------------------------------------------------