add new file

lib/rate/rate.go

@@ -0,0 +1,111 @@
+package rate
+
+import (
+	"errors"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// Rate is an implementation of the token bucket added regularly
+type Rate struct {
+	bucketSize        int64
+	bucketSurplusSize int64
+	bucketAddSize     int64
+	stopChan          chan bool
+	nowRate           int64
+	cond              *sync.Cond
+	hasStop           bool
+	hasStart          bool
+}
+
+// NewRate return token bucket with specified rate
+func NewRate(addSize int64) *Rate {
+	r := &Rate{
+		bucketSize:        addSize * 2,
+		bucketSurplusSize: 0,
+		bucketAddSize:     addSize,
+		stopChan:          make(chan bool),
+		cond:              sync.NewCond(new(sync.Mutex)),
+	}
+	return r
+}
+
+// Start is used to add token regularly
+func (r *Rate) Start() {
+	if !r.hasStart {
+		r.hasStart = true
+		go r.session()
+	}
+}
+
+func (r *Rate) add(size int64) {
+	if res := r.bucketSize - r.bucketSurplusSize; res < r.bucketAddSize {
+		atomic.AddInt64(&r.bucketSurplusSize, res)
+		return
+	}
+	atomic.AddInt64(&r.bucketSurplusSize, size)
+}
+
+// Write is called when add token to bucket
+func (r *Rate) Write(size int64) {
+	r.add(size)
+}
+
+// Stop is called when not use the rate bucket
+func (r *Rate) Stop() {
+	if r.hasStart {
+		r.stopChan <- true
+		r.hasStop = true
+		r.cond.Broadcast()
+	}
+}
+
+// Get is called when get token from bucket
+func (r *Rate) Get(size int64) error {
+	if r.hasStop {
+		return errors.New("the rate has closed")
+	}
+	if r.bucketSurplusSize >= size {
+		atomic.AddInt64(&r.bucketSurplusSize, -size)
+		return nil
+	}
+	for {
+		r.cond.L.Lock()
+		r.cond.Wait()
+		if r.bucketSurplusSize >= size {
+			r.cond.L.Unlock()
+			atomic.AddInt64(&r.bucketSurplusSize, -size)
+			return nil
+		}
+		if r.hasStop {
+			return errors.New("the rate has closed")
+		}
+		r.cond.L.Unlock()
+	}
+}
+
+// GetNowRate returns the current rate
+// Just a rough number
+func (r *Rate) GetNowRate() int64 {
+	return r.nowRate
+}
+
+func (r *Rate) session() {
+	ticker := time.NewTicker(time.Second * 1)
+	for {
+		select {
+		case <-ticker.C:
+			if rs := r.bucketAddSize - r.bucketSurplusSize; rs > 0 {
+				r.nowRate = rs
+			} else {
+				r.nowRate = r.bucketSize - r.bucketSurplusSize
+			}
+			r.add(r.bucketAddSize)
+			r.cond.Broadcast()
+		case <-r.stopChan:
+			ticker.Stop()
+			return
+		}
+	}
+}

lib/rate/rate_test.go

@@ -0,0 +1,52 @@
+package rate
+
+import (
+	"math"
+	"testing"
+	"time"
+)
+
+func TestRate_GetWrite(t *testing.T) {
+	r := NewRate(1024)
+	r.Write(2048)
+	n := 0
+	go func() {
+		for {
+			r.Get(1024)
+			n += 1024
+		}
+	}()
+	select {
+	case <-time.After(time.Second):
+		r.Stop()
+		if n != 2048 {
+			t.Fatal("get token error", n)
+		}
+	}
+}
+
+func TestRate_StartGetRate(t *testing.T) {
+	r := NewRate(1024)
+	r.Start()
+	n := 0
+	go func() {
+		for {
+			err := r.Get(1024)
+			if err != nil {
+				return
+			}
+			n += 1024
+		}
+	}()
+	select {
+	case <-time.After(time.Second * 5):
+		r.Stop()
+		time.Sleep(time.Second * 2)
+		if n < 4*1024 || n > 1024*5 {
+			t.Fatal("get error", n)
+		}
+		if math.Abs(float64(r.GetNowRate()-1024)) > 100 {
+			t.Fatal("rate error", n)
+		}
+	}
+}