Create & Init Project...

library/rate/vegas/BUILD

@@ -0,0 +1,42 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_test",
+    "go_library",
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["vegas_test.go"],
+    embed = [":go_default_library"],
+    rundir = ".",
+    tags = ["automanaged"],
+    deps = ["//library/rate:go_default_library"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = [
+        "sample.go",
+        "vegas.go",
+    ],
+    importpath = "go-common/library/rate/vegas",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = ["//library/rate:go_default_library"],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

library/rate/vegas/sample.go

@@ -0,0 +1,46 @@
+package vegas
+
+import (
+	"sync/atomic"
+)
+
+type sample struct {
+	count       int64
+	maxInFlight int64
+	drop        int64
+	// nanoseconds
+	totalRTT int64
+}
+
+func (s *sample) Add(rtt int64, inFlight int64, drop bool) {
+	if drop {
+		atomic.StoreInt64(&s.drop, 1)
+	}
+	for max := atomic.LoadInt64(&s.maxInFlight); max < inFlight; max = atomic.LoadInt64(&s.maxInFlight) {
+		if atomic.CompareAndSwapInt64(&s.maxInFlight, max, inFlight) {
+			break
+		}
+	}
+	atomic.AddInt64(&s.totalRTT, rtt)
+	atomic.AddInt64(&s.count, 1)
+}
+
+func (s *sample) RTT() int64 {
+	count := atomic.LoadInt64(&s.count)
+	if count == 0 {
+		return 0
+	}
+	return atomic.LoadInt64(&s.totalRTT) / count
+}
+
+func (s *sample) MaxInFlight() int64 {
+	return atomic.LoadInt64(&s.maxInFlight)
+}
+
+func (s *sample) Count() int64 {
+	return atomic.LoadInt64(&s.count)
+}
+
+func (s *sample) Drop() bool {
+	return atomic.LoadInt64(&s.drop) == 1
+}

library/rate/vegas/vegas.go

@@ -0,0 +1,137 @@
+package vegas
+
+import (
+	"math"
+	"math/rand"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"go-common/library/rate"
+)
+
+const (
+	_minWindowTime = int64(time.Millisecond * 500)
+	_maxWindowTime = int64(time.Millisecond * 2000)
+
+	_minLimit = 8
+	_maxLimit = 2048
+)
+
+// Stat is the Statistics of vegas.
+type Stat struct {
+	Limit    int64
+	InFlight int64
+	MinRTT   time.Duration
+	LastRTT  time.Duration
+}
+
+// New new a rate vegas.
+func New() *Vegas {
+	v := &Vegas{
+		probes: 100,
+		limit:  _minLimit,
+	}
+	v.sample.Store(&sample{})
+	return v
+}
+
+// Vegas tcp vegas.
+type Vegas struct {
+	limit      int64
+	inFlight   int64
+	updateTime int64
+	minRTT     int64
+
+	sample atomic.Value
+	mu     sync.Mutex
+	probes int64
+}
+
+// Stat return the statistics of vegas.
+func (v *Vegas) Stat() Stat {
+	return Stat{
+		Limit:    atomic.LoadInt64(&v.limit),
+		InFlight: atomic.LoadInt64(&v.inFlight),
+		MinRTT:   time.Duration(atomic.LoadInt64(&v.minRTT)),
+		LastRTT:  time.Duration(v.sample.Load().(*sample).RTT()),
+	}
+}
+
+// Acquire No matter success or not,done() must be called at last.
+func (v *Vegas) Acquire() (done func(time.Time, rate.Op), success bool) {
+	inFlight := atomic.AddInt64(&v.inFlight, 1)
+	if inFlight <= atomic.LoadInt64(&v.limit) {
+		success = true
+	}
+
+	return func(start time.Time, op rate.Op) {
+		atomic.AddInt64(&v.inFlight, -1)
+		if op == rate.Ignore {
+			return
+		}
+		end := time.Now().UnixNano()
+		rtt := end - start.UnixNano()
+
+		s := v.sample.Load().(*sample)
+		if op == rate.Drop {
+			s.Add(rtt, inFlight, true)
+		} else if op == rate.Success {
+			s.Add(rtt, inFlight, false)
+		}
+		if end > atomic.LoadInt64(&v.updateTime) && s.Count() >= 16 {
+			v.mu.Lock()
+			defer v.mu.Unlock()
+			if v.sample.Load().(*sample) != s {
+				return
+			}
+			v.sample.Store(&sample{})
+
+			lastRTT := s.RTT()
+			if lastRTT <= 0 {
+				return
+			}
+			updateTime := end + lastRTT*5
+			if lastRTT*5 < _minWindowTime {
+				updateTime = end + _minWindowTime
+			} else if lastRTT*5 > _maxWindowTime {
+				updateTime = end + _maxWindowTime
+			}
+			atomic.StoreInt64(&v.updateTime, updateTime)
+			limit := atomic.LoadInt64(&v.limit)
+			queue := float64(limit) * (1 - float64(v.minRTT)/float64(lastRTT))
+			v.probes--
+			if v.probes <= 0 {
+				maxFlight := s.MaxInFlight()
+				if maxFlight*2 < v.limit || maxFlight <= _minLimit {
+					v.probes = 3*limit + rand.Int63n(3*limit)
+					v.minRTT = lastRTT
+				}
+			}
+			if v.minRTT == 0 || lastRTT < v.minRTT {
+				v.minRTT = lastRTT
+			}
+			var newLimit float64
+			threshold := math.Sqrt(float64(limit)) / 2
+			if s.Drop() {
+				newLimit = float64(limit) - threshold
+			} else if s.MaxInFlight()*2 < v.limit {
+				return
+			} else {
+				if queue < threshold {
+					newLimit = float64(limit) + 6*threshold
+				} else if queue < 2*threshold {
+					newLimit = float64(limit) + 3*threshold
+				} else if queue < 3*threshold {
+					newLimit = float64(limit) + threshold
+				} else if queue > 6*threshold {
+					newLimit = float64(limit) - threshold
+				} else {
+					return
+				}
+			}
+			newLimit = math.Max(_minLimit, math.Min(_maxLimit, newLimit))
+			atomic.StoreInt64(&v.limit, int64(newLimit))
+		}
+	}, success
+}

library/rate/vegas/vegas_test.go

@@ -0,0 +1,59 @@
+package vegas
+
+import (
+	"testing"
+	"time"
+
+	"go-common/library/rate"
+)
+
+func worker(qps int64, ch chan struct{}) {
+	for {
+		<-ch
+		time.Sleep(time.Duration(int64(time.Second) / qps))
+	}
+}
+
+func TestRateSuccess(t *testing.T) {
+	ch := make(chan struct{})
+	go worker(100, ch)
+	failed := producer(New(), 100, ch)
+	if failed > 0 {
+		t.Fatalf("Should be rejected 0 time,but (%d)", failed)
+	}
+}
+
+func TestRateFail(t *testing.T) {
+	ch := make(chan struct{})
+	go worker(100, ch)
+	failed := producer(New(), 200, ch)
+	if failed < 900 {
+		t.Fatalf("Should be rejected more than 900 times,but (%d)", failed)
+	}
+}
+
+func TestRateFailMuch(t *testing.T) {
+	ch := make(chan struct{})
+	go worker(10, ch)
+	failed := producer(New(), 200, ch)
+	if failed < 1600 {
+		t.Fatalf("Should be rejected more than 1600 times,but (%d)", failed)
+	}
+}
+
+func producer(v *Vegas, qps int64, ch chan struct{}) (failed int) {
+	for i := 0; i < int(qps)*10; i++ {
+		go func() {
+			start := time.Now()
+			done, success := v.Acquire()
+			defer done(start, rate.Success)
+			if success {
+				ch <- struct{}{}
+			} else {
+				failed++
+			}
+		}()
+		time.Sleep(time.Duration(int64(time.Second) / qps))
+	}
+	return
+}