Repository: zeromicro/go-zero
Pull request #2476 (2022-10-01 12:42:41)

chore: better shedding algorithm, make sure recover from shedding


Patch:
diff --git a/core/load/adaptiveshedder.go b/core/load/adaptiveshedder.go
index c25743b9..f925ec4d 100644
--- a/core/load/adaptiveshedder.go
+++ b/core/load/adaptiveshedder.go
@@ -17,7 +17,7 @@ import (
 const (
 	defaultBuckets = 50
 	defaultWindow  = time.Second * 5
-	// using 1000m notation, 900m is like 80%, keep it as var for unit test
+	// using 1000m notation, 900m is like 90%, 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
@@ -70,7 +70,7 @@ type (
 		flying          int64
 		avgFlying       float64
 		avgFlyingLock   syncx.SpinLock
-		dropTime        *syncx.AtomicDuration
+		overloadTime    *syncx.AtomicDuration
 		droppedRecently *syncx.AtomicBool
 		passCounter     *collection.RollingWindow
 		rtCounter       *collection.RollingWindow
@@ -106,7 +106,7 @@ func NewAdaptiveShedder(opts ...ShedderOption) Shedder {
 	return &adaptiveShedder{
 		cpuThreshold:    options.cpuThreshold,
 		windows:         int64(time.Second / bucketDuration),
-		dropTime:        syncx.NewAtomicDuration(),
+		overloadTime:    syncx.NewAtomicDuration(),
 		droppedRecently: syncx.NewAtomicBool(),
 		passCounter: collection.NewRollingWindow(options.buckets, bucketDuration,
 			collection.IgnoreCurrentBucket()),
@@ -118,7 +118,6 @@ func NewAdaptiveShedder(opts ...ShedderOption) Shedder {
 // Allow implements Shedder.Allow.
 func (as *adaptiveShedder) Allow() (Promise, error) {
 	if as.shouldDrop() {
-		as.dropTime.Set(timex.Now())
 		as.droppedRecently.Set(true)
 
 		return nil, ErrServiceOverloaded
@@ -215,21 +214,26 @@ func (as *adaptiveShedder) stillHot() bool {
 		return false
 	}
 
-	dropTime := as.dropTime.Load()
-	if dropTime == 0 {
+	overloadTime := as.overloadTime.Load()
+	if overloadTime == 0 {
 		return false
 	}
 
-	hot := timex.Since(dropTime) < coolOffDuration
-	if !hot {
-		as.droppedRecently.Set(false)
+	if timex.Since(overloadTime) < coolOffDuration {
+		return true
 	}
 
-	return hot
+	as.droppedRecently.Set(false)
+	return false
 }
 
 func (as *adaptiveShedder) systemOverloaded() bool {
-	return systemOverloadChecker(as.cpuThreshold)
+	if !systemOverloadChecker(as.cpuThreshold) {
+		return false
+	}
+
+	as.overloadTime.Set(timex.Now())
+	return true
 }
 
 // WithBuckets customizes the Shedder with given number of buckets.
diff --git a/core/stores/redis/redis.go b/core/stores/redis/redis.go
index c4d9c401..535f523a 100644
--- a/core/stores/redis/redis.go
+++ b/core/stores/redis/redis.go
@@ -267,6 +267,31 @@ func (s *Redis) BlpopExCtx(ctx context.Context, node RedisNode, key string) (str
 	return vals[1], true, nil
 }
 
+// BlpopWithTimeout uses passed in redis connection to execute blpop command.
+// Control blocking query timeout
+func (s *Redis) BlpopWithTimeout(node RedisNode, timeout time.Duration, key string) (string, error) {
+	return s.BlpopWithTimeoutCtx(context.Background(), node, timeout, key)
+}
+
+// BlpopWithTimeoutCtx uses passed in redis connection to execute blpop command.
+// Control blocking query timeout
+func (s *Redis) BlpopWithTimeoutCtx(ctx context.Context, node RedisNode, timeout time.Duration, key string) (string, error) {
+	if node == nil {
+		return "", ErrNilNode
+	}
+
+	vals, err := node.BLPop(ctx, timeout, key).Result()
+	if err != nil {
+		return "", err
+	}
+
+	if len(vals) < 2 {
+		return "", fmt.Errorf("no value on key: %s", key)
+	}
+
+	return vals[1], nil
+}
+
 // Decr is the implementation of redis decr command.
 func (s *Redis) Decr(key string) (int64, error) {
 	return s.DecrCtx(context.Background(), key)


Test patch:
diff --git a/core/load/adaptiveshedder_test.go b/core/load/adaptiveshedder_test.go
index cb711bbe..94789322 100644
--- a/core/load/adaptiveshedder_test.go
+++ b/core/load/adaptiveshedder_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/zeromicro/go-zero/core/mathx"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/core/syncx"
+	"github.com/zeromicro/go-zero/core/timex"
 )
 
 const (
@@ -136,7 +137,7 @@ func TestAdaptiveShedderShouldDrop(t *testing.T) {
 		passCounter:     passCounter,
 		rtCounter:       rtCounter,
 		windows:         buckets,
-		dropTime:        syncx.NewAtomicDuration(),
+		overloadTime:    syncx.NewAtomicDuration(),
 		droppedRecently: syncx.NewAtomicBool(),
 	}
 	// cpu >=  800, inflight < maxPass
@@ -190,12 +191,15 @@ func TestAdaptiveShedderStillHot(t *testing.T) {
 		passCounter:     passCounter,
 		rtCounter:       rtCounter,
 		windows:         buckets,
-		dropTime:        syncx.NewAtomicDuration(),
+		overloadTime:    syncx.NewAtomicDuration(),
 		droppedRecently: syncx.ForAtomicBool(true),
 	}
 	assert.False(t, shedder.stillHot())
-	shedder.dropTime.Set(-coolOffDuration * 2)
+	shedder.overloadTime.Set(-coolOffDuration * 2)
 	assert.False(t, shedder.stillHot())
+	shedder.droppedRecently.Set(true)
+	shedder.overloadTime.Set(timex.Now())
+	assert.True(t, shedder.stillHot())
 }
 
 func BenchmarkAdaptiveShedder_Allow(b *testing.B) {
diff --git a/core/stores/redis/redis_test.go b/core/stores/redis/redis_test.go
index 4f960909..3dd2a576 100644
--- a/core/stores/redis/redis_test.go
+++ b/core/stores/redis/redis_test.go
@@ -1117,6 +1117,17 @@ func TestRedisBlpopEx(t *testing.T) {
 	})
 }
 
+func TestRedisBlpopWithTimeout(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		client.Ping()
+		var node mockedNode
+		_, err := client.BlpopWithTimeout(nil, 10*time.Second, "foo")
+		assert.NotNil(t, err)
+		_, err = client.BlpopWithTimeout(node, 10*time.Second, "foo")
+		assert.NotNil(t, err)
+	})
+}
+
 func TestRedisGeo(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		client.Ping()