Repository: alibaba/sentinel-golang
Pull request #370 (2021-01-04 12:00:17)

Rule Manager: support resource level rule replacement in flow module

### Describe what this PR does / why we need it
Rule Manager: support resource level rule replacement in flow module

### Does this pull request fix one issue?

<!--If that, add "Fixes #xxxx" below in the next line. For example, Fixes #15. Otherwise, add "NONE" -->

### Describe how you did it


### Describe how to verify it


### Special notes for reviews

codecov-io: # [Codecov](https://codecov.io/gh/alibaba/sentinel-golang/pull/370?src=pr&el=h1) Report
> Merging [#370](https://codecov.io/gh/alibaba/sentinel-golang/pull/370?src=pr&el=desc) (bde8dcb) into [master](https://codecov.io/gh/alibaba/sentinel-golang/commit/c0595301e04248716c66d8da248d38fd01def6b0?el=desc) (c059530) will **decrease** coverage by `0.61%`.
> The diff coverage is `7.27%`.

[![Impacted file tree graph](https://codecov.io/gh/alibaba/sentinel-golang/pull/370/graphs/tree.svg?width=650&height=150&src=pr&token=SrpzbhKA3E)](https://codecov.io/gh/alibaba/sentinel-golang/pull/370?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #370      +/-   ##
==========================================
- Coverage   50.20%   49.58%   -0.62%     
==========================================
  Files          70       70              
  Lines        4089     4140      +51     
==========================================
  Hits         2053     2053              
- Misses       1759     1810      +51     
  Partials      277      277              
```


| [Impacted Files](https://codecov.io/gh/alibaba/sentinel-golang/pull/370?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [core/flow/rule\_manager.go](https://codecov.io/gh/alibaba/sentinel-golang/pull/370/diff?src=pr&el=tree#diff-Y29yZS9mbG93L3J1bGVfbWFuYWdlci5nbw==) | `51.64% <5.55%> (-8.42%)` | :arrow_down: |
| [core/flow/rule.go](https://codecov.io/gh/alibaba/sentinel-golang/pull/370/diff?src=pr&el=tree#diff-Y29yZS9mbG93L3J1bGUuZ28=) | `27.45% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/alibaba/sentinel-golang/pull/370?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/alibaba/sentinel-golang/pull/370?src=pr&el=footer). Last update [c059530...bde8dcb](https://codecov.io/gh/alibaba/sentinel-golang/pull/370?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

louyuting: @liqiangz  Could you please review it?

Patch:
diff --git a/core/flow/rule.go b/core/flow/rule.go
index 983210f..31c6f14 100644
--- a/core/flow/rule.go
+++ b/core/flow/rule.go
@@ -145,7 +145,7 @@ func (r *Rule) isStatReusable(newRule *Rule) bool {
 }
 
 func (r *Rule) needStatistic() bool {
-	return !((r.TokenCalculateStrategy == Direct && r.ControlBehavior == Throttling) || (r.TokenCalculateStrategy == MemoryAdaptive && r.ControlBehavior == Throttling))
+	return r.TokenCalculateStrategy == WarmUp || r.ControlBehavior == Reject
 }
 
 func (r *Rule) String() string {
diff --git a/core/flow/rule_manager.go b/core/flow/rule_manager.go
index 2bec703..041a2c2 100644
--- a/core/flow/rule_manager.go
+++ b/core/flow/rule_manager.go
@@ -50,7 +50,7 @@ var (
 		readOnlyMetric:    base.NopReadStat(),
 		writeOnlyMetric:   base.NopWriteStat(),
 	}
-	currentRules  = make([]*Rule, 0)
+	currentRules  = make(map[string][]*Rule, 0)
 	updateRuleMux = new(sync.Mutex)
 )
 
@@ -175,7 +175,7 @@ func logRuleUpdate(m map[string][]*Rule) {
 	}
 }
 
-func onRuleUpdate(rules []*Rule) (err error) {
+func onRuleUpdate(rawResRulesMap map[string][]*Rule) (err error) {
 	defer func() {
 		if r := recover(); r != nil {
 			var ok bool
@@ -186,23 +186,26 @@ func onRuleUpdate(rules []*Rule) (err error) {
 		}
 	}()
 
-	resRulesMap := make(map[string][]*Rule, len(rules))
-	for _, rule := range rules {
-		if err := IsValidRule(rule); err != nil {
-			logging.Warn("[Flow onRuleUpdate] Ignoring invalid flow rule", "rule", rule, "reason", err.Error())
-			continue
+	// ignore invalid rules
+	validResRulesMap := make(map[string][]*Rule, len(rawResRulesMap))
+	for res, rules := range rawResRulesMap {
+		validResRules := make([]*Rule, 0, len(rules))
+		for _, rule := range rules {
+			if err := IsValidRule(rule); err != nil {
+				logging.Warn("[Flow onRuleUpdate] Ignoring invalid flow rule", "rule", rule, "reason", err.Error())
+				continue
+			}
+			validResRules = append(validResRules, rule)
 		}
-		resRules, exist := resRulesMap[rule.Resource]
-		if !exist {
-			resRules = make([]*Rule, 0, 1)
+		if len(validResRules) > 0 {
+			validResRulesMap[res] = validResRules
 		}
-		resRulesMap[rule.Resource] = append(resRules, rule)
 	}
-	m := make(TrafficControllerMap, len(resRulesMap))
+
 	start := util.CurrentTimeNano()
 
 	tcMux.RLock()
-	tcMapClone := make(TrafficControllerMap, len(resRulesMap))
+	tcMapClone := make(TrafficControllerMap, len(validResRulesMap))
 	for res, tcs := range tcMap {
 		resTcClone := make([]*TrafficShapingController, 0, len(tcs))
 		resTcClone = append(resTcClone, tcs...)
@@ -210,8 +213,12 @@ func onRuleUpdate(rules []*Rule) (err error) {
 	}
 	tcMux.RUnlock()
 
-	for res, rulesOfRes := range resRulesMap {
-		m[res] = buildResourceTrafficShapingController(res, rulesOfRes, tcMapClone)
+	m := make(TrafficControllerMap, len(validResRulesMap))
+	for res, rulesOfRes := range validResRulesMap {
+		newTcsOfRes := buildResourceTrafficShapingController(res, rulesOfRes, tcMapClone[res])
+		if len(newTcsOfRes) > 0 {
+			m[res] = newTcsOfRes
+		}
 	}
 
 	for res, tcs := range m {
@@ -224,27 +231,110 @@ func onRuleUpdate(rules []*Rule) (err error) {
 
 	tcMux.Lock()
 	tcMap = m
-	currentRules = rules
 	tcMux.Unlock()
+	currentRules = rawResRulesMap
 
 	logging.Debug("[Flow onRuleUpdate] Time statistic(ns) for updating flow rule", "timeCost", util.CurrentTimeNano()-start)
-	logRuleUpdate(resRulesMap)
+	logRuleUpdate(validResRulesMap)
 	return nil
 }
 
 // LoadRules loads the given flow rules to the rule manager, while all previous rules will be replaced.
 // the first returned value indicates whether do real load operation, if the rules is the same with previous rules, return false
 func LoadRules(rules []*Rule) (bool, error) {
-	// TODO: rethink the design
+	resRulesMap := make(map[string][]*Rule, 16)
+	for _, rule := range rules {
+		resRules, exist := resRulesMap[rule.Resource]
+		if !exist {
+			resRules = make([]*Rule, 0, 1)
+		}
+		resRulesMap[rule.Resource] = append(resRules, rule)
+	}
+
 	updateRuleMux.Lock()
 	defer updateRuleMux.Unlock()
-	isEqual := reflect.DeepEqual(currentRules, rules)
+	isEqual := reflect.DeepEqual(currentRules, resRulesMap)
 	if isEqual {
 		logging.Info("[Flow] Load rules is the same with current rules, so ignore load operation.")
 		return false, nil
 	}
+	err := onRuleUpdate(resRulesMap)
+	return true, err
+}
+
+func onResourceRuleUpdate(res string, rawResRules []*Rule) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			var ok bool
+			err, ok = r.(error)
+			if !ok {
+				err = fmt.Errorf("%v", r)
+			}
+		}
+	}()
 
-	err := onRuleUpdate(rules)
+	validResRules := make([]*Rule, 0, len(rawResRules))
+	for _, rule := range rawResRules {
+		if err := IsValidRule(rule); err != nil {
+			logging.Warn("[Flow onResourceRuleUpdate] Ignoring invalid flow rule", "rule", rule, "reason", err.Error())
+			continue
+		}
+		validResRules = append(validResRules, rule)
+	}
+
+	start := util.CurrentTimeNano()
+	oldResTcs := make([]*TrafficShapingController, 0)
+	tcMux.RLock()
+	oldResTcs = append(oldResTcs, tcMap[res]...)
+	tcMux.RUnlock()
+	newResTcs := buildResourceTrafficShapingController(res, validResRules, oldResTcs)
+
+	if len(newResTcs) > 0 {
+		// update resource slot chain
+		misc.RegisterRuleCheckSlotForResource(res, DefaultSlot)
+		misc.RegisterStatSlotForResource(res, DefaultStandaloneStatSlot)
+	}
+
+	tcMux.Lock()
+	if len(newResTcs) == 0 {
+		delete(tcMap, res)
+	} else {
+		tcMap[res] = newResTcs
+	}
+	tcMux.Unlock()
+	currentRules[res] = rawResRules
+	logging.Debug("[Flow onRuleUpdate] Time statistic(ns) for updating flow rule", "timeCost", util.CurrentTimeNano()-start)
+	logging.Info("[Flow] load resource level rules", "resource", res, "validResRules", validResRules)
+	return nil
+}
+
+// LoadRulesOfResource loads the given resource's flow rules to the rule manager, while all previous resource's rules will be replaced.
+// the first returned value indicates whether do real load operation, if the rules is the same with previous resource's rules, return false
+func LoadRulesOfResource(res string, rules []*Rule) (bool, error) {
+	if len(res) == 0 {
+		return false, errors.New("empty resource")
+	}
+	updateRuleMux.Lock()
+	defer updateRuleMux.Unlock()
+	// clear resource rules
+	if len(rules) == 0 {
+		// clear resource's currentRules
+		delete(currentRules, res)
+		// clear tcMap
+		tcMux.Lock()
+		delete(tcMap, res)
+		tcMux.Unlock()
+		logging.Info("[Flow] clear resource level rules", "resource", res)
+		return true, nil
+	}
+	// load resource level rules
+	isEqual := reflect.DeepEqual(currentRules[res], rules)
+	if isEqual {
+		logging.Info("[Flow] Load resource level rules is the same with current resource level rules, so ignore load operation.")
+		return false, nil
+	}
+
+	err := onResourceRuleUpdate(res, rules)
 	return true, err
 }
 
@@ -302,6 +392,12 @@ func ClearRules() error {
 	return err
 }
 
+// ClearRulesOfResource clears resource level rules in flow module.
+func ClearRulesOfResource(res string) error {
+	_, err := LoadRulesOfResource(res, nil)
+	return err
+}
+
 func rulesFrom(m TrafficControllerMap) []*Rule {
 	rules := make([]*Rule, 0, 8)
 	if len(m) == 0 {
@@ -458,19 +554,13 @@ func calculateReuseIndexFor(r *Rule, oldResTcs []*TrafficShapingController) (equ
 }
 
 // buildResourceTrafficShapingController builds TrafficShapingController slice from rules. the resource of rules must be equals to res
-func buildResourceTrafficShapingController(res string, rulesOfRes []*Rule, m TrafficControllerMap) []*TrafficShapingController {
+func buildResourceTrafficShapingController(res string, rulesOfRes []*Rule, oldResTcs []*TrafficShapingController) []*TrafficShapingController {
 	newTcsOfRes := make([]*TrafficShapingController, 0, len(rulesOfRes))
-	emptyTcs := make([]*TrafficShapingController, 0, 0)
 	for _, rule := range rulesOfRes {
 		if res != rule.Resource {
 			logging.Error(errors.Errorf("unmatched resource name expect: %s, actual: %s", res, rule.Resource), "Unmatched resource name in flow.buildResourceTrafficShapingController()", "rule", rule)
 			continue
 		}
-		oldResTcs, exist := m[res]
-		if !exist {
-			oldResTcs = emptyTcs
-		}
-
 		equalIdx, reuseStatIdx := calculateReuseIndexFor(rule, oldResTcs)
 
 		// First check equals scenario
@@ -479,7 +569,7 @@ func buildResourceTrafficShapingController(res string, rulesOfRes []*Rule, m Tra
 			equalOldTc := oldResTcs[equalIdx]
 			newTcsOfRes = append(newTcsOfRes, equalOldTc)
 			// remove old tc from oldResTcs
-			m[res] = append(oldResTcs[:equalIdx], oldResTcs[equalIdx+1:]...)
+			oldResTcs = append(oldResTcs[:equalIdx], oldResTcs[equalIdx+1:]...)
 			continue
 		}
 
@@ -505,7 +595,7 @@ func buildResourceTrafficShapingController(res string, rulesOfRes []*Rule, m Tra
 		}
 		if reuseStatIdx >= 0 {
 			// remove old tc from oldResTcs
-			m[res] = append(oldResTcs[:reuseStatIdx], oldResTcs[reuseStatIdx+1:]...)
+			oldResTcs = append(oldResTcs[:reuseStatIdx], oldResTcs[reuseStatIdx+1:]...)
 		}
 		newTcsOfRes = append(newTcsOfRes, tc)
 	}


Test patch:
diff --git a/core/flow/rule_manager_test.go b/core/flow/rule_manager_test.go
index 6eb59fd..3025aae 100644
--- a/core/flow/rule_manager_test.go
+++ b/core/flow/rule_manager_test.go
@@ -23,6 +23,10 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func clearData() {
+	tcMap = make(TrafficControllerMap)
+	currentRules = make(map[string][]*Rule, 0)
+}
 func TestSetAndRemoveTrafficShapingGenerator(t *testing.T) {
 	tsc := &TrafficShapingController{}
 
@@ -61,7 +65,7 @@ func TestSetAndRemoveTrafficShapingGenerator(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotContains(t, tcGenFuncMap, cs)
 
-	_, _ = LoadRules([]*Rule{})
+	clearData()
 }
 
 func TestIsValidFlowRule(t *testing.T) {
@@ -124,9 +128,7 @@ func TestGetRules(t *testing.T) {
 			assert.True(t, reflect.DeepEqual(&rs1[0], r2))
 			assert.True(t, reflect.DeepEqual(&rs1[1], r1))
 		}
-		if err := ClearRules(); err != nil {
-			t.Fatal(err)
-		}
+		clearData()
 	})
 
 	t.Run("getRules", func(t *testing.T) {
@@ -169,9 +171,7 @@ func TestGetRules(t *testing.T) {
 		assert.True(t, len(tcMap["abc2"]) == 1 && !tcMap["abc2"][0].boundStat.reuseResourceStat)
 		assert.True(t, reflect.DeepEqual(tcMap["abc2"][0].boundStat.readOnlyMetric, nopStat.readOnlyMetric))
 		assert.True(t, reflect.DeepEqual(tcMap["abc2"][0].boundStat.writeOnlyMetric, nopStat.writeOnlyMetric))
-		if err := ClearRules(); err != nil {
-			t.Fatal(err)
-		}
+		clearData()
 	})
 }
 
@@ -263,12 +263,14 @@ func Test_buildResourceTrafficShapingController(t *testing.T) {
 			MaxQueueingTimeMs:      10,
 		}
 		assert.True(t, len(tcMap["abc1"]) == 0)
-		tcs := buildResourceTrafficShapingController("abc1", []*Rule{r1, r2}, tcMap)
+		tcs := buildResourceTrafficShapingController("abc1", []*Rule{r1, r2}, tcMap["abc1"])
 		assert.True(t, len(tcs) == 2)
 		assert.True(t, tcs[0].BoundRule() == r1)
 		assert.True(t, tcs[1].BoundRule() == r2)
 		assert.True(t, reflect.DeepEqual(tcs[0].BoundRule(), r1))
 		assert.True(t, reflect.DeepEqual(tcs[1].BoundRule(), r2))
+
+		clearData()
 	})
 
 	t.Run("Test_buildResourceTrafficShapingController_reuse_stat", func(t *testing.T) {
@@ -413,7 +415,7 @@ func Test_buildResourceTrafficShapingController(t *testing.T) {
 			StatIntervalInMs:       50000,
 		}
 
-		tcs := buildResourceTrafficShapingController("abc1", []*Rule{r12, r22, r32, r42}, tcMap)
+		tcs := buildResourceTrafficShapingController("abc1", []*Rule{r12, r22, r32, r42}, tcMap["abc1"])
 		assert.True(t, len(tcs) == 4)
 
 		assert.True(t, tcs[0].BoundRule() == r12)
@@ -430,6 +432,8 @@ func Test_buildResourceTrafficShapingController(t *testing.T) {
 		assert.True(t, tcs[1].boundStat != stat2)
 		assert.True(t, tcs[2] == fakeTc3)
 		assert.True(t, tcs[3].boundStat == stat4)
+
+		clearData()
 	})
 }
 
@@ -454,6 +458,8 @@ func TestLoadRules(t *testing.T) {
 		})
 		assert.Nil(t, err)
 		assert.False(t, ok)
+
+		clearData()
 	})
 }
 
@@ -493,3 +499,138 @@ func TestIsValidRule(t *testing.T) {
 	rule1.MemHighWaterMarkBytes = 300 * 1024 * 1024
 	assert.Nil(t, IsValidRule(rule1))
 }
+
+func TestLoadRulesOfResource(t *testing.T) {
+	r11 := &Rule{
+		Resource:               "abc1",
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+		Threshold:              10,
+	}
+	r12 := &Rule{
+		Resource:               "abc1",
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+		Threshold:              20,
+	}
+	r21 := &Rule{
+		Resource:               "abc2",
+		Threshold:              10,
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+	}
+	r22 := &Rule{
+		Resource:               "abc2",
+		Threshold:              20,
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+	}
+
+	succ, err := LoadRules([]*Rule{r11, r12, r21, r22})
+	assert.True(t, succ && err == nil)
+
+	t.Run("LoadRulesOfResource_empty_resource", func(t *testing.T) {
+		succ, err = LoadRulesOfResource("", []*Rule{r11, r12})
+		assert.True(t, !succ && err != nil)
+	})
+
+	t.Run("LoadRulesOfResource_cache_hit", func(t *testing.T) {
+		r111 := *r11
+		r122 := *r12
+		succ, err = LoadRulesOfResource("abc1", []*Rule{&r111, &r122})
+		assert.True(t, !succ && err == nil)
+	})
+
+	t.Run("LoadRulesOfResource_clear", func(t *testing.T) {
+		succ, err = LoadRulesOfResource("abc1", []*Rule{})
+		assert.True(t, succ && err == nil)
+		assert.True(t, len(tcMap["abc1"]) == 0 && len(currentRules["abc1"]) == 0)
+		assert.True(t, len(tcMap["abc2"]) == 2 && len(currentRules["abc2"]) == 2)
+	})
+	clearData()
+}
+
+func Test_onResourceRuleUpdate(t *testing.T) {
+	r11 := Rule{
+		Resource:               "abc1",
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+		Threshold:              10,
+	}
+	r12 := Rule{
+		Resource:               "abc1",
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+		Threshold:              20,
+	}
+	r21 := Rule{
+		Resource:               "abc2",
+		Threshold:              10,
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+	}
+	r22 := Rule{
+		Resource:               "abc2",
+		Threshold:              20,
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+	}
+
+	succ, err := LoadRules([]*Rule{&r11, &r12, &r21, &r22})
+	assert.True(t, succ && err == nil)
+
+	t.Run("Test_onResourceRuleUpdate_normal", func(t *testing.T) {
+		r111 := r11
+		r111.Threshold = 100
+		err = onResourceRuleUpdate("abc1", []*Rule{&r111})
+
+		assert.True(t, len(tcMap["abc1"]) == 1)
+		assert.True(t, len(currentRules["abc1"]) == 1)
+		assert.True(t, tcMap["abc1"][0].rule == &r111)
+
+		assert.True(t, len(tcMap["abc2"]) == 2)
+		assert.True(t, len(currentRules["abc2"]) == 2)
+
+		clearData()
+	})
+}
+
+func TestClearRulesOfResource(t *testing.T) {
+	r11 := Rule{
+		Resource:               "abc1",
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+		Threshold:              10,
+	}
+	r12 := Rule{
+		Resource:               "abc1",
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+		Threshold:              20,
+	}
+	r21 := Rule{
+		Resource:               "abc2",
+		Threshold:              10,
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+	}
+	r22 := Rule{
+		Resource:               "abc2",
+		Threshold:              20,
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+	}
+
+	succ, err := LoadRules([]*Rule{&r11, &r12, &r21, &r22})
+	assert.True(t, succ && err == nil)
+
+	t.Run("TestClearRulesOfResource_normal", func(t *testing.T) {
+		assert.True(t, ClearRulesOfResource("abc1") == nil)
+
+		assert.True(t, len(tcMap["abc1"]) == 0)
+		assert.True(t, len(currentRules["abc1"]) == 0)
+		assert.True(t, len(tcMap["abc2"]) == 2)
+		assert.True(t, len(currentRules["abc2"]) == 2)
+		clearData()
+	})
+}