support regex in flow and isolation for resource name (#595)

Author: StarJourneyMingsing

core/flow/rule.go

@@ -117,6 +117,8 @@ type Rule struct {
 	HighMemUsageThreshold int64 `json:"highMemUsageThreshold"`
 	MemLowWaterMarkBytes  int64 `json:"memLowWaterMarkBytes"`
 	MemHighWaterMarkBytes int64 `json:"memHighWaterMarkBytes"`
+	// Regex indicates whether the rule is a regex rule
+	Regex bool `json:"regex"`
 }
 
 func (r *Rule) isEqualsTo(newRule *Rule) bool {

core/flow/rule_manager.go

@@ -17,6 +17,7 @@ package flow
 import (
 	"fmt"
 	"reflect"
+	"regexp"
 	"sync"
 
 	"github.com/alibaba/sentinel-golang/core/base"
@@ -41,10 +42,12 @@ type trafficControllerGenKey struct {
 type TrafficControllerMap map[string][]*TrafficShapingController
 
 var (
-	tcGenFuncMap = make(map[trafficControllerGenKey]TrafficControllerGenFunc, 6)
-	tcMap        = make(TrafficControllerMap)
-	tcMux        = new(sync.RWMutex)
-	nopStat      = &standaloneStatistic{
+	tcGenFuncMap    = make(map[trafficControllerGenKey]TrafficControllerGenFunc, 6)
+	tcMap           = make(TrafficControllerMap)
+	tcRegexMap      = make(TrafficControllerMap)
+	tcRegexCacheMap = make(TrafficControllerMap)
+	tcMux           = new(sync.RWMutex)
+	nopStat         = &standaloneStatistic{
 		reuseResourceStat: false,
 		readOnlyMetric:    base.NopReadStat(),
 		writeOnlyMetric:   base.NopWriteStat(),
@@ -204,24 +207,35 @@ func onRuleUpdate(rawResRulesMap map[string][]*Rule) (err error) {
 	start := util.CurrentTimeNano()
 
 	tcMux.RLock()
-	tcMapClone := make(TrafficControllerMap, len(validResRulesMap))
+	tcMapClone := make(TrafficControllerMap, len(tcMap)+len(tcRegexMap))
 	for res, tcs := range tcMap {
 		resTcClone := make([]*TrafficShapingController, 0, len(tcs))
 		resTcClone = append(resTcClone, tcs...)
 		tcMapClone[res] = resTcClone
 	}
+	for res, tcs := range tcRegexMap {
+		resTcClone := make([]*TrafficShapingController, 0, len(tcs))
+		resTcClone = append(resTcClone, tcs...)
+		tcMapClone[res] = resTcClone
+	}
 	tcMux.RUnlock()
 
-	m := make(TrafficControllerMap, len(validResRulesMap))
+	m := make(TrafficControllerMap)
+	regexM := make(TrafficControllerMap)
 	for res, rulesOfRes := range validResRulesMap {
-		newTcsOfRes := buildResourceTrafficShapingController(res, rulesOfRes, tcMapClone[res])
+		newTcsOfRes, newRegexTcsOfRes := buildResourceTrafficShapingController(res, rulesOfRes, tcMapClone[res])
 		if len(newTcsOfRes) > 0 {
 			m[res] = newTcsOfRes
 		}
+		if len(newRegexTcsOfRes) > 0 {
+			regexM[res] = newRegexTcsOfRes
+		}
 	}
 
 	tcMux.Lock()
 	tcMap = m
+	tcRegexMap = regexM
+	tcRegexCacheMap = make(TrafficControllerMap)
 	tcMux.Unlock()
 	currentRules = rawResRulesMap
 
@@ -277,15 +291,22 @@ func onResourceRuleUpdate(res string, rawResRules []*Rule) (err error) {
 	oldResTcs := make([]*TrafficShapingController, 0)
 	tcMux.RLock()
 	oldResTcs = append(oldResTcs, tcMap[res]...)
+	oldResTcs = append(oldResTcs, tcRegexMap[res]...)
 	tcMux.RUnlock()
-	newResTcs := buildResourceTrafficShapingController(res, validResRules, oldResTcs)
+	newResTcs, newRegexResTcs := buildResourceTrafficShapingController(res, validResRules, oldResTcs)
 
 	tcMux.Lock()
 	if len(newResTcs) == 0 {
 		delete(tcMap, res)
 	} else {
 		tcMap[res] = newResTcs
 	}
+	if len(newRegexResTcs) == 0 {
+		delete(tcRegexMap, res)
+	} else {
+		tcRegexMap[res] = newRegexResTcs
+	}
+	tcRegexCacheMap = make(TrafficControllerMap)
 	tcMux.Unlock()
 	currentRules[res] = rawResRules
 	logging.Debug("[Flow onResourceRuleUpdate] Time statistic(ns) for updating flow rule", "timeCost", util.CurrentTimeNano()-start)
@@ -308,6 +329,8 @@ func LoadRulesOfResource(res string, rules []*Rule) (bool, error) {
 		// clear tcMap
 		tcMux.Lock()
 		delete(tcMap, res)
+		delete(tcRegexMap, res)
+		tcRegexCacheMap = make(TrafficControllerMap)
 		tcMux.Unlock()
 		logging.Info("[Flow] clear resource level rules", "resource", res)
 		return true, nil
@@ -329,19 +352,31 @@ func getRules() []*Rule {
 	tcMux.RLock()
 	defer tcMux.RUnlock()
 
-	return rulesFrom(tcMap)
+	return append(rulesFrom(tcMap), rulesFrom(tcRegexMap)...)
 }
 
 // getRulesOfResource returns specific resource's rules。Any changes of rules take effect for flow module
 // getRulesOfResource is an internal interface.
 func getRulesOfResource(res string) []*Rule {
 	tcMux.RLock()
-	defer tcMux.RUnlock()
 
-	resTcs, exist := tcMap[res]
-	if !exist {
-		return nil
+	resTcs := make([]*TrafficShapingController, 0)
+	if tcs, ok := tcMap[res]; ok {
+		resTcs = append(resTcs, tcs...)
+	}
+	cacheRegexTcs, ok := tcRegexCacheMap[res]
+	tcMux.RUnlock()
+
+	if ok {
+		resTcs = append(resTcs, cacheRegexTcs...)
+	} else {
+		tcMux.Lock()
+		regexTcs := regexTcMatches(res)
+		resTcs = append(resTcs, regexTcs...)
+		tcRegexCacheMap[res] = regexTcs
+		tcMux.Unlock()
 	}
+
 	ret := make([]*Rule, 0, len(resTcs))
 	for _, tc := range resTcs {
 		ret = append(ret, tc.BoundRule())
@@ -507,9 +542,28 @@ func RemoveTrafficShapingGenerator(tokenCalculateStrategy TokenCalculateStrategy
 
 func getTrafficControllerListFor(name string) []*TrafficShapingController {
 	tcMux.RLock()
-	defer tcMux.RUnlock()
+	rules := make([]*TrafficShapingController, 0)
+	rules = append(rules, tcMap[name]...)
+	regexCacheRules, ok := tcRegexCacheMap[name]
+	tcMux.RUnlock()
+
+	if ok {
+		rules = append(rules, regexCacheRules...)
+		return rules
+	}
 
-	return tcMap[name]
+	tcMux.Lock()
+	defer tcMux.Unlock()
+	// double check
+	if regexCacheRules, ok = tcRegexCacheMap[name]; ok {
+		rules = append(rules, regexCacheRules...)
+		return rules
+	}
+	regexRules := regexTcMatches(name)
+	tcRegexCacheMap[name] = regexRules
+	rules = append(rules, regexRules...)
+
+	return rules
 }
 
 func calculateReuseIndexFor(r *Rule, oldResTcs []*TrafficShapingController) (equalIdx, reuseStatIdx int) {
@@ -539,8 +593,10 @@ 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, oldResTcs []*TrafficShapingController) []*TrafficShapingController {
+func buildResourceTrafficShapingController(res string, rulesOfRes []*Rule,
+	oldResTcs []*TrafficShapingController) (simpleControllers []*TrafficShapingController, regexControllers []*TrafficShapingController) {
 	newTcsOfRes := make([]*TrafficShapingController, 0, len(rulesOfRes))
+	newRegexTcsOfRes := make([]*TrafficShapingController, 0, len(rulesOfRes))
 	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)
@@ -582,9 +638,15 @@ func buildResourceTrafficShapingController(res string, rulesOfRes []*Rule, oldRe
 			// remove old tc from oldResTcs
 			oldResTcs = append(oldResTcs[:reuseStatIdx], oldResTcs[reuseStatIdx+1:]...)
 		}
-		newTcsOfRes = append(newTcsOfRes, tc)
+
+		if tc.rule != nil && tc.rule.Regex {
+			newRegexTcsOfRes = append(newRegexTcsOfRes, tc)
+		} else {
+			newTcsOfRes = append(newTcsOfRes, tc)
+		}
+
 	}
-	return newTcsOfRes
+	return newTcsOfRes, newRegexTcsOfRes
 }
 
 // IsValidRule checks whether the given Rule is valid.
@@ -649,3 +711,35 @@ func IsValidRule(rule *Rule) error {
 
 	return nil
 }
+
+func regexTcMatches(resource string) []*TrafficShapingController {
+	result := make([]*TrafficShapingController, 0)
+	for pattern, controllers := range tcRegexMap {
+		re := regexp.MustCompile(pattern)
+		if !re.MatchString(resource) {
+			continue
+		}
+		controllersCopy := make([]*TrafficShapingController, len(controllers))
+		for resourceName, controller := range controllers {
+			generator, supported := tcGenFuncMap[trafficControllerGenKey{
+				tokenCalculateStrategy: controller.rule.TokenCalculateStrategy,
+				controlBehavior:        controller.rule.ControlBehavior,
+			}]
+			if !supported || generator == nil {
+				logging.Error(errors.New("get trafficShapingController copy failed, unsupported flow control strategy"),
+					"Ignoring the rule due to unsupported control behavior in flow."+
+						"buildResourceTrafficShapingController()", "rule", controller.rule)
+				continue
+			}
+			if tc, e := generator(controller.rule, nil); e == nil && tc != nil {
+				controllersCopy[resourceName] = tc
+			} else {
+				logging.Error(errors.New("get trafficShapingController copy failed, bad generated traffic controller"),
+					"Ignoring the rule due to bad generated traffic controller in flow."+
+						"buildResourceTrafficShapingController()", "rule", controller.rule)
+			}
+		}
+		result = append(result, controllersCopy...)
+	}
+	return result
+}

core/flow/rule_manager_test.go

@@ -263,7 +263,7 @@ func Test_buildResourceTrafficShapingController(t *testing.T) {
 			MaxQueueingTimeMs:      10,
 		}
 		assert.True(t, len(tcMap["abc1"]) == 0)
-		tcs := buildResourceTrafficShapingController("abc1", []*Rule{r1, r2}, tcMap["abc1"])
+		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)
@@ -415,7 +415,7 @@ func Test_buildResourceTrafficShapingController(t *testing.T) {
 			StatIntervalInMs:       50000,
 		}
 
-		tcs := buildResourceTrafficShapingController("abc1", []*Rule{r12, r22, r32, r42}, tcMap["abc1"])
+		tcs, _ := buildResourceTrafficShapingController("abc1", []*Rule{r12, r22, r32, r42}, tcMap["abc1"])
 		assert.True(t, len(tcs) == 4)
 
 		assert.True(t, tcs[0].BoundRule() == r12)

core/flow/slot_test.go

@@ -66,3 +66,68 @@ func Test_FlowSlot_StandaloneStat(t *testing.T) {
 	}
 	assert.True(t, getTrafficControllerListFor("abc")[0].boundStat.readOnlyMetric.GetSum(base.MetricEventPass) == 50)
 }
+
+func Test_RegexFlowSlot_StandaloneStat(t *testing.T) {
+	slot := &Slot{}
+	statSLot := &StandaloneStatSlot{}
+	res1 := base.NewResourceWrapper("abc/123", base.ResTypeCommon, base.Inbound)
+	res2 := base.NewResourceWrapper("abc/456", base.ResTypeCommon, base.Inbound)
+	resNode1 := stat.GetOrCreateResourceNode("abc/123", base.ResTypeCommon)
+	resNode2 := stat.GetOrCreateResourceNode("abc/456", base.ResTypeCommon)
+	ctx1 := &base.EntryContext{
+		Resource: res1,
+		StatNode: resNode1,
+		Input: &base.SentinelInput{
+			BatchCount: 1,
+		},
+		RuleCheckResult: nil,
+		Data:            nil,
+	}
+	ctx2 := &base.EntryContext{
+		Resource: res2,
+		StatNode: resNode2,
+		Input: &base.SentinelInput{
+			BatchCount: 1,
+		},
+		RuleCheckResult: nil,
+		Data:            nil,
+	}
+
+	r1 := &Rule{
+		Resource:               "abc/*",
+		TokenCalculateStrategy: Direct,
+		ControlBehavior:        Reject,
+		// Use standalone statistic, using single-bucket-sliding-windows
+		StatIntervalInMs: 20000,
+		Threshold:        100,
+		RelationStrategy: CurrentResource,
+		Regex:            true,
+	}
+	_, e := LoadRules([]*Rule{r1})
+	if e != nil {
+		logging.Error(e, "")
+		t.Fail()
+		return
+	}
+
+	for i := 0; i < 80; i++ {
+		ret := slot.Check(ctx1)
+		if ret != nil {
+			t.Fail()
+			return
+		}
+		statSLot.OnEntryPassed(ctx1)
+	}
+	for i := 0; i < 80; i++ {
+		ret := slot.Check(ctx2)
+		if ret != nil {
+			t.Fail()
+			return
+		}
+		statSLot.OnEntryPassed(ctx2)
+	}
+	assert.True(t, getTrafficControllerListFor("abc/123")[0].boundStat.readOnlyMetric.GetSum(base.
+		MetricEventPass) == 80)
+	assert.True(t, getTrafficControllerListFor("abc/456")[0].boundStat.readOnlyMetric.GetSum(base.
+		MetricEventPass) == 80)
+}

core/isolation/rule.go

@@ -46,13 +46,16 @@ type Rule struct {
 	// Currently Concurrency is supported for concurrency limiting.
 	MetricType MetricType `json:"metricType"`
 	Threshold  uint32     `json:"threshold"`
+	// Regex indicates whether the rule is a regex rule
+	Regex bool `json:"regex"`
 }
 
 func (r *Rule) String() string {
 	b, err := json.Marshal(r)
 	if err != nil {
 		// Return the fallback string
-		return fmt.Sprintf("{Id=%s, Resource=%s, MetricType=%s, Threshold=%d}", r.ID, r.Resource, r.MetricType.String(), r.Threshold)
+		return fmt.Sprintf("{Id=%s, Resource=%s, MetricType=%s, Threshold=%d, Regex=%v}", r.ID, r.Resource,
+			r.MetricType.String(), r.Threshold, r.Regex)
 	}
 	return string(b)
 }