Switch to using the dep tool and update all the dependencies
This commit is contained in:
Nick Craig-Wood
+191
@@ -0,0 +1,191 @@
|
||||
package tb
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"runtime"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestNewBucket(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
b := NewBucket(10, -1)
|
||||
b.Take(10)
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
if w, g := int64(0), b.Take(1); w != g {
|
||||
t.Fatal("Expected no filling when freq == -1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBucket_Take_single(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
b := NewBucket(10, 0)
|
||||
defer b.Close()
|
||||
|
||||
ex := [...]int64{5, 5, 1, 1, 5, 4, 1, 0}
|
||||
for i := 0; i < len(ex)-1; i += 2 {
|
||||
if got, want := b.Take(ex[i]), ex[i+1]; got != want {
|
||||
t.Errorf("Want: %d, Got: %d", want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBucket_Put_single(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
b := NewBucket(10, 0)
|
||||
defer b.Close()
|
||||
|
||||
b.Take(10)
|
||||
|
||||
ex := [...]int64{5, 5, 10, 5, 15, 0}
|
||||
for i := 0; i < len(ex)-1; i += 2 {
|
||||
if got, want := b.Put(ex[i]), ex[i+1]; got != want {
|
||||
t.Errorf("Want: %d, Got: %d", want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBucket_Take_multi(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
b := NewBucket(10, 0)
|
||||
defer b.Close()
|
||||
|
||||
exs := [2][]int64{{4, 4, 2, 2}, {2, 2, 1, 1}}
|
||||
for i := 0; i < 2; i++ {
|
||||
go func(i int) {
|
||||
for j := 0; j < len(exs[i])-1; j += 2 {
|
||||
if got, want := b.Take(exs[i][j]), exs[i][j+1]; got != want {
|
||||
t.Errorf("Want: %d, Got: %d", want, got)
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBucket_Put_multi(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
b := NewBucket(10, 0)
|
||||
defer b.Close()
|
||||
|
||||
b.Take(10)
|
||||
|
||||
exs := [2][]int64{{4, 4, 2, 2}, {2, 2, 1, 1}}
|
||||
for i := 0; i < 2; i++ {
|
||||
go func(i int) {
|
||||
for j := 0; j < len(exs[i])-1; j += 2 {
|
||||
if got, want := b.Put(exs[i][j]), exs[i][j+1]; got != want {
|
||||
t.Errorf("Want: %d, Got: %d", want, got)
|
||||
}
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBucket_Take_throughput(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
if testing.Short() {
|
||||
t.Skip("Skipping test in short mode")
|
||||
}
|
||||
|
||||
runtime.GOMAXPROCS(2)
|
||||
|
||||
b := NewBucket(1000, 0)
|
||||
defer b.Close()
|
||||
|
||||
b.Take(1000)
|
||||
|
||||
var (
|
||||
out int64
|
||||
began = time.Now()
|
||||
)
|
||||
for out < 1000 {
|
||||
out += b.Take(1000 - out)
|
||||
}
|
||||
|
||||
ended := time.Since(began)
|
||||
if int(ended.Seconds()) != 1 {
|
||||
t.Errorf("Want 1000 tokens to take 1s. Got: %d", int(ended.Seconds()))
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkBucket_Take_sequential(b *testing.B) {
|
||||
bucket := NewBucket(int64(b.N), 0)
|
||||
defer bucket.Close()
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
bucket.Take(8)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkBucket_Put_sequential(b *testing.B) {
|
||||
bucket := NewBucket(int64(b.N), 0)
|
||||
defer bucket.Close()
|
||||
|
||||
bucket.Take(int64(b.N))
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
bucket.Put(8)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBucket_Wait(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := map[*Bucket]time.Duration{
|
||||
NewBucket(250, 100*time.Millisecond): 7 * time.Second,
|
||||
NewBucket(500, 100*time.Millisecond): 3 * time.Second,
|
||||
NewBucket(1e3, 500*time.Millisecond): 1 * time.Second,
|
||||
NewBucket(1e3, 20*time.Millisecond): 1 * time.Second,
|
||||
NewBucket(1e3, 1*time.Millisecond): 1 * time.Second,
|
||||
NewBucket(1e3, 0): 1 * time.Second,
|
||||
NewBucket(2e3, 0): 0,
|
||||
NewBucket(3e3, 0): 0,
|
||||
}
|
||||
errors := make(chan error, len(cases))
|
||||
|
||||
for bucket, wait := range cases {
|
||||
go func(bucket *Bucket, wait time.Duration) {
|
||||
defer bucket.Close()
|
||||
|
||||
start := time.Now()
|
||||
got := bucket.Wait(2000)
|
||||
took := time.Since(start)
|
||||
if int(wait.Seconds()) != int(got.Seconds()) {
|
||||
errors <- fmt.Errorf("bucket.Wait(2000) with cap=%d, freq=%s: Want: %s, Got %s",
|
||||
bucket.capacity, bucket.freq, wait, got)
|
||||
} else if took < wait-time.Second || took > wait+time.Second {
|
||||
// took is the actual time the bucket.Wait() took
|
||||
// wait is the time we expected it to take
|
||||
// if took is more than 1 second different from wait, then return an error
|
||||
errors <- fmt.Errorf("bucket.Wait(2000) with cap=%d, freq=%s: Waited for %v which isn't within 1 second of %v",
|
||||
bucket.capacity, bucket.freq, wait, took)
|
||||
} else {
|
||||
errors <- nil
|
||||
}
|
||||
}(bucket, wait)
|
||||
}
|
||||
|
||||
for i := 0; i < cap(errors); i++ {
|
||||
if err := <-errors; err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBucket_Close(t *testing.T) {
|
||||
b := NewBucket(10000, 0)
|
||||
b.Close()
|
||||
b.Take(10000)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
if want, got := int64(0), b.Take(1); want != got {
|
||||
t.Errorf("Want: %d Got: %d", want, got)
|
||||
}
|
||||
}
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
package examples
|
||||
|
||||
import (
|
||||
"github.com/tsenart/tb"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"time"
|
||||
)
|
||||
|
||||
func main() {
|
||||
ln, err := net.Listen("tcp", ":6789")
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
th := tb.NewThrottler(100 * time.Millisecond)
|
||||
|
||||
echo := func(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
|
||||
host, port, err := net.SplitHostPort(conn.RemoteAddr().String())
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
// Throttle to 10 connection per second from the same host
|
||||
// Handle non-conformity by dropping the connection
|
||||
if th.Halt(host, 1, 10) {
|
||||
log.Printf("Throttled %s", host)
|
||||
return
|
||||
}
|
||||
log.Printf("Echoing payload from %s:%s", host, port)
|
||||
io.Copy(conn, conn)
|
||||
}
|
||||
|
||||
for {
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
go echo(conn)
|
||||
}
|
||||
}
|
||||
+45
@@ -0,0 +1,45 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"github.com/tsenart/tb"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
type roundTripperFunc func(r *http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) {
|
||||
return f(r)
|
||||
}
|
||||
|
||||
// ByteThrottledRoundTripper wraps another RoundTripper rt,
|
||||
// throttling all requests to the specified byte rate.
|
||||
func ByteThrottledRoundTripper(rt http.RoundTripper, rate int64) http.RoundTripper {
|
||||
freq := time.Duration(1 * time.Millisecond)
|
||||
bucket := tb.NewBucket(rate, freq)
|
||||
|
||||
return roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
got := bucket.Take(r.ContentLength)
|
||||
for got < r.ContentLength {
|
||||
got += bucket.Take(r.ContentLength - got)
|
||||
time.Sleep(freq)
|
||||
}
|
||||
return rt.RoundTrip(r)
|
||||
})
|
||||
}
|
||||
|
||||
// ReqThrottledRoundTripper wraps another RoundTripper rt,
|
||||
// throttling all requests to the specified request rate.
|
||||
func ReqThrottledRoundTripper(rt http.RoundTripper, rate int64) http.RoundTripper {
|
||||
freq := time.Duration(1e9 / rate)
|
||||
bucket := tb.NewBucket(rate, freq)
|
||||
|
||||
return roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
got := bucket.Take(1)
|
||||
for got != 1 {
|
||||
got = bucket.Take(1)
|
||||
time.Sleep(freq)
|
||||
}
|
||||
return rt.RoundTrip(r)
|
||||
})
|
||||
}
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"github.com/tsenart/tb"
|
||||
"net"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
var byteThrottler = tb.NewThrottler(25 * time.Millisecond)
|
||||
|
||||
// ByteThrottledHandler wraps an http.Handler with per host byte throttling to
|
||||
// the specified byte rate, responding with 429 when throttled.
|
||||
func ByteThrottledHandler(h http.Handler, rate int64) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
host, _, _ := net.SplitHostPort(r.RemoteAddr)
|
||||
if byteThrottler.Halt(host, r.ContentLength, rate) {
|
||||
http.Error(w, "Too many requests", 429)
|
||||
return
|
||||
}
|
||||
h.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
var reqThrottler = tb.NewThrottler(5 * time.Millisecond)
|
||||
|
||||
// ReqThrottledHandler wraps an http.Handler with per host request throttling
|
||||
// to the specified request rate, responding with 429 when throttled.
|
||||
func ReqThrottledHandler(h http.Handler, rate int64) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
host, _, _ := net.SplitHostPort(r.RemoteAddr)
|
||||
if reqThrottler.Halt(host, 1, rate) {
|
||||
http.Error(w, "Too many requests", 429)
|
||||
return
|
||||
}
|
||||
h.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
package io
|
||||
|
||||
import (
|
||||
"github.com/tsenart/tb"
|
||||
"io"
|
||||
"time"
|
||||
)
|
||||
|
||||
// NewThrottledWriter is an io.Writer wrapping another io.Writer with
|
||||
// byte rate throttling, flushing block bytes at a time.
|
||||
func NewThrottledWriter(rate, block int64, w io.Writer) io.Writer {
|
||||
return &throttledWriter{rate, block, w, tb.NewBucket(rate, -1)}
|
||||
}
|
||||
|
||||
type throttledWriter struct {
|
||||
rate, block int64
|
||||
w io.Writer
|
||||
b *tb.Bucket
|
||||
}
|
||||
|
||||
func (tw *throttledWriter) Write(p []byte) (n int, err error) {
|
||||
for wr := 0; wr < len(p); {
|
||||
var got int64
|
||||
for got < tw.block {
|
||||
if got += tw.b.Take(tw.block - got); got != tw.block {
|
||||
time.Sleep(time.Duration((1e9 / tw.rate) * (tw.block - got)))
|
||||
}
|
||||
}
|
||||
if n, err = tw.w.Write(p[wr : wr+int(got)]); err != nil {
|
||||
return wr, err
|
||||
}
|
||||
wr += n
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
package tb
|
||||
|
||||
import (
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestThrottler_Bucket(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
th := NewThrottler(0)
|
||||
defer th.Close()
|
||||
|
||||
b := th.Bucket("a", 1000)
|
||||
|
||||
ex := [...]int64{100, 100, 1000, 900, 1, 0}
|
||||
for i := 0; i < len(ex)-1; i += 2 {
|
||||
if got, want := b.Take(ex[i]), ex[i+1]; got != want {
|
||||
t.Errorf("Want: %d, Got: %d", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < len(ex)-1; i += 2 {
|
||||
if got, want := b.Put(ex[i]), ex[i+1]; got != want {
|
||||
t.Errorf("Want: %d, Got: %d", want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestThrottler_Halt(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
th := NewThrottler(0)
|
||||
defer th.Close()
|
||||
|
||||
if th.Halt("a", 1000, 1000) {
|
||||
t.Fatal("Didn't expect halt")
|
||||
}
|
||||
|
||||
if !th.Halt("a", 1, 1000) {
|
||||
t.Fatal("Expected halt")
|
||||
}
|
||||
|
||||
if th.Halt("b", 1000, 1000) {
|
||||
t.Fatal("Didn't expect halt")
|
||||
}
|
||||
}
|
||||
|
||||
func TestThrottler_Wait(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
th := NewThrottler(1 * time.Millisecond)
|
||||
defer th.Close()
|
||||
|
||||
if wait := th.Wait("a", 1000, 1000); wait > 0 {
|
||||
t.Fatal("Didn't expect wait")
|
||||
}
|
||||
|
||||
if wait := th.Wait("a", 2000, 1000); int(wait.Seconds()) != 2 {
|
||||
t.Fatalf("Expected wait of 2s. Got: %s", wait)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkThrottler_Bucket(b *testing.B) {
|
||||
keys := make([]string, 10000)
|
||||
for i := 0; i < len(keys); i++ {
|
||||
keys[i] = strconv.Itoa(i)
|
||||
}
|
||||
|
||||
th := NewThrottler(1 * time.Millisecond)
|
||||
defer th.Close()
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
th.Bucket(keys[i%(len(keys)-1)], 1000)
|
||||
}
|
||||
}
|
||||
|
||||
func ExampleThrottler() {
|
||||
ln, err := net.Listen("tcp", ":6789")
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
th := NewThrottler(100 * time.Millisecond)
|
||||
defer th.Close()
|
||||
|
||||
echo := func(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
|
||||
host, port, err := net.SplitHostPort(conn.RemoteAddr().String())
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
// Throttle to 10 connection per second from the same host
|
||||
// Handle non-conformity by dropping the connection
|
||||
if th.Halt(host, 1, 10) {
|
||||
log.Printf("Throttled %s", host)
|
||||
return
|
||||
}
|
||||
log.Printf("Echoing payload from %s:%s", host, port)
|
||||
io.Copy(conn, conn)
|
||||
}
|
||||
|
||||
for {
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
go echo(conn)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user