Repository: alibaba/sentinel-golang
Pull request #372 (2021-01-05 16:13:14)

Rule Manager: support resource level rule replacement in isolation module
<!--  Thanks for submitting a pull request! Here are some tips for you:
1. Please make sure you have read and understood the contributing guidelines: https://github.com/alibaba/blob/master/CONTRIBUTING.md
2. Please make sure the PR has a corresponding issue.
-->

### Describe what this PR does / why we need it
Rule Manager: support resource level rule replacement in isolation 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/372?src=pr&el=h1) Report
> Merging [#372](https://codecov.io/gh/alibaba/sentinel-golang/pull/372?src=pr&el=desc) (9779694) into [master](https://codecov.io/gh/alibaba/sentinel-golang/commit/1eaa5c5cc8ea53d25036bf0eb36ba42719cc5b1b?el=desc) (1eaa5c5) will **increase** coverage by `0.30%`.
> The diff coverage is `85.96%`.

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

```diff
@@            Coverage Diff             @@
##           master     #372      +/-   ##
==========================================
+ Coverage   50.42%   50.72%   +0.30%     
==========================================
  Files          70       70              
  Lines        4141     4183      +42     
==========================================
+ Hits         2088     2122      +34     
- Misses       1771     1776       +5     
- Partials      282      285       +3     
```


| [Impacted Files](https://codecov.io/gh/alibaba/sentinel-golang/pull/372?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [core/isolation/rule\_manager.go](https://codecov.io/gh/alibaba/sentinel-golang/pull/372/diff?src=pr&el=tree#diff-Y29yZS9pc29sYXRpb24vcnVsZV9tYW5hZ2VyLmdv) | `69.72% <85.71%> (+7.03%)` | :arrow_up: |
| [core/flow/rule\_manager.go](https://codecov.io/gh/alibaba/sentinel-golang/pull/372/diff?src=pr&el=tree#diff-Y29yZS9mbG93L3J1bGVfbWFuYWdlci5nbw==) | `61.09% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/alibaba/sentinel-golang/pull/372?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/372?src=pr&el=footer). Last update [1eaa5c5...9779694](https://codecov.io/gh/alibaba/sentinel-golang/pull/372?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/core/flow/rule_manager.go b/core/flow/rule_manager.go
index 041a2c2..4e5814a 100644
--- a/core/flow/rule_manager.go
+++ b/core/flow/rule_manager.go
@@ -303,7 +303,7 @@ func onResourceRuleUpdate(res string, rawResRules []*Rule) (err error) {
 	}
 	tcMux.Unlock()
 	currentRules[res] = rawResRules
-	logging.Debug("[Flow onRuleUpdate] Time statistic(ns) for updating flow rule", "timeCost", util.CurrentTimeNano()-start)
+	logging.Debug("[Flow onResourceRuleUpdate] Time statistic(ns) for updating flow rule", "timeCost", util.CurrentTimeNano()-start)
 	logging.Info("[Flow] load resource level rules", "resource", res, "validResRules", validResRules)
 	return nil
 }
diff --git a/core/isolation/rule_manager.go b/core/isolation/rule_manager.go
index f1654f9..399a2fe 100644
--- a/core/isolation/rule_manager.go
+++ b/core/isolation/rule_manager.go
@@ -27,63 +27,139 @@ import (
 var (
 	ruleMap       = make(map[string][]*Rule)
 	rwMux         = &sync.RWMutex{}
-	currentRules  = make([]*Rule, 0)
+	currentRules  = make(map[string][]*Rule, 0)
 	updateRuleMux = new(sync.Mutex)
 )
 
 // LoadRules loads the given isolation 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) {
+	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("[Isolation] Load rules is the same with current rules, so ignore load operation.")
 		return false, nil
 	}
 
-	err := onRuleUpdate(rules)
+	err := onRuleUpdate(resRulesMap)
 	return true, err
 }
 
-func onRuleUpdate(rules []*Rule) (err error) {
-	m := make(map[string][]*Rule, len(rules))
-	for _, r := range rules {
-		if e := IsValid(r); e != nil {
-			logging.Error(e, "Invalid isolation rule in isolation.LoadRules()", "rule", r)
-			continue
+func onRuleUpdate(rawResRulesMap map[string][]*Rule) (err error) {
+	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("[Isolation onRuleUpdate] Ignoring invalid isolation rule", "rule", rule, "reason", err.Error())
+				continue
+			}
+			validResRules = append(validResRules, rule)
 		}
-		resRules, ok := m[r.Resource]
-		if !ok {
-			resRules = make([]*Rule, 0, 1)
+		if len(validResRules) > 0 {
+			validResRulesMap[res] = validResRules
 		}
-		m[r.Resource] = append(resRules, r)
 	}
 
 	start := util.CurrentTimeNano()
 
-	for res, rs := range m {
+	for res, rs := range validResRulesMap {
 		if len(rs) > 0 {
 			// update resource slot chain
 			misc.RegisterRuleCheckSlotForResource(res, DefaultSlot)
 		}
 	}
 	rwMux.Lock()
-	ruleMap = m
-	currentRules = rules
+	ruleMap = validResRulesMap
 	rwMux.Unlock()
+	currentRules = rawResRulesMap
 
-	logging.Debug("[Isolation LoadRules] Time statistic(ns) for updating isolation rule", "timeCost", util.CurrentTimeNano()-start)
-	logRuleUpdate(m)
+	logging.Debug("[Isolation onRuleUpdate] Time statistic(ns) for updating isolation rule", "timeCost", util.CurrentTimeNano()-start)
+	logRuleUpdate(validResRulesMap)
 	return
 }
 
+// LoadRulesOfResource loads the given resource's isolation 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 ruleMap
+		rwMux.Lock()
+		delete(ruleMap, res)
+		rwMux.Unlock()
+		logging.Info("[Isolation] clear resource level rules", "resource", res)
+		return true, nil
+	}
+	// load resource level rules
+	isEqual := reflect.DeepEqual(currentRules[res], rules)
+	if isEqual {
+		logging.Info("[Isolation] 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
+}
+
+func onResourceRuleUpdate(res string, rawResRules []*Rule) (err error) {
+	validResRules := make([]*Rule, 0, len(rawResRules))
+	for _, rule := range rawResRules {
+		if err := IsValidRule(rule); err != nil {
+			logging.Warn("[Isolation onResourceRuleUpdate] Ignoring invalid flow rule", "rule", rule, "reason", err.Error())
+			continue
+		}
+		validResRules = append(validResRules, rule)
+	}
+
+	start := util.CurrentTimeNano()
+	if len(validResRules) > 0 {
+		// update resource slot chain
+		misc.RegisterRuleCheckSlotForResource(res, DefaultSlot)
+	}
+
+	rwMux.Lock()
+	if len(validResRules) == 0 {
+		delete(ruleMap, res)
+	} else {
+		ruleMap[res] = validResRules
+	}
+	rwMux.Unlock()
+	currentRules[res] = rawResRules
+	logging.Debug("[Isolation onResourceRuleUpdate] Time statistic(ns) for updating flow rule", "timeCost", util.CurrentTimeNano()-start)
+	logging.Info("[Isolation] load resource level rules", "resource", res, "validResRules", validResRules)
+	return nil
+}
+
 // ClearRules clears all the rules in isolation module.
 func ClearRules() error {
 	_, err := LoadRules(nil)
 	return err
 }
 
+// ClearRulesOfResource clears resource level rules in isolation module.
+func ClearRulesOfResource(res string) error {
+	_, err := LoadRulesOfResource(res, nil)
+	return err
+}
+
 // GetRules returns all the rules based on copy.
 // It doesn't take effect for isolation module if user changes the rule.
 func GetRules() []Rule {
@@ -157,7 +233,7 @@ func logRuleUpdate(m map[string][]*Rule) {
 }
 
 // IsValidRule checks whether the given Rule is valid.
-func IsValid(r *Rule) error {
+func IsValidRule(r *Rule) error {
 	if r == nil {
 		return errors.New("nil isolation rule")
 	}


Test patch:
diff --git a/core/isolation/rule_manager_test.go b/core/isolation/rule_manager_test.go
index d177bc1..471d41d 100644
--- a/core/isolation/rule_manager_test.go
+++ b/core/isolation/rule_manager_test.go
@@ -21,6 +21,11 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func clearData() {
+	ruleMap = make(map[string][]*Rule)
+	currentRules = make(map[string][]*Rule, 0)
+}
+
 func TestLoadRules(t *testing.T) {
 	t.Run("TestLoadRules_1", func(t *testing.T) {
 		logging.ResetGlobalLoggerLevel(logging.DebugLevel)
@@ -47,8 +52,7 @@ func TestLoadRules(t *testing.T) {
 		assert.True(t, len(ruleMap["abc2"]) == 1)
 		assert.True(t, ruleMap["abc2"][0] == r2)
 
-		err = ClearRules()
-		assert.True(t, err == nil)
+		clearData()
 	})
 
 	t.Run("loadSameRules", func(t *testing.T) {
@@ -69,5 +73,153 @@ func TestLoadRules(t *testing.T) {
 		})
 		assert.Nil(t, err)
 		assert.False(t, ok)
+		clearData()
+	})
+
+	t.Run("TestClearRules_normal", func(t *testing.T) {
+		r1 := &Rule{
+			Resource:   "abc1",
+			MetricType: Concurrency,
+			Threshold:  100,
+		}
+		r2 := &Rule{
+			Resource:   "abc2",
+			MetricType: Concurrency,
+			Threshold:  200,
+		}
+		r3 := &Rule{
+			Resource:   "abc3",
+			MetricType: MetricType(1),
+			Threshold:  200,
+		}
+
+		succ, err := LoadRules([]*Rule{r1, r2, r3})
+		assert.True(t, succ && err == nil)
+
+		assert.True(t, ClearRules() == nil)
+
+		assert.True(t, len(ruleMap["abc1"]) == 0)
+		assert.True(t, len(currentRules["abc1"]) == 0)
+		assert.True(t, len(ruleMap["abc2"]) == 0)
+		assert.True(t, len(currentRules["abc2"]) == 0)
+		assert.True(t, len(ruleMap["abc3"]) == 0)
+		assert.True(t, len(currentRules["abc3"]) == 0)
+		clearData()
+	})
+}
+
+func TestLoadRulesOfResource(t *testing.T) {
+	r1 := &Rule{
+		Resource:   "abc1",
+		MetricType: Concurrency,
+		Threshold:  100,
+	}
+	r2 := &Rule{
+		Resource:   "abc2",
+		MetricType: Concurrency,
+		Threshold:  200,
+	}
+	r3 := &Rule{
+		Resource:   "abc3",
+		MetricType: MetricType(1),
+		Threshold:  200,
+	}
+
+	r4 := &Rule{
+		Resource:   "abc2",
+		MetricType: MetricType(1),
+		Threshold:  300,
+	}
+
+	succ, err := LoadRules([]*Rule{r1, r2, r3, r4})
+	assert.True(t, succ && err == nil)
+
+	t.Run("LoadRulesOfResource_empty_resource", func(t *testing.T) {
+		succ, err = LoadRulesOfResource("", []*Rule{r1, r2})
+		assert.True(t, !succ && err != nil)
+	})
+
+	t.Run("LoadRulesOfResource_cache_hit", func(t *testing.T) {
+		r111 := *r2
+		r122 := *r4
+		succ, err = LoadRulesOfResource("abc2", []*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(ruleMap["abc1"]) == 0 && len(currentRules["abc1"]) == 0)
+		assert.True(t, len(ruleMap["abc2"]) == 1 && len(currentRules["abc2"]) == 2)
+	})
+	clearData()
+}
+
+func Test_ResourceRuleUpdate(t *testing.T) {
+	logging.ResetGlobalLoggerLevel(logging.DebugLevel)
+	t.Run("Test_onResourceRuleUpdate_normal", func(t *testing.T) {
+		r1 := &Rule{
+			Resource:   "abc1",
+			MetricType: Concurrency,
+			Threshold:  100,
+		}
+		r2 := &Rule{
+			Resource:   "abc2",
+			MetricType: Concurrency,
+			Threshold:  200,
+		}
+		r3 := &Rule{
+			Resource:   "abc3",
+			MetricType: MetricType(1),
+			Threshold:  200,
+		}
+
+		succ, err := LoadRules([]*Rule{r1, r2, r3})
+		assert.True(t, succ && err == nil)
+
+		r111 := r1
+		r111.Threshold = 100
+		err = onResourceRuleUpdate("abc1", []*Rule{r111})
+
+		assert.True(t, err == nil)
+		assert.True(t, len(ruleMap["abc1"]) == 1)
+		assert.True(t, len(currentRules["abc1"]) == 1)
+		assert.True(t, ruleMap["abc1"][0] == r111)
+
+		assert.True(t, len(ruleMap["abc2"]) == 1)
+		assert.True(t, len(currentRules["abc2"]) == 1)
+
+		clearData()
+	})
+
+	t.Run("TestClearRulesOfResource_normal", func(t *testing.T) {
+		r1 := &Rule{
+			Resource:   "abc1",
+			MetricType: Concurrency,
+			Threshold:  100,
+		}
+		r2 := &Rule{
+			Resource:   "abc2",
+			MetricType: Concurrency,
+			Threshold:  200,
+		}
+		r3 := &Rule{
+			Resource:   "abc3",
+			MetricType: MetricType(1),
+			Threshold:  200,
+		}
+
+		succ, err := LoadRules([]*Rule{r1, r2, r3})
+		assert.True(t, succ && err == nil)
+
+		assert.True(t, ClearRulesOfResource("abc1") == nil)
+
+		assert.True(t, len(ruleMap["abc1"]) == 0)
+		assert.True(t, len(currentRules["abc1"]) == 0)
+		assert.True(t, len(ruleMap["abc2"]) == 1)
+		assert.True(t, len(currentRules["abc2"]) == 1)
+		assert.True(t, len(ruleMap["abc3"]) == 0)
+		assert.True(t, len(currentRules["abc3"]) == 1)
+		clearData()
 	})
 }