Fix bug: incorrectly defined storeToken using uint in warm-up flow control (#218)

louyuting · web-flow · commit 1a987a8edc2e · 2020-09-01T10:08:44.000+08:00
* Add flow rule checking logic for WarmUpColdFactor
diff --git a/core/flow/rule_manager.go b/core/flow/rule_manager.go
@@ -223,6 +223,9 @@ func checkControlBehaviorField(rule *FlowRule) error {
 		if rule.WarmUpPeriodSec <= 0 {
 			return errors.New("invalid warmUpPeriodSec")
 		}
+		if rule.WarmUpColdFactor == 1 {
+			return errors.New("WarmUpColdFactor must be great than 1")
+		}
 		return nil
 	case WarmUpThrottling:
 		if rule.WarmUpPeriodSec <= 0 {
diff --git a/core/flow/tc_warm_up.go b/core/flow/tc_warm_up.go
@@ -18,14 +18,14 @@ type WarmUpTrafficShapingCalculator struct {
 	warningToken      uint64
 	maxToken          uint64
 	slope             float64
-	storedTokens      *uint64
-	lastFilledTime    *uint64
+	storedTokens      int64
+	lastFilledTime    uint64
 }
 
 func NewWarmUpTrafficShapingCalculator(rule *FlowRule) *WarmUpTrafficShapingCalculator {
-	if rule.WarmUpColdFactor == 0 {
+	if rule.WarmUpColdFactor <= 1 {
 		rule.WarmUpColdFactor = config.DefaultWarmUpColdFactor
-		logging.Warnf("[NewWarmUpTrafficShapingCalculator] No set WarmUpColdFactor,use default values: %d", config.DefaultWarmUpColdFactor)
+		logging.Warnf("[NewWarmUpTrafficShapingCalculator] invalid WarmUpColdFactor,use default values: %d", config.DefaultWarmUpColdFactor)
 	}
 
 	warningToken := uint64((float64(rule.WarmUpPeriodSec) * rule.Count) / float64(rule.WarmUpColdFactor-1))
@@ -41,64 +41,67 @@ func NewWarmUpTrafficShapingCalculator(rule *FlowRule) *WarmUpTrafficShapingCalc
 		maxToken:          maxToken,
 		slope:             slope,
 		threshold:         rule.Count,
-		storedTokens:      new(uint64),
-		lastFilledTime:    new(uint64),
+		storedTokens:      0,
+		lastFilledTime:    0,
 	}
 
 	return warmUpTrafficShapingCalculator
 }
 
-func (d *WarmUpTrafficShapingCalculator) CalculateAllowedTokens(node base.StatNode, acquireCount uint32, flag int32) float64 {
+func (c *WarmUpTrafficShapingCalculator) CalculateAllowedTokens(node base.StatNode, acquireCount uint32, flag int32) float64 {
 	previousQps := node.GetPreviousQPS(base.MetricEventPass)
-	d.syncToken(previousQps)
+	c.syncToken(previousQps)
 
-	restToken := atomic.LoadUint64(d.storedTokens)
-	if restToken >= d.warningToken {
-		aboveToken := restToken - d.warningToken
-		warningQps := math.Nextafter(1.0/(float64(aboveToken)*d.slope+1.0/d.threshold), math.MaxFloat64)
+	restToken := atomic.LoadInt64(&c.storedTokens)
+	if restToken < 0 {
+		restToken = 0
+	}
+	if restToken >= int64(c.warningToken) {
+		aboveToken := restToken - int64(c.warningToken)
+		warningQps := math.Nextafter(1.0/(float64(aboveToken)*c.slope+1.0/c.threshold), math.MaxFloat64)
 		return warningQps
 	} else {
-		return d.threshold
+		return c.threshold
 	}
 }
 
-func (d *WarmUpTrafficShapingCalculator) syncToken(passQps float64) {
+func (c *WarmUpTrafficShapingCalculator) syncToken(passQps float64) {
 	currentTime := util.CurrentTimeMillis()
 	currentTime = currentTime - currentTime%1000
 
-	oldLastFillTime := atomic.LoadUint64(d.lastFilledTime)
+	oldLastFillTime := atomic.LoadUint64(&c.lastFilledTime)
 	if currentTime <= oldLastFillTime {
 		return
 	}
 
-	oldValue := atomic.LoadUint64(d.storedTokens)
-	newValue := d.coolDownTokens(currentTime, passQps)
+	oldValue := atomic.LoadInt64(&c.storedTokens)
+	newValue := c.coolDownTokens(currentTime, passQps)
 
-	if atomic.CompareAndSwapUint64(d.storedTokens, oldValue, newValue) {
-		if currentValue := atomic.AddUint64(d.storedTokens, uint64(0-passQps)); currentValue < 0 {
-			atomic.StoreUint64(d.storedTokens, 0)
+	if atomic.CompareAndSwapInt64(&c.storedTokens, oldValue, newValue) {
+		if currentValue := atomic.AddInt64(&c.storedTokens, int64(-passQps)); currentValue < 0 {
+			atomic.StoreInt64(&c.storedTokens, 0)
 		}
-		atomic.StoreUint64(d.lastFilledTime, currentTime)
+		atomic.StoreUint64(&c.lastFilledTime, currentTime)
 	}
 }
 
-func (d *WarmUpTrafficShapingCalculator) coolDownTokens(currentTime uint64, passQps float64) uint64 {
-	oldValue := atomic.LoadUint64(d.storedTokens)
+func (c *WarmUpTrafficShapingCalculator) coolDownTokens(currentTime uint64, passQps float64) int64 {
+	oldValue := atomic.LoadInt64(&c.storedTokens)
 	newValue := oldValue
 
 	// Prerequisites for adding a token:
 	// When token consumption is much lower than the warning line
-	if oldValue < d.warningToken {
-		newValue = uint64(float64(oldValue) + (float64(currentTime)-float64(atomic.LoadUint64(d.lastFilledTime)))*d.threshold/1000)
-	} else if oldValue > d.warningToken {
-		if passQps < float64(uint32(d.threshold)/d.coldFactor) {
-			newValue = uint64(float64(oldValue) + float64(currentTime-atomic.LoadUint64(d.lastFilledTime))*d.threshold/1000)
+	if oldValue < int64(c.warningToken) {
+		newValue = int64(float64(oldValue) + (float64(currentTime)-float64(atomic.LoadUint64(&c.lastFilledTime)))*c.threshold/1000)
+	} else if oldValue > int64(c.warningToken) {
+		if passQps < float64(uint32(c.threshold)/c.coldFactor) {
+			newValue = int64(float64(oldValue) + float64(currentTime-atomic.LoadUint64(&c.lastFilledTime))*c.threshold/1000)
 		}
 	}
 
-	if newValue <= d.maxToken {
+	if newValue <= int64(c.maxToken) {
 		return newValue
 	} else {
-		return d.maxToken
+		return int64(c.maxToken)
 	}
 }
diff --git a/example/warm_up/qps_warm_up_example.go b/example/warm_up/qps_warm_up_example.go
@@ -10,6 +10,7 @@ import (
 	sentinel "github.com/alibaba/sentinel-golang/api"
 	"github.com/alibaba/sentinel-golang/core/base"
 	"github.com/alibaba/sentinel-golang/core/flow"
+	"github.com/alibaba/sentinel-golang/logging"
 	"github.com/alibaba/sentinel-golang/util"
 )
 
@@ -25,17 +26,19 @@ func main() {
 	counter := Counter{pass: new(int64), block: new(int64), total: new(int64)}
 	// We should initialize Sentinel first.
 	err := sentinel.InitDefault()
+	logging.ResetGlobalLogger(logging.NewConsoleLogger("flow-warmup-test"))
 	if err != nil {
 		log.Fatalf("Unexpected error: %+v", err)
 	}
 
 	_, err = flow.LoadRules([]*flow.FlowRule{
 		{
-			Resource:        "some-test",
-			MetricType:      flow.QPS,
-			Count:           100,
-			ControlBehavior: flow.WarmUp,
-			WarmUpPeriodSec: 10,
+			Resource:         "some-test",
+			MetricType:       flow.QPS,
+			Count:            100,
+			ControlBehavior:  flow.WarmUp,
+			WarmUpPeriodSec:  10,
+			WarmUpColdFactor: 3,
 		},
 	})
 	if err != nil {
