initial import

core/load/adaptiveshedder.go

@@ -0,0 +1,248 @@
+package load
+
+import (
+	"errors"
+	"fmt"
+	"math"
+	"sync/atomic"
+	"time"
+
+	"zero/core/collection"
+	"zero/core/logx"
+	"zero/core/stat"
+	"zero/core/syncx"
+	"zero/core/timex"
+)
+
+const (
+	defaultBuckets = 50
+	defaultWindow  = time.Second * 5
+	// using 1000m notation, 900m is like 80%, keep it as var for unit test
+	defaultCpuThreshold = 900
+	defaultMinRt        = float64(time.Second / time.Millisecond)
+	// moving average hyperparameter beta for calculating requests on the fly
+	flyingBeta      = 0.9
+	coolOffDuration = time.Second
+)
+
+var (
+	ErrServiceOverloaded = errors.New("service overloaded")
+
+	// default to be enabled
+	enabled = syncx.ForAtomicBool(true)
+	// make it a variable for unit test
+	systemOverloadChecker = func(cpuThreshold int64) bool {
+		return stat.CpuUsage() >= cpuThreshold
+	}
+)
+
+type (
+	Promise interface {
+		Pass()
+		Fail()
+	}
+
+	Shedder interface {
+		Allow() (Promise, error)
+	}
+
+	ShedderOption func(opts *shedderOptions)
+
+	shedderOptions struct {
+		window       time.Duration
+		buckets      int
+		cpuThreshold int64
+	}
+
+	adaptiveShedder struct {
+		cpuThreshold    int64
+		windows         int64
+		flying          int64
+		avgFlying       float64
+		avgFlyingLock   syncx.SpinLock
+		dropTime        *syncx.AtomicDuration
+		droppedRecently *syncx.AtomicBool
+		passCounter     *collection.RollingWindow
+		rtCounter       *collection.RollingWindow
+	}
+)
+
+func Disable() {
+	enabled.Set(false)
+}
+
+func NewAdaptiveShedder(opts ...ShedderOption) Shedder {
+	if !enabled.True() {
+		return newNopShedder()
+	}
+
+	options := shedderOptions{
+		window:       defaultWindow,
+		buckets:      defaultBuckets,
+		cpuThreshold: defaultCpuThreshold,
+	}
+	for _, opt := range opts {
+		opt(&options)
+	}
+	bucketDuration := options.window / time.Duration(options.buckets)
+	return &adaptiveShedder{
+		cpuThreshold:    options.cpuThreshold,
+		windows:         int64(time.Second / bucketDuration),
+		dropTime:        syncx.NewAtomicDuration(),
+		droppedRecently: syncx.NewAtomicBool(),
+		passCounter: collection.NewRollingWindow(options.buckets, bucketDuration,
+			collection.IgnoreCurrentBucket()),
+		rtCounter: collection.NewRollingWindow(options.buckets, bucketDuration,
+			collection.IgnoreCurrentBucket()),
+	}
+}
+
+func (as *adaptiveShedder) Allow() (Promise, error) {
+	if as.shouldDrop() {
+		as.dropTime.Set(timex.Now())
+		as.droppedRecently.Set(true)
+
+		return nil, ErrServiceOverloaded
+	}
+
+	as.addFlying(1)
+
+	return &promise{
+		start:   timex.Now(),
+		shedder: as,
+	}, nil
+}
+
+func (as *adaptiveShedder) addFlying(delta int64) {
+	flying := atomic.AddInt64(&as.flying, delta)
+	// update avgFlying when the request is finished.
+	// this strategy makes avgFlying have a little bit lag against flying, and smoother.
+	// when the flying requests increase rapidly, avgFlying increase slower, accept more requests.
+	// when the flying requests drop rapidly, avgFlying drop slower, accept less requests.
+	// it makes the service to serve as more requests as possible.
+	if delta < 0 {
+		as.avgFlyingLock.Lock()
+		as.avgFlying = as.avgFlying*flyingBeta + float64(flying)*(1-flyingBeta)
+		as.avgFlyingLock.Unlock()
+	}
+}
+
+func (as *adaptiveShedder) highThru() bool {
+	as.avgFlyingLock.Lock()
+	avgFlying := as.avgFlying
+	as.avgFlyingLock.Unlock()
+	maxFlight := as.maxFlight()
+	return int64(avgFlying) > maxFlight && atomic.LoadInt64(&as.flying) > maxFlight
+}
+
+func (as *adaptiveShedder) maxFlight() int64 {
+	// windows = buckets per second
+	// maxQPS = maxPASS * windows
+	// minRT = min average response time in milliseconds
+	// maxQPS * minRT / milliseconds_per_second
+	return int64(math.Max(1, float64(as.maxPass()*as.windows)*(as.minRt()/1e3)))
+}
+
+func (as *adaptiveShedder) maxPass() int64 {
+	var result float64 = 1
+
+	as.passCounter.Reduce(func(b *collection.Bucket) {
+		if b.Sum > result {
+			result = b.Sum
+		}
+	})
+
+	return int64(result)
+}
+
+func (as *adaptiveShedder) minRt() float64 {
+	var result = defaultMinRt
+
+	as.rtCounter.Reduce(func(b *collection.Bucket) {
+		if b.Count <= 0 {
+			return
+		}
+
+		avg := math.Round(b.Sum / float64(b.Count))
+		if avg < result {
+			result = avg
+		}
+	})
+
+	return result
+}
+
+func (as *adaptiveShedder) shouldDrop() bool {
+	if as.systemOverloaded() || as.stillHot() {
+		if as.highThru() {
+			flying := atomic.LoadInt64(&as.flying)
+			as.avgFlyingLock.Lock()
+			avgFlying := as.avgFlying
+			as.avgFlyingLock.Unlock()
+			msg := fmt.Sprintf(
+				"dropreq, cpu: %d, maxPass: %d, minRt: %.2f, hot: %t, flying: %d, avgFlying: %.2f",
+				stat.CpuUsage(), as.maxPass(), as.minRt(), as.stillHot(), flying, avgFlying)
+			logx.Error(msg)
+			stat.Report(msg)
+			return true
+		}
+	}
+
+	return false
+}
+
+func (as *adaptiveShedder) stillHot() bool {
+	if !as.droppedRecently.True() {
+		return false
+	}
+
+	dropTime := as.dropTime.Load()
+	if dropTime == 0 {
+		return false
+	}
+
+	hot := timex.Since(dropTime) < coolOffDuration
+	if !hot {
+		as.droppedRecently.Set(false)
+	}
+
+	return hot
+}
+
+func (as *adaptiveShedder) systemOverloaded() bool {
+	return systemOverloadChecker(as.cpuThreshold)
+}
+
+func WithBuckets(buckets int) ShedderOption {
+	return func(opts *shedderOptions) {
+		opts.buckets = buckets
+	}
+}
+
+func WithCpuThreshold(threshold int64) ShedderOption {
+	return func(opts *shedderOptions) {
+		opts.cpuThreshold = threshold
+	}
+}
+
+func WithWindow(window time.Duration) ShedderOption {
+	return func(opts *shedderOptions) {
+		opts.window = window
+	}
+}
+
+type promise struct {
+	start   time.Duration
+	shedder *adaptiveShedder
+}
+
+func (p *promise) Fail() {
+	p.shedder.addFlying(-1)
+}
+
+func (p *promise) Pass() {
+	rt := float64(timex.Since(p.start)) / float64(time.Millisecond)
+	p.shedder.addFlying(-1)
+	p.shedder.rtCounter.Add(math.Ceil(rt))
+	p.shedder.passCounter.Add(1)
+}

core/load/adaptiveshedder_test.go

@@ -0,0 +1,205 @@
+package load
+
+import (
+	"math/rand"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"zero/core/collection"
+	"zero/core/logx"
+	"zero/core/mathx"
+	"zero/core/stat"
+	"zero/core/syncx"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	buckets        = 10
+	bucketDuration = time.Millisecond * 50
+)
+
+func init() {
+	stat.SetReporter(nil)
+}
+
+func TestAdaptiveShedder(t *testing.T) {
+	shedder := NewAdaptiveShedder(WithWindow(bucketDuration), WithBuckets(buckets), WithCpuThreshold(100))
+	var wg sync.WaitGroup
+	var drop int64
+	proba := mathx.NewProba()
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 30; i++ {
+				promise, err := shedder.Allow()
+				if err != nil {
+					atomic.AddInt64(&drop, 1)
+				} else {
+					count := rand.Intn(5)
+					time.Sleep(time.Millisecond * time.Duration(count))
+					if proba.TrueOnProba(0.01) {
+						promise.Fail()
+					} else {
+						promise.Pass()
+					}
+				}
+			}
+		}()
+	}
+	wg.Wait()
+}
+
+func TestAdaptiveShedderMaxPass(t *testing.T) {
+	passCounter := newRollingWindow()
+	for i := 1; i <= 10; i++ {
+		passCounter.Add(float64(i * 100))
+		time.Sleep(bucketDuration)
+	}
+	shedder := &adaptiveShedder{
+		passCounter:     passCounter,
+		droppedRecently: syncx.NewAtomicBool(),
+	}
+	assert.Equal(t, int64(1000), shedder.maxPass())
+
+	// default max pass is equal to 1.
+	passCounter = newRollingWindow()
+	shedder = &adaptiveShedder{
+		passCounter:     passCounter,
+		droppedRecently: syncx.NewAtomicBool(),
+	}
+	assert.Equal(t, int64(1), shedder.maxPass())
+}
+
+func TestAdaptiveShedderMinRt(t *testing.T) {
+	rtCounter := newRollingWindow()
+	for i := 0; i < 10; i++ {
+		if i > 0 {
+			time.Sleep(bucketDuration)
+		}
+		for j := i*10 + 1; j <= i*10+10; j++ {
+			rtCounter.Add(float64(j))
+		}
+	}
+	shedder := &adaptiveShedder{
+		rtCounter: rtCounter,
+	}
+	assert.Equal(t, float64(6), shedder.minRt())
+
+	// default max min rt is equal to maxFloat64.
+	rtCounter = newRollingWindow()
+	shedder = &adaptiveShedder{
+		rtCounter:       rtCounter,
+		droppedRecently: syncx.NewAtomicBool(),
+	}
+	assert.Equal(t, defaultMinRt, shedder.minRt())
+}
+
+func TestAdaptiveShedderMaxFlight(t *testing.T) {
+	passCounter := newRollingWindow()
+	rtCounter := newRollingWindow()
+	for i := 0; i < 10; i++ {
+		if i > 0 {
+			time.Sleep(bucketDuration)
+		}
+		passCounter.Add(float64((i + 1) * 100))
+		for j := i*10 + 1; j <= i*10+10; j++ {
+			rtCounter.Add(float64(j))
+		}
+	}
+	shedder := &adaptiveShedder{
+		passCounter:     passCounter,
+		rtCounter:       rtCounter,
+		windows:         buckets,
+		droppedRecently: syncx.NewAtomicBool(),
+	}
+	assert.Equal(t, int64(54), shedder.maxFlight())
+}
+
+func TestAdaptiveShedderShouldDrop(t *testing.T) {
+	logx.Disable()
+	passCounter := newRollingWindow()
+	rtCounter := newRollingWindow()
+	for i := 0; i < 10; i++ {
+		if i > 0 {
+			time.Sleep(bucketDuration)
+		}
+		passCounter.Add(float64((i + 1) * 100))
+		for j := i*10 + 1; j <= i*10+10; j++ {
+			rtCounter.Add(float64(j))
+		}
+	}
+	shedder := &adaptiveShedder{
+		passCounter:     passCounter,
+		rtCounter:       rtCounter,
+		windows:         buckets,
+		droppedRecently: syncx.NewAtomicBool(),
+	}
+	// cpu >=  800, inflight < maxPass
+	systemOverloadChecker = func(int64) bool {
+		return true
+	}
+	shedder.avgFlying = 50
+	assert.False(t, shedder.shouldDrop())
+
+	// cpu >=  800, inflight > maxPass
+	shedder.avgFlying = 80
+	shedder.flying = 50
+	assert.False(t, shedder.shouldDrop())
+
+	// cpu >=  800, inflight > maxPass
+	shedder.avgFlying = 80
+	shedder.flying = 80
+	assert.True(t, shedder.shouldDrop())
+
+	// cpu < 800, inflight > maxPass
+	systemOverloadChecker = func(int64) bool {
+		return false
+	}
+	shedder.avgFlying = 80
+	assert.False(t, shedder.shouldDrop())
+}
+
+func BenchmarkAdaptiveShedder_Allow(b *testing.B) {
+	logx.Disable()
+
+	bench := func(b *testing.B) {
+		var shedder = NewAdaptiveShedder()
+		proba := mathx.NewProba()
+		for i := 0; i < 6000; i++ {
+			p, err := shedder.Allow()
+			if err == nil {
+				time.Sleep(time.Millisecond)
+				if proba.TrueOnProba(0.01) {
+					p.Fail()
+				} else {
+					p.Pass()
+				}
+			}
+		}
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			p, err := shedder.Allow()
+			if err == nil {
+				p.Pass()
+			}
+		}
+	}
+
+	systemOverloadChecker = func(int64) bool {
+		return true
+	}
+	b.Run("high load", bench)
+	systemOverloadChecker = func(int64) bool {
+		return false
+	}
+	b.Run("low load", bench)
+}
+
+func newRollingWindow() *collection.RollingWindow {
+	return collection.NewRollingWindow(buckets, bucketDuration, collection.IgnoreCurrentBucket())
+}

core/load/nopshedder.go

@@ -0,0 +1,21 @@
+package load
+
+type nopShedder struct {
+}
+
+func newNopShedder() Shedder {
+	return nopShedder{}
+}
+
+func (s nopShedder) Allow() (Promise, error) {
+	return nopPromise{}, nil
+}
+
+type nopPromise struct {
+}
+
+func (p nopPromise) Pass() {
+}
+
+func (p nopPromise) Fail() {
+}

core/load/nopshedder_test.go

@@ -0,0 +1,21 @@
+package load
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNopShedder(t *testing.T) {
+	Disable()
+	shedder := NewAdaptiveShedder()
+	for i := 0; i < 1000; i++ {
+		p, err := shedder.Allow()
+		assert.Nil(t, err)
+		p.Fail()
+	}
+
+	p, err := shedder.Allow()
+	assert.Nil(t, err)
+	p.Pass()
+}

core/load/sheddergroup.go

@@ -0,0 +1,36 @@
+package load
+
+import (
+	"io"
+
+	"zero/core/syncx"
+)
+
+type ShedderGroup struct {
+	options []ShedderOption
+	manager *syncx.ResourceManager
+}
+
+func NewShedderGroup(opts ...ShedderOption) *ShedderGroup {
+	return &ShedderGroup{
+		options: opts,
+		manager: syncx.NewResourceManager(),
+	}
+}
+
+func (g *ShedderGroup) GetShedder(key string) Shedder {
+	shedder, _ := g.manager.GetResource(key, func() (closer io.Closer, e error) {
+		return nopCloser{
+			Shedder: NewAdaptiveShedder(g.options...),
+		}, nil
+	})
+	return shedder.(Shedder)
+}
+
+type nopCloser struct {
+	Shedder
+}
+
+func (c nopCloser) Close() error {
+	return nil
+}

core/load/sheddergroup_test.go

@@ -0,0 +1,15 @@
+package load
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGroup(t *testing.T) {
+	group := NewShedderGroup()
+	t.Run("get", func(t *testing.T) {
+		limiter := group.GetShedder("test")
+		assert.NotNil(t, limiter)
+	})
+}

core/load/sheddingstat.go

@@ -0,0 +1,68 @@
+package load
+
+import (
+	"sync/atomic"
+	"time"
+
+	"zero/core/logx"
+	"zero/core/stat"
+)
+
+type (
+	SheddingStat struct {
+		name  string
+		total int64
+		pass  int64
+		drop  int64
+	}
+
+	snapshot struct {
+		Total int64
+		Pass  int64
+		Drop  int64
+	}
+)
+
+func NewSheddingStat(name string) *SheddingStat {
+	st := &SheddingStat{
+		name: name,
+	}
+	go st.run()
+	return st
+}
+
+func (s *SheddingStat) IncrementTotal() {
+	atomic.AddInt64(&s.total, 1)
+}
+
+func (s *SheddingStat) IncrementPass() {
+	atomic.AddInt64(&s.pass, 1)
+}
+
+func (s *SheddingStat) IncrementDrop() {
+	atomic.AddInt64(&s.drop, 1)
+}
+
+func (s *SheddingStat) reset() snapshot {
+	return snapshot{
+		Total: atomic.SwapInt64(&s.total, 0),
+		Pass:  atomic.SwapInt64(&s.pass, 0),
+		Drop:  atomic.SwapInt64(&s.drop, 0),
+	}
+}
+
+func (s *SheddingStat) run() {
+	ticker := time.NewTicker(time.Minute)
+	defer ticker.Stop()
+	for range ticker.C {
+		c := stat.CpuUsage()
+		st := s.reset()
+		if st.Drop == 0 {
+			logx.Statf("(%s) shedding_stat [1m], cpu: %d, total: %d, pass: %d, drop: %d",
+				s.name, c, st.Total, st.Pass, st.Drop)
+		} else {
+			logx.Statf("(%s) shedding_stat_drop [1m], cpu: %d, total: %d, pass: %d, drop: %d",
+				s.name, c, st.Total, st.Pass, st.Drop)
+		}
+	}
+}

core/load/sheddingstat_test.go

@@ -0,0 +1,24 @@
+package load
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSheddingStat(t *testing.T) {
+	st := NewSheddingStat("any")
+	for i := 0; i < 3; i++ {
+		st.IncrementTotal()
+	}
+	for i := 0; i < 5; i++ {
+		st.IncrementPass()
+	}
+	for i := 0; i < 7; i++ {
+		st.IncrementDrop()
+	}
+	result := st.reset()
+	assert.Equal(t, int64(3), result.Total)
+	assert.Equal(t, int64(5), result.Pass)
+	assert.Equal(t, int64(7), result.Drop)
+}