Repository: alibaba/sentinel-golang
Pull request #223 (2020-09-03 13:11:26)

Break down flow.ControlBehavior to flow.TokenCalculateStrategy and flow.ControlBehavior
### Describe what this PR does / why we need it
Break down flow.ControlBehavior to flow.AllocationStrategy and flow.ControlEffect
```go
const (
	Default AllocationStrategy = iota
	WarmUp
)

const (
	Reject ControlEffect = iota
	Throttling
)
```
### 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


Patch:
diff --git a/core/flow/rule.go b/core/flow/rule.go
index 632c659..5932945 100644
--- a/core/flow/rule.go
+++ b/core/flow/rule.go
@@ -15,55 +15,103 @@ const (
 	QPS
 )
 
+func (s MetricType) String() string {
+	switch s {
+	case Concurrency:
+		return "Concurrency"
+	case QPS:
+		return "QPS"
+	default:
+		return "Undefined"
+	}
+}
+
 // RelationStrategy indicates the flow control strategy based on the relation of invocations.
 type RelationStrategy int32
 
 const (
-	// Direct means flow control by current resource directly.
-	Direct RelationStrategy = iota
+	// CurrentResource means flow control by current resource directly.
+	CurrentResource RelationStrategy = iota
 	// AssociatedResource means flow control by the associated resource rather than current resource.
 	AssociatedResource
 )
 
-// ControlBehavior indicates the traffic shaping behaviour.
+func (s RelationStrategy) String() string {
+	switch s {
+	case CurrentResource:
+		return "CurrentResource"
+	case AssociatedResource:
+		return "AssociatedResource"
+	default:
+		return "Undefined"
+	}
+}
+
+type TokenCalculateStrategy int32
+
+const (
+	Direct TokenCalculateStrategy = iota
+	WarmUp
+)
+
+func (s TokenCalculateStrategy) String() string {
+	switch s {
+	case Direct:
+		return "Direct"
+	case WarmUp:
+		return "WarmUp"
+	default:
+		return "Undefined"
+	}
+}
+
 type ControlBehavior int32
 
 const (
 	Reject ControlBehavior = iota
-	WarmUp
 	Throttling
-	WarmUpThrottling
 )
 
+func (s ControlBehavior) String() string {
+	switch s {
+	case Reject:
+		return "Reject"
+	case Throttling:
+		return "Throttling"
+	default:
+		return "Undefined"
+	}
+}
+
 // Rule describes the strategy of flow control.
 type Rule struct {
 	// ID represents the unique ID of the rule (optional).
 	ID uint64 `json:"id,omitempty"`
 
 	// Resource represents the resource name.
-	Resource   string     `json:"resource"`
-	MetricType MetricType `json:"metricType"`
-	// Count represents the threshold.
-	Count           float64         `json:"count"`
-	ControlBehavior ControlBehavior `json:"controlBehavior"`
-
-	RelationStrategy  RelationStrategy `json:"relationStrategy"`
-	RefResource       string           `json:"refResource"`
-	MaxQueueingTimeMs uint32           `json:"maxQueueingTimeMs"`
-	WarmUpPeriodSec   uint32           `json:"warmUpPeriodSec"`
-	WarmUpColdFactor  uint32           `json:"warmUpColdFactor"`
+	Resource               string                 `json:"resource"`
+	MetricType             MetricType             `json:"metricType"`
+	TokenCalculateStrategy TokenCalculateStrategy `json:"tokenCalculateStrategy"`
+	ControlBehavior        ControlBehavior        `json:"controlBehavior"`
+	Count                  float64                `json:"count"`
+	RelationStrategy       RelationStrategy       `json:"relationStrategy"`
+	RefResource            string                 `json:"refResource"`
+	MaxQueueingTimeMs      uint32                 `json:"maxQueueingTimeMs"`
+	WarmUpPeriodSec        uint32                 `json:"warmUpPeriodSec"`
+	WarmUpColdFactor       uint32                 `json:"warmUpColdFactor"`
 }
 
-func (f *Rule) String() string {
-	b, err := json.Marshal(f)
+func (r *Rule) String() string {
+	b, err := json.Marshal(r)
 	if err != nil {
 		// Return the fallback string
-		return fmt.Sprintf("Rule{resource=%s, id=%d, metricType=%d, threshold=%.2f}",
-			f.Resource, f.ID, f.MetricType, f.Count)
+		return fmt.Sprintf("Rule{Resource=%s, MetricType=%s, TokenCalculateStrategy=%s, ControlBehavior=%s, "+
+			"Count=%.2f, RelationStrategy=%s, WarmUpPeriodSec=%d, WarmUpColdFactor=%d, MaxQueueingTimeMs=%d}",
+			r.Resource, r.MetricType, r.TokenCalculateStrategy, r.ControlBehavior, r.Count, r.RelationStrategy, r.WarmUpPeriodSec, r.WarmUpColdFactor, r.MaxQueueingTimeMs)
 	}
 	return string(b)
 }
 
-func (f *Rule) ResourceName() string {
-	return f.Resource
+func (r *Rule) ResourceName() string {
+	return r.Resource
 }
diff --git a/core/flow/rule_manager.go b/core/flow/rule_manager.go
index 7549192..dde6957 100644
--- a/core/flow/rule_manager.go
+++ b/core/flow/rule_manager.go
@@ -13,34 +13,62 @@ import (
 // TrafficControllerGenFunc represents the TrafficShapingController generator function of a specific control behavior.
 type TrafficControllerGenFunc func(*Rule) *TrafficShapingController
 
+type trafficControllerGenKey struct {
+	tokenCalculateStrategy TokenCalculateStrategy
+	controlBehavior        ControlBehavior
+}
+
 // TrafficControllerMap represents the map storage for TrafficShapingController.
 type TrafficControllerMap map[string][]*TrafficShapingController
 
 var (
-	tcGenFuncMap = make(map[ControlBehavior]TrafficControllerGenFunc)
+	tcGenFuncMap = make(map[trafficControllerGenKey]TrafficControllerGenFunc)
 	tcMap        = make(TrafficControllerMap)
 	tcMux        = new(sync.RWMutex)
 )
 
 func init() {
 	// Initialize the traffic shaping controller generator map for existing control behaviors.
-	tcGenFuncMap[Reject] = func(rule *Rule) *TrafficShapingController {
-		return NewTrafficShapingController(NewDefaultTrafficShapingCalculator(rule.Count), NewDefaultTrafficShapingChecker(rule), rule)
+	tcGenFuncMap[trafficControllerGenKey{
+		tokenCalculateStrategy: Direct,
+		controlBehavior:        Reject,
+	}] = func(rule *Rule) *TrafficShapingController {
+		return NewTrafficShapingController(NewDirectTrafficShapingCalculator(rule.Count), NewDefaultTrafficShapingChecker(rule), rule)
 	}
-	tcGenFuncMap[Throttling] = func(rule *Rule) *TrafficShapingController {
-		return NewTrafficShapingController(NewDefaultTrafficShapingCalculator(rule.Count), NewThrottlingChecker(rule.MaxQueueingTimeMs), rule)
+	tcGenFuncMap[trafficControllerGenKey{
+		tokenCalculateStrategy: Direct,
+		controlBehavior:        Throttling,
+	}] = func(rule *Rule) *TrafficShapingController {
+		return NewTrafficShapingController(NewDirectTrafficShapingCalculator(rule.Count), NewThrottlingChecker(rule.MaxQueueingTimeMs), rule)
 	}
-	tcGenFuncMap[WarmUp] = func(rule *Rule) *TrafficShapingController {
+	tcGenFuncMap[trafficControllerGenKey{
+		tokenCalculateStrategy: WarmUp,
+		controlBehavior:        Reject,
+	}] = func(rule *Rule) *TrafficShapingController {
 		return NewTrafficShapingController(NewWarmUpTrafficShapingCalculator(rule), NewDefaultTrafficShapingChecker(rule), rule)
 	}
+	tcGenFuncMap[trafficControllerGenKey{
+		tokenCalculateStrategy: WarmUp,
+		controlBehavior:        Throttling,
+	}] = func(rule *Rule) *TrafficShapingController {
+		return NewTrafficShapingController(NewWarmUpTrafficShapingCalculator(rule), NewThrottlingChecker(rule.MaxQueueingTimeMs), rule)
+	}
 }
 
 func logRuleUpdate(m TrafficControllerMap) {
 	bs, err := json.Marshal(rulesFrom(m))
 	if err != nil {
-		logging.Info("[FlowRuleManager] Flow rules loaded")
+		if len(m) == 0 {
+			logging.Info("[FlowRuleManager] Flow rules were cleared")
+		} else {
+			logging.Info("[FlowRuleManager] Flow rules were loaded")
+		}
 	} else {
-		logging.Infof("[FlowRuleManager] Flow rules loaded: %s", bs)
+		if len(m) == 0 {
+			logging.Info("[FlowRuleManager] Flow rules were cleared")
+		} else {
+			logging.Infof("[FlowRuleManager] Flow rules were loaded: %s", bs)
+		}
 	}
 }
 
@@ -79,13 +107,55 @@ func LoadRules(rules []*Rule) (bool, error) {
 	return true, err
 }
 
-func GetRules() []*Rule {
+// getRules returns all the rules。Any changes of rules take effect for flow module
+// getRules is an internal interface.
+func getRules() []*Rule {
 	tcMux.RLock()
 	defer tcMux.RUnlock()
 
 	return rulesFrom(tcMap)
 }
 
+// getResRules returns specific resource's rules。Any changes of rules take effect for flow module
+// getResRules is an internal interface.
+func getResRules(res string) []*Rule {
+	tcMux.RLock()
+	defer tcMux.RUnlock()
+
+	resTcs, exist := tcMap[res]
+	if !exist {
+		return nil
+	}
+	ret := make([]*Rule, 0, len(resTcs))
+	for _, tc := range resTcs {
+		ret = append(ret, tc.Rule())
+	}
+	return ret
+}
+
+// GetRules returns all the rules based on copy.
+// It doesn't take effect for flow module if user changes the rule.
+func GetRules() []Rule {
+	rules := getRules()
+	ret := make([]Rule, 0, len(rules))
+	for _, rule := range rules {
+		ret = append(ret, *rule)
+	}
+	return ret
+}
+
+// GetResRules returns specific resource's rules based on copy.
+// It doesn't take effect for flow module if user changes the rule.
+func GetResRules(res string) []Rule {
+	rules := getResRules(res)
+	ret := make([]Rule, 0, len(rules))
+	for _, rule := range rules {
+		ret = append(ret, *rule)
+	}
+	return ret
+}
+
+// ClearRules clears all the rules in flow module.
 func ClearRules() error {
 	_, err := LoadRules(nil)
 	return err
@@ -109,30 +179,43 @@ func rulesFrom(m TrafficControllerMap) []*Rule {
 	return rules
 }
 
-// SetTrafficShapingGenerator sets the traffic controller generator for the given control behavior.
-// Note that modifying the generator of default control behaviors is not allowed.
-func SetTrafficShapingGenerator(cb ControlBehavior, generator TrafficControllerGenFunc) error {
+// SetTrafficShapingGenerator sets the traffic controller generator for the given TokenCalculateStrategy and ControlBehavior.
+// Note that modifying the generator of default control strategy is not allowed.
+func SetTrafficShapingGenerator(tokenCalculateStrategy TokenCalculateStrategy, controlBehavior ControlBehavior, generator TrafficControllerGenFunc) error {
 	if generator == nil {
 		return errors.New("nil generator")
 	}
-	if cb >= Reject && cb <= WarmUpThrottling {
-		return errors.New("not allowed to replace the generator for default control behaviors")
+
+	if tokenCalculateStrategy >= Direct && tokenCalculateStrategy <= WarmUp {
+		return errors.New("not allowed to replace the generator for default control strategy")
+	}
+	if controlBehavior >= Reject && controlBehavior <= Throttling {
+		return errors.New("not allowed to replace the generator for default control strategy")
 	}
 	tcMux.Lock()
 	defer tcMux.Unlock()
 
-	tcGenFuncMap[cb] = generator
+	tcGenFuncMap[trafficControllerGenKey{
+		tokenCalculateStrategy: tokenCalculateStrategy,
+		controlBehavior:        controlBehavior,
+	}] = generator
 	return nil
 }
 
-func RemoveTrafficShapingGenerator(cb ControlBehavior) error {
-	if cb >= Reject && cb <= WarmUpThrottling {
-		return errors.New("not allowed to replace the generator for default control behaviors")
+func RemoveTrafficShapingGenerator(tokenCalculateStrategy TokenCalculateStrategy, controlBehavior ControlBehavior) error {
+	if tokenCalculateStrategy >= Direct && tokenCalculateStrategy <= WarmUp {
+		return errors.New("not allowed to replace the generator for default control strategy")
+	}
+	if controlBehavior >= Reject && controlBehavior <= Throttling {
+		return errors.New("not allowed to replace the generator for default control strategy")
 	}
 	tcMux.Lock()
 	defer tcMux.Unlock()
 
-	delete(tcGenFuncMap, cb)
+	delete(tcGenFuncMap, trafficControllerGenKey{
+		tokenCalculateStrategy: tokenCalculateStrategy,
+		controlBehavior:        controlBehavior,
+	})
 	return nil
 }
 
@@ -151,11 +234,14 @@ func buildFlowMap(rules []*Rule) TrafficControllerMap {
 	}
 
 	for _, rule := range rules {
-		if err := IsValidFlowRule(rule); err != nil {
+		if err := IsValidRule(rule); err != nil {
 			logging.Warnf("Ignoring invalid flow rule: %v, reason: %s", rule, err.Error())
 			continue
 		}
-		generator, supported := tcGenFuncMap[rule.ControlBehavior]
+		generator, supported := tcGenFuncMap[trafficControllerGenKey{
+			tokenCalculateStrategy: rule.TokenCalculateStrategy,
+			controlBehavior:        rule.ControlBehavior,
+		}]
 		if !supported {
 			logging.Warnf("Ignoring the rule due to unsupported control behavior: %v", rule)
 			continue
@@ -176,8 +262,8 @@ func buildFlowMap(rules []*Rule) TrafficControllerMap {
 	return m
 }
 
-// IsValidFlowRule checks whether the given Rule is valid.
-func IsValidFlowRule(rule *Rule) error {
+// IsValidRule checks whether the given Rule is valid.
+func IsValidRule(rule *Rule) error {
 	if rule == nil {
 		return errors.New("nil Rule")
 	}
@@ -193,33 +279,21 @@ func IsValidFlowRule(rule *Rule) error {
 	if rule.RelationStrategy < 0 {
 		return errors.New("invalid relation strategy")
 	}
-	if rule.ControlBehavior < 0 {
-		return errors.New("invalid control behavior")
+	if rule.TokenCalculateStrategy < 0 || rule.ControlBehavior < 0 {
+		return errors.New("invalid control strategy")
 	}
 
 	if rule.RelationStrategy == AssociatedResource && rule.RefResource == "" {
 		return errors.New("Bad flow rule: invalid control behavior")
 	}
 
-	return checkControlBehaviorField(rule)
-}
-
-func checkControlBehaviorField(rule *Rule) error {
-	switch rule.ControlBehavior {
-	case WarmUp:
+	if rule.TokenCalculateStrategy == WarmUp {
 		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 {
-			return errors.New("invalid warmUpPeriodSec")
-		}
-		return nil
-	default:
 	}
 	return nil
 }
diff --git a/core/flow/tc_default.go b/core/flow/tc_default.go
index c767e69..e8ac91f 100644
--- a/core/flow/tc_default.go
+++ b/core/flow/tc_default.go
@@ -4,15 +4,15 @@ import (
 	"github.com/alibaba/sentinel-golang/core/base"
 )
 
-type DefaultTrafficShapingCalculator struct {
+type DirectTrafficShapingCalculator struct {
 	threshold float64
 }
 
-func NewDefaultTrafficShapingCalculator(threshold float64) *DefaultTrafficShapingCalculator {
-	return &DefaultTrafficShapingCalculator{threshold: threshold}
+func NewDirectTrafficShapingCalculator(threshold float64) *DirectTrafficShapingCalculator {
+	return &DirectTrafficShapingCalculator{threshold: threshold}
 }
 
-func (d *DefaultTrafficShapingCalculator) CalculateAllowedTokens(base.StatNode, uint32, int32) float64 {
+func (d *DirectTrafficShapingCalculator) CalculateAllowedTokens(base.StatNode, uint32, int32) float64 {
 	return d.threshold
 }
 
diff --git a/example/qps/qps_limit_example.go b/example/qps/qps_limit_example.go
index e7d5270..9ccd96e 100644
--- a/example/qps/qps_limit_example.go
+++ b/example/qps/qps_limit_example.go
@@ -21,10 +21,11 @@ func main() {
 
 	_, err = flow.LoadRules([]*flow.Rule{
 		{
-			Resource:        "some-test",
-			MetricType:      flow.QPS,
-			Count:           10,
-			ControlBehavior: flow.Reject,
+			Resource:               "some-test",
+			MetricType:             flow.QPS,
+			Count:                  10,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
 		},
 	})
 	if err != nil {
diff --git a/example/warm_up/qps_warm_up_example.go b/example/warm_up/qps_warm_up_example.go
index 0f5563a..9dddc89 100644
--- a/example/warm_up/qps_warm_up_example.go
+++ b/example/warm_up/qps_warm_up_example.go
@@ -33,12 +33,13 @@ func main() {
 
 	_, err = flow.LoadRules([]*flow.Rule{
 		{
-			Resource:         "some-test",
-			MetricType:       flow.QPS,
-			Count:            100,
-			ControlBehavior:  flow.WarmUp,
-			WarmUpPeriodSec:  10,
-			WarmUpColdFactor: 3,
+			Resource:               "some-test",
+			MetricType:             flow.QPS,
+			Count:                  100,
+			TokenCalculateStrategy: flow.WarmUp,
+			ControlBehavior:        flow.Reject,
+			WarmUpPeriodSec:        10,
+			WarmUpColdFactor:       3,
 		},
 	})
 	if err != nil {


Test patch:
diff --git a/adapter/echo/middleware_test.go b/adapter/echo/middleware_test.go
index 686f0bf..9e4cbeb 100644
--- a/adapter/echo/middleware_test.go
+++ b/adapter/echo/middleware_test.go
@@ -20,16 +20,18 @@ func initSentinel(t *testing.T) {
 
 	_, err = flow.LoadRules([]*flow.Rule{
 		{
-			Resource:        "GET:/ping",
-			MetricType:      flow.QPS,
-			Count:           1,
-			ControlBehavior: flow.Reject,
+			Resource:               "GET:/ping",
+			MetricType:             flow.QPS,
+			Count:                  1,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
 		},
 		{
-			Resource:        "/api/:uid",
-			MetricType:      flow.QPS,
-			Count:           0,
-			ControlBehavior: flow.Reject,
+			Resource:               "/api/:uid",
+			MetricType:             flow.QPS,
+			Count:                  0,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
 		},
 	})
 	if err != nil {
diff --git a/adapter/gin/middleware_test.go b/adapter/gin/middleware_test.go
index a459331..9f31217 100644
--- a/adapter/gin/middleware_test.go
+++ b/adapter/gin/middleware_test.go
@@ -20,16 +20,18 @@ func initSentinel(t *testing.T) {
 
 	_, err = flow.LoadRules([]*flow.Rule{
 		{
-			Resource:        "GET:/ping",
-			MetricType:      flow.QPS,
-			Count:           1,
-			ControlBehavior: flow.Reject,
+			Resource:               "GET:/ping",
+			MetricType:             flow.QPS,
+			Count:                  1,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
 		},
 		{
-			Resource:        "/api/users/:id",
-			MetricType:      flow.QPS,
-			Count:           0,
-			ControlBehavior: flow.Reject,
+			Resource:               "/api/users/:id",
+			MetricType:             flow.QPS,
+			Count:                  0,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
 		},
 	})
 	if err != nil {
diff --git a/adapter/grpc/client_test.go b/adapter/grpc/client_test.go
index 3b4d4f4..26368f6 100644
--- a/adapter/grpc/client_test.go
+++ b/adapter/grpc/client_test.go
@@ -25,10 +25,11 @@ func TestUnaryClientIntercept(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        "client:" + method,
-				MetricType:      flow.QPS,
-				Count:           1,
-				ControlBehavior: flow.Reject,
+				Resource:               "client:" + method,
+				MetricType:             flow.QPS,
+				Count:                  1,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
@@ -43,10 +44,11 @@ func TestUnaryClientIntercept(t *testing.T) {
 	t.Run("fail", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        "client:" + method,
-				MetricType:      flow.QPS,
-				Count:           0,
-				ControlBehavior: flow.Reject,
+				Resource:               "client:" + method,
+				MetricType:             flow.QPS,
+				Count:                  0,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
@@ -69,10 +71,11 @@ func TestStreamClientIntercept(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        "client:/grpc.testing.TestService/StreamingOutputCall",
-				MetricType:      flow.QPS,
-				Count:           1,
-				ControlBehavior: flow.Reject,
+				Resource:               "client:/grpc.testing.TestService/StreamingOutputCall",
+				MetricType:             flow.QPS,
+				Count:                  1,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
@@ -89,10 +92,11 @@ func TestStreamClientIntercept(t *testing.T) {
 	t.Run("fail", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        "client:/grpc.testing.TestService/StreamingOutputCall",
-				MetricType:      flow.QPS,
-				Count:           0,
-				ControlBehavior: flow.Reject,
+				Resource:               "client:/grpc.testing.TestService/StreamingOutputCall",
+				MetricType:             flow.QPS,
+				Count:                  0,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
diff --git a/adapter/grpc/server_test.go b/adapter/grpc/server_test.go
index 8292921..d9ec2bc 100644
--- a/adapter/grpc/server_test.go
+++ b/adapter/grpc/server_test.go
@@ -31,10 +31,11 @@ func TestStreamServerIntercept(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        "/grpc.testing.TestService/StreamingInputCall",
-				MetricType:      flow.QPS,
-				Count:           1,
-				ControlBehavior: flow.Reject,
+				Resource:               "/grpc.testing.TestService/StreamingInputCall",
+				MetricType:             flow.QPS,
+				Count:                  1,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
@@ -49,10 +50,11 @@ func TestStreamServerIntercept(t *testing.T) {
 	t.Run("fail", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        "/grpc.testing.TestService/StreamingInputCall",
-				MetricType:      flow.QPS,
-				Count:           0,
-				ControlBehavior: flow.Reject,
+				Resource:               "/grpc.testing.TestService/StreamingInputCall",
+				MetricType:             flow.QPS,
+				Count:                  0,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
@@ -73,10 +75,11 @@ func TestUnaryServerIntercept(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        "/grpc.testing.TestService/UnaryCall",
-				MetricType:      flow.QPS,
-				Count:           1,
-				ControlBehavior: flow.Reject,
+				Resource:               "/grpc.testing.TestService/UnaryCall",
+				MetricType:             flow.QPS,
+				Count:                  1,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
@@ -101,10 +104,11 @@ func TestUnaryServerIntercept(t *testing.T) {
 	t.Run("fail", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        "/grpc.testing.TestService/UnaryCall",
-				MetricType:      flow.QPS,
-				Count:           0,
-				ControlBehavior: flow.Reject,
+				Resource:               "/grpc.testing.TestService/UnaryCall",
+				MetricType:             flow.QPS,
+				Count:                  0,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
diff --git a/adapter/micro/client_test.go b/adapter/micro/client_test.go
index 14b588e..b8cb781 100644
--- a/adapter/micro/client_test.go
+++ b/adapter/micro/client_test.go
@@ -47,10 +47,11 @@ func TestClientLimiter(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        req.Method(),
-				MetricType:      flow.QPS,
-				Count:           1,
-				ControlBehavior: flow.Reject,
+				Resource:               req.Method(),
+				MetricType:             flow.QPS,
+				Count:                  1,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
diff --git a/adapter/micro/server_test.go b/adapter/micro/server_test.go
index 9a3912d..b0a0c7b 100644
--- a/adapter/micro/server_test.go
+++ b/adapter/micro/server_test.go
@@ -57,10 +57,11 @@ func TestServerLimiter(t *testing.T) {
 
 	_, err = flow.LoadRules([]*flow.Rule{
 		{
-			Resource:        req.Method(),
-			MetricType:      flow.QPS,
-			Count:           1,
-			ControlBehavior: flow.Reject,
+			Resource:               req.Method(),
+			MetricType:             flow.QPS,
+			Count:                  1,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
 		},
 	})
 
@@ -71,10 +72,11 @@ func TestServerLimiter(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
 		var _, err = flow.LoadRules([]*flow.Rule{
 			{
-				Resource:        req.Method(),
-				MetricType:      flow.QPS,
-				Count:           1,
-				ControlBehavior: flow.Reject,
+				Resource:               req.Method(),
+				MetricType:             flow.QPS,
+				Count:                  1,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
 			},
 		})
 		assert.Nil(t, err)
diff --git a/core/flow/rule_manager_test.go b/core/flow/rule_manager_test.go
index b308e0d..f11c4a9 100644
--- a/core/flow/rule_manager_test.go
+++ b/core/flow/rule_manager_test.go
@@ -1,23 +1,23 @@
 package flow
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestSetAndRemoveTrafficShapingGenerator(t *testing.T) {
-	tsc := NewTrafficShapingController(nil, nil, nil)
+	tsc := &TrafficShapingController{}
 
-	err := SetTrafficShapingGenerator(Reject, func(_ *Rule) *TrafficShapingController {
+	err := SetTrafficShapingGenerator(Direct, Reject, func(_ *Rule) *TrafficShapingController {
 		return tsc
 	})
 	assert.Error(t, err, "default control behaviors are not allowed to be modified")
-	err = RemoveTrafficShapingGenerator(Reject)
+	err = RemoveTrafficShapingGenerator(Direct, Reject)
 	assert.Error(t, err, "default control behaviors are not allowed to be removed")
 
-	cb := ControlBehavior(9999)
-	err = SetTrafficShapingGenerator(cb, func(_ *Rule) *TrafficShapingController {
+	err = SetTrafficShapingGenerator(TokenCalculateStrategy(111), ControlBehavior(112), func(_ *Rule) *TrafficShapingController {
 		return tsc
 	})
 	assert.NoError(t, err)
@@ -25,21 +25,27 @@ func TestSetAndRemoveTrafficShapingGenerator(t *testing.T) {
 	resource := "test-customized-tc"
 	_, err = LoadRules([]*Rule{
 		{
-			ID:              10,
-			Count:           20,
-			MetricType:      QPS,
-			Resource:        resource,
-			ControlBehavior: cb,
+			ID:                     10,
+			Count:                  20,
+			MetricType:             QPS,
+			Resource:               resource,
+			TokenCalculateStrategy: TokenCalculateStrategy(111),
+			ControlBehavior:        ControlBehavior(112),
 		},
 	})
+
+	cs := trafficControllerGenKey{
+		tokenCalculateStrategy: TokenCalculateStrategy(111),
+		controlBehavior:        ControlBehavior(112),
+	}
 	assert.NoError(t, err)
-	assert.Contains(t, tcGenFuncMap, cb)
+	assert.Contains(t, tcGenFuncMap, cs)
 	assert.NotZero(t, len(tcMap[resource]))
 	assert.Equal(t, tsc, tcMap[resource][0])
 
-	err = RemoveTrafficShapingGenerator(cb)
+	err = RemoveTrafficShapingGenerator(TokenCalculateStrategy(111), ControlBehavior(112))
 	assert.NoError(t, err)
-	assert.NotContains(t, tcGenFuncMap, cb)
+	assert.NotContains(t, tcGenFuncMap, cs)
 
 	_, _ = LoadRules([]*Rule{})
 }
@@ -47,11 +53,107 @@ func TestSetAndRemoveTrafficShapingGenerator(t *testing.T) {
 func TestIsValidFlowRule(t *testing.T) {
 	badRule1 := &Rule{ID: 1, Count: 1, MetricType: QPS, Resource: ""}
 	badRule2 := &Rule{ID: 1, Count: -1.9, MetricType: QPS, Resource: "test"}
-	badRule3 := &Rule{Count: 5, MetricType: QPS, Resource: "test", ControlBehavior: WarmUp}
-	goodRule1 := &Rule{Count: 10, MetricType: QPS, Resource: "test", ControlBehavior: Throttling}
+	badRule3 := &Rule{Count: 5, MetricType: QPS, Resource: "test", TokenCalculateStrategy: WarmUp, ControlBehavior: Reject}
+	goodRule1 := &Rule{Count: 10, MetricType: QPS, Resource: "test", TokenCalculateStrategy: WarmUp, ControlBehavior: Throttling, WarmUpPeriodSec: 10}
+
+	assert.Error(t, IsValidRule(badRule1))
+	assert.Error(t, IsValidRule(badRule2))
+	assert.Error(t, IsValidRule(badRule3))
+	assert.NoError(t, IsValidRule(goodRule1))
+}
+
+func TestGetRules(t *testing.T) {
+	t.Run("GetRules", func(t *testing.T) {
+		if err := ClearRules(); err != nil {
+			t.Fatal(err)
+		}
+		r1 := &Rule{
+			ID:                     0,
+			Resource:               "abc1",
+			MetricType:             0,
+			Count:                  0,
+			RelationStrategy:       0,
+			TokenCalculateStrategy: Direct,
+			ControlBehavior:        Reject,
+			RefResource:            "",
+			WarmUpPeriodSec:        0,
+			MaxQueueingTimeMs:      0,
+		}
+		r2 := &Rule{
+			ID:                     1,
+			Resource:               "abc2",
+			MetricType:             0,
+			Count:                  0,
+			RelationStrategy:       0,
+			TokenCalculateStrategy: Direct,
+			ControlBehavior:        Throttling,
+			RefResource:            "",
+			WarmUpPeriodSec:        0,
+			MaxQueueingTimeMs:      0,
+		}
+		if _, err := LoadRules([]*Rule{r1, r2}); err != nil {
+			t.Fatal(err)
+		}
 
-	assert.Error(t, IsValidFlowRule(badRule1))
-	assert.Error(t, IsValidFlowRule(badRule2))
-	assert.Error(t, IsValidFlowRule(badRule3))
-	assert.NoError(t, IsValidFlowRule(goodRule1))
+		rs1 := GetRules()
+		if rs1[0].Resource == "abc1" {
+			assert.True(t, &rs1[0] != r1)
+			assert.True(t, &rs1[1] != r2)
+			assert.True(t, reflect.DeepEqual(&rs1[0], r1))
+			assert.True(t, reflect.DeepEqual(&rs1[1], r2))
+		} else {
+			assert.True(t, &rs1[0] != r2)
+			assert.True(t, &rs1[1] != r1)
+			assert.True(t, reflect.DeepEqual(&rs1[0], r2))
+			assert.True(t, reflect.DeepEqual(&rs1[1], r1))
+		}
+		if err := ClearRules(); err != nil {
+			t.Fatal(err)
+		}
+	})
+
+	t.Run("getRules", func(t *testing.T) {
+		r1 := &Rule{
+			ID:                     0,
+			Resource:               "abc1",
+			MetricType:             0,
+			Count:                  0,
+			RelationStrategy:       0,
+			TokenCalculateStrategy: Direct,
+			ControlBehavior:        Reject,
+			RefResource:            "",
+			WarmUpPeriodSec:        0,
+			MaxQueueingTimeMs:      0,
+		}
+		r2 := &Rule{
+			ID:                     1,
+			Resource:               "abc2",
+			MetricType:             0,
+			Count:                  0,
+			RelationStrategy:       0,
+			TokenCalculateStrategy: Direct,
+			ControlBehavior:        Throttling,
+			RefResource:            "",
+			WarmUpPeriodSec:        0,
+			MaxQueueingTimeMs:      0,
+		}
+		if _, err := LoadRules([]*Rule{r1, r2}); err != nil {
+			t.Fatal(err)
+		}
+		rs2 := getRules()
+		if rs2[0].Resource == "abc1" {
+			assert.True(t, rs2[0] == r1)
+			assert.True(t, rs2[1] == r2)
+			assert.True(t, reflect.DeepEqual(rs2[0], r1))
+			assert.True(t, reflect.DeepEqual(rs2[1], r2))
+		} else {
+			assert.True(t, rs2[0] == r2)
+			assert.True(t, rs2[1] == r1)
+			assert.True(t, reflect.DeepEqual(rs2[0], r2))
+			assert.True(t, reflect.DeepEqual(rs2[1], r1))
+		}
+		if err := ClearRules(); err != nil {
+			t.Fatal(err)
+		}
+	})
 }
diff --git a/example/qps/qps_limit_example_test.go b/example/qps/qps_limit_example_test.go
index f821e4a..98ea497 100644
--- a/example/qps/qps_limit_example_test.go
+++ b/example/qps/qps_limit_example_test.go
@@ -28,10 +28,11 @@ func doTest() {
 
 	_, err = flow.LoadRules([]*flow.Rule{
 		{
-			Resource:        "some-test",
-			MetricType:      flow.QPS,
-			Count:           100,
-			ControlBehavior: flow.Reject,
+			Resource:               "some-test",
+			MetricType:             flow.QPS,
+			Count:                  100,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
 		},
 	})
 	if err != nil {
diff --git a/example/warm_up/qps_warm_up_example_test.go b/example/warm_up/qps_warm_up_example_test.go
index 075502b..0f2c9f4 100644
--- a/example/warm_up/qps_warm_up_example_test.go
+++ b/example/warm_up/qps_warm_up_example_test.go
@@ -30,11 +30,12 @@ func doTest() {
 
 	_, err = flow.LoadRules([]*flow.Rule{
 		{
-			Resource:        "some-test",
-			MetricType:      flow.QPS,
-			Count:           100,
-			ControlBehavior: flow.WarmUp,
-			WarmUpPeriodSec: 10,
+			Resource:               "some-test",
+			MetricType:             flow.QPS,
+			Count:                  100,
+			TokenCalculateStrategy: flow.WarmUp,
+			ControlBehavior:        flow.Reject,
+			WarmUpPeriodSec:        10,
 		},
 	})
 	if err != nil {
diff --git a/ext/datasource/helper_test.go b/ext/datasource/helper_test.go
index 5984bf9..e36ba1e 100644
--- a/ext/datasource/helper_test.go
+++ b/ext/datasource/helper_test.go
@@ -47,38 +47,41 @@ func TestFlowRulesJsonConverter(t *testing.T) {
 		flowRules := got.([]*flow.Rule)
 		assert.True(t, len(flowRules) == 3)
 		r1 := &flow.Rule{
-			Resource:          "abc",
-			MetricType:        flow.Concurrency,
-			Count:             100,
-			RelationStrategy:  flow.Direct,
-			ControlBehavior:   flow.Reject,
-			RefResource:       "refDefault",
-			WarmUpPeriodSec:   10,
-			MaxQueueingTimeMs: 1000,
+			Resource:               "abc",
+			MetricType:             flow.Concurrency,
+			Count:                  100,
+			RelationStrategy:       flow.CurrentResource,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
+			RefResource:            "refDefault",
+			WarmUpPeriodSec:        10,
+			MaxQueueingTimeMs:      1000,
 		}
 		assert.True(t, reflect.DeepEqual(flowRules[0], r1))
 
 		r2 := &flow.Rule{
-			Resource:          "abc",
-			MetricType:        flow.QPS,
-			Count:             200,
-			RelationStrategy:  flow.AssociatedResource,
-			ControlBehavior:   flow.WarmUp,
-			RefResource:       "refDefault",
-			WarmUpPeriodSec:   20,
-			MaxQueueingTimeMs: 2000,
+			Resource:               "abc",
+			MetricType:             flow.QPS,
+			Count:                  200,
+			RelationStrategy:       flow.AssociatedResource,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Throttling,
+			RefResource:            "refDefault",
+			WarmUpPeriodSec:        20,
+			MaxQueueingTimeMs:      2000,
 		}
 		assert.True(t, reflect.DeepEqual(flowRules[1], r2))
 
 		r3 := &flow.Rule{
-			Resource:          "abc",
-			MetricType:        flow.QPS,
-			Count:             300,
-			RelationStrategy:  flow.Direct,
-			ControlBehavior:   flow.WarmUp,
-			RefResource:       "refDefault",
-			WarmUpPeriodSec:   30,
-			MaxQueueingTimeMs: 3000,
+			Resource:               "abc",
+			MetricType:             flow.QPS,
+			Count:                  300,
+			RelationStrategy:       flow.CurrentResource,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Throttling,
+			RefResource:            "refDefault",
+			WarmUpPeriodSec:        30,
+			MaxQueueingTimeMs:      3000,
 		}
 		assert.True(t, reflect.DeepEqual(flowRules[2], r3))
 	})
@@ -89,15 +92,16 @@ func TestFlowRulesUpdater(t *testing.T) {
 		flow.ClearRules()
 		flow.LoadRules([]*flow.Rule{
 			{
-				ID:                0,
-				Resource:          "abc",
-				MetricType:        0,
-				Count:             0,
-				RelationStrategy:  0,
-				ControlBehavior:   0,
-				RefResource:       "",
-				WarmUpPeriodSec:   0,
-				MaxQueueingTimeMs: 0,
+				ID:                     0,
+				Resource:               "abc",
+				MetricType:             0,
+				Count:                  0,
+				RelationStrategy:       0,
+				TokenCalculateStrategy: flow.Direct,
+				ControlBehavior:        flow.Reject,
+				RefResource:            "",
+				WarmUpPeriodSec:        0,
+				MaxQueueingTimeMs:      0,
 			}})
 		assert.True(t, len(flow.GetRules()) == 1, "Fail to prepare test data.")
 		err := FlowRulesUpdater(nil)
@@ -121,15 +125,16 @@ func TestFlowRulesUpdater(t *testing.T) {
 		flow.ClearRules()
 		p := make([]flow.Rule, 0)
 		fw := flow.Rule{
-			ID:                0,
-			Resource:          "aaaa",
-			MetricType:        0,
-			Count:             0,
-			RelationStrategy:  0,
-			ControlBehavior:   0,
-			RefResource:       "",
-			WarmUpPeriodSec:   0,
-			MaxQueueingTimeMs: 0,
+			ID:                     0,
+			Resource:               "aaaa",
+			MetricType:             0,
+			Count:                  0,
+			RelationStrategy:       0,
+			TokenCalculateStrategy: flow.Direct,
+			ControlBehavior:        flow.Reject,
+			RefResource:            "",
+			WarmUpPeriodSec:        0,
+			MaxQueueingTimeMs:      0,
 		}
 		p = append(p, fw)
 		err := FlowRulesUpdater(p)
diff --git a/tests/testdata/extension/helper/FlowRule.json b/tests/testdata/extension/helper/FlowRule.json
index ad37323..d4f2b73 100644
--- a/tests/testdata/extension/helper/FlowRule.json
+++ b/tests/testdata/extension/helper/FlowRule.json
@@ -4,6 +4,7 @@
         "metricType": 0,
         "count": 100.0,
         "relationStrategy": 0,
+        "tokenCalculateStrategy": 0,
         "controlBehavior": 0,
         "refResource": "refDefault",
         "warmUpPeriodSec": 10,
@@ -14,6 +15,7 @@
         "metricType": 1,
         "count": 200.0,
         "relationStrategy": 1,
+        "tokenCalculateStrategy": 0,
         "controlBehavior": 1,
         "refResource": "refDefault",
         "warmUpPeriodSec": 20,
@@ -24,6 +26,7 @@
         "metricType": 1,
         "count": 300.0,
         "relationStrategy": 0,
+        "tokenCalculateStrategy": 0,
         "controlBehavior": 1,
         "refResource": "refDefault",
         "warmUpPeriodSec": 30,