Repository: alibaba/sentinel-golang
Pull request #363 (2020-12-22 11:16:53)

feat: support hotspot param extract with key
<!--  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

see : https://github.com/alibaba/sentinel-golang/issues/355
### 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/363?src=pr&el=h1) Report
> :exclamation: No coverage uploaded for pull request base (`mosn-customized@19fad24`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference#section-missing-base-commit).
> The diff coverage is `n/a`.

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

```diff
@@                Coverage Diff                 @@
##             mosn-customized     #363   +/-   ##
==================================================
  Coverage                   ?   49.90%           
==================================================
  Files                      ?       74           
  Lines                      ?     4078           
  Branches                   ?        0           
==================================================
  Hits                       ?     2035           
  Misses                     ?     1766           
  Partials                   ?      277           
```



------

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


Patch:
diff --git a/core/hotspot/concurrency_stat_slot.go b/core/hotspot/concurrency_stat_slot.go
index ad64ae1..ba6c642 100644
--- a/core/hotspot/concurrency_stat_slot.go
+++ b/core/hotspot/concurrency_stat_slot.go
@@ -15,8 +15,11 @@
 package hotspot
 
 import (
+	"context"
 	"sync/atomic"
 
+	"golang.org/x/sync/errgroup"
+
 	"github.com/alibaba/sentinel-golang/core/base"
 	"github.com/alibaba/sentinel-golang/logging"
 )
@@ -44,26 +47,35 @@ func (s *ConcurrencyStatSlot) Order() uint32 {
 
 func (c *ConcurrencyStatSlot) OnEntryPassed(ctx *base.EntryContext) {
 	res := ctx.Resource.Name()
-	args := ctx.Input.Args
+	g, _ := errgroup.WithContext(context.Background())
 	tcs := getTrafficControllersFor(res)
+
 	for _, tc := range tcs {
 		if tc.BoundRule().MetricType != Concurrency {
 			continue
 		}
-		arg := matchArg(tc, args)
-		if arg == nil {
+		args := tc.ExtractArgs(ctx)
+		if args == nil || len(args) == 0 {
 			continue
 		}
-		metric := tc.BoundMetric()
-		concurrencyPtr, existed := metric.ConcurrencyCounter.Get(arg)
-		if !existed || concurrencyPtr == nil {
-			if logging.DebugEnabled() {
-				logging.Debug("[ConcurrencyStatSlot OnEntryPassed] Parameter does not exist in ConcurrencyCounter.", "argument", arg)
-			}
-			continue
+		for _, arg := range args {
+			arg := arg // https://golang.org/doc/faq#closures_and_goroutines
+			g.Go(func() error {
+				metric := tc.BoundMetric()
+				concurrencyPtr, existed := metric.ConcurrencyCounter.Get(arg)
+				if !existed || concurrencyPtr == nil {
+					if logging.DebugEnabled() {
+						logging.Debug("[ConcurrencyStatSlot OnEntryPassed] Parameter does not exist in ConcurrencyCounter.", "argument", arg)
+					}
+					return nil
+				}
+				atomic.AddInt64(concurrencyPtr, 1)
+				return nil
+			})
 		}
-		atomic.AddInt64(concurrencyPtr, 1)
 	}
+	g.Wait()
+	return
 }
 
 func (c *ConcurrencyStatSlot) OnEntryBlocked(ctx *base.EntryContext, blockError *base.BlockError) {
@@ -72,24 +84,33 @@ func (c *ConcurrencyStatSlot) OnEntryBlocked(ctx *base.EntryContext, blockError
 
 func (c *ConcurrencyStatSlot) OnCompleted(ctx *base.EntryContext) {
 	res := ctx.Resource.Name()
-	args := ctx.Input.Args
+	g, _ := errgroup.WithContext(context.Background())
 	tcs := getTrafficControllersFor(res)
+
 	for _, tc := range tcs {
 		if tc.BoundRule().MetricType != Concurrency {
 			continue
 		}
-		arg := matchArg(tc, args)
-		if arg == nil {
+		args := tc.ExtractArgs(ctx)
+		if args == nil || len(args) == 0 {
 			continue
 		}
-		metric := tc.BoundMetric()
-		concurrencyPtr, existed := metric.ConcurrencyCounter.Get(arg)
-		if !existed || concurrencyPtr == nil {
-			if logging.DebugEnabled() {
-				logging.Debug("[ConcurrencyStatSlot OnCompleted] Parameter does not exist in ConcurrencyCounter.", "argument", arg)
-			}
-			continue
+		for _, arg := range args {
+			arg := arg // https://golang.org/doc/faq#closures_and_goroutines
+			g.Go(func() error {
+				metric := tc.BoundMetric()
+				concurrencyPtr, existed := metric.ConcurrencyCounter.Get(arg)
+				if !existed || concurrencyPtr == nil {
+					if logging.DebugEnabled() {
+						logging.Debug("[ConcurrencyStatSlot OnCompleted] Parameter does not exist in ConcurrencyCounter.", "argument", arg)
+					}
+					return nil
+				}
+				atomic.AddInt64(concurrencyPtr, -1)
+				return nil
+			})
 		}
-		atomic.AddInt64(concurrencyPtr, -1)
 	}
+	g.Wait()
+	return
 }
diff --git a/core/hotspot/rule.go b/core/hotspot/rule.go
index 97917bb..b726587 100644
--- a/core/hotspot/rule.go
+++ b/core/hotspot/rule.go
@@ -80,6 +80,9 @@ type Rule struct {
 	// if ParamIndex is great than or equals to zero, ParamIndex means the <ParamIndex>-th parameter
 	// if ParamIndex is the negative, ParamIndex means the reversed <ParamIndex>-th parameter
 	ParamIndex int `json:"paramIndex"`
+	// ParamKey is the key in context attachments map.
+	// ParamKey can be used as a supplement to ParamIndex to facilitate rules to quickly obtain parameter from a large number of parameters
+	ParamKey string `json:"paramKey"`
 	// Threshold is the threshold to trigger rejection
 	Threshold int64 `json:"threshold"`
 	// MaxQueueingTimeMs only takes effect when ControlBehavior is Throttling and MetricType is QPS
diff --git a/core/hotspot/slot.go b/core/hotspot/slot.go
index e6ac4c9..95920b1 100644
--- a/core/hotspot/slot.go
+++ b/core/hotspot/slot.go
@@ -15,10 +15,13 @@
 package hotspot
 
 import (
+	"context"
 	"time"
 
+	"golang.org/x/sync/errgroup"
+	"gopkg.in/errgo.v2/fmt/errors"
+
 	"github.com/alibaba/sentinel-golang/core/base"
-	"github.com/alibaba/sentinel-golang/logging"
 )
 
 const (
@@ -41,58 +44,43 @@ func (s *Slot) Order() uint32 {
 	return RuleCheckSlotOrder
 }
 
-// matchArg matches the arg from args based on TrafficShapingController
-// return nil if match failed.
-func matchArg(tc TrafficShapingController, args []interface{}) interface{} {
-	if tc == nil {
-		return nil
-	}
-	idx := tc.BoundParamIndex()
-	if idx < 0 {
-		idx = len(args) + idx
-	}
-	if idx < 0 {
-		if logging.DebugEnabled() {
-			logging.Debug("[Slot matchArg] The param index of hotspot traffic shaping controller is invalid", "args", args, "paramIndex", tc.BoundParamIndex())
-		}
-		return nil
-	}
-	if idx >= len(args) {
-		if logging.DebugEnabled() {
-			logging.Debug("[Slot matchArg] The argument in index doesn't exist", "args", args, "paramIndex", tc.BoundParamIndex())
-		}
-		return nil
-	}
-	return args[idx]
-}
-
 func (s *Slot) Check(ctx *base.EntryContext) *base.TokenResult {
 	res := ctx.Resource.Name()
-	args := ctx.Input.Args
 	batch := int64(ctx.Input.BatchCount)
+	g, _ := errgroup.WithContext(context.Background())
 
 	result := ctx.RuleCheckResult
 	tcs := getTrafficControllersFor(res)
 	for _, tc := range tcs {
-		arg := matchArg(tc, args)
-		if arg == nil {
-			continue
-		}
-		r := canPassCheck(tc, arg, batch)
-		if r == nil {
+		args := tc.ExtractArgs(ctx)
+		if args == nil || len(args) == 0 {
 			continue
 		}
-		if r.Status() == base.ResultStatusBlocked {
-			return r
-		}
-		if r.Status() == base.ResultStatusShouldWait {
-			if nanosToWait := r.NanosToWait(); nanosToWait > 0 {
-				// Handle waiting action.
-				time.Sleep(nanosToWait)
-			}
-			continue
+		for _, arg := range args {
+			arg := arg // https://golang.org/doc/faq#closures_and_goroutines
+			g.Go(func() error {
+				r := canPassCheck(tc, arg, batch)
+				if r == nil {
+					return nil
+				}
+				if r.Status() == base.ResultStatusBlocked {
+					return errors.Newf("concurrent canPassCheck err=%v", r.BlockError())
+				}
+				if r.Status() == base.ResultStatusShouldWait {
+					if nanosToWait := r.NanosToWait(); nanosToWait > 0 {
+						// Handle waiting action.
+						time.Sleep(nanosToWait)
+					}
+					return nil
+				}
+				return nil
+			})
 		}
 	}
+	if err := g.Wait(); err != nil {
+		result.ResetToBlockedWithMessage(base.BlockTypeHotSpotParamFlow, err.Error())
+		return result
+	}
 	return result
 }
 
diff --git a/core/hotspot/traffic_shaping.go b/core/hotspot/traffic_shaping.go
index 7ba5865..9f14f05 100644
--- a/core/hotspot/traffic_shaping.go
+++ b/core/hotspot/traffic_shaping.go
@@ -33,6 +33,8 @@ type TrafficShapingController interface {
 
 	BoundParamIndex() int
 
+	ExtractArgs(ctx *base.EntryContext) []interface{}
+
 	BoundMetric() *ParamsMetric
 
 	BoundRule() *Rule
@@ -44,6 +46,7 @@ type baseTrafficShapingController struct {
 	res           string
 	metricType    MetricType
 	paramIndex    int
+	paramKey      string
 	threshold     int64
 	specificItems map[interface{}]int64
 	durationInSec int64
@@ -60,6 +63,7 @@ func newBaseTrafficShapingControllerWithMetric(r *Rule, metric *ParamsMetric) *b
 		res:           r.Resource,
 		metricType:    r.MetricType,
 		paramIndex:    r.ParamIndex,
+		paramKey:      r.ParamKey,
 		threshold:     r.Threshold,
 		specificItems: r.SpecificItems,
 		durationInSec: r.DurationInSec,
@@ -155,6 +159,74 @@ func (c *baseTrafficShapingController) BoundParamIndex() int {
 	return c.paramIndex
 }
 
+// ExtractArgs matches the arg from ctx based on TrafficShapingController
+// return nil if match failed.
+func (c *baseTrafficShapingController) ExtractArgs(ctx *base.EntryContext) []interface{} {
+	if c == nil {
+		return nil
+	}
+	args := make([]interface{}, 0, len(ctx.Input.Args)+len(ctx.Input.Attachments))
+
+	indexArg := c.extractArgs(ctx)
+	if indexArg != nil {
+		args = append(args, indexArg)
+	}
+	indexArg = c.extractAttachmentArgs(ctx)
+	if indexArg != nil {
+		args = append(args, indexArg)
+	}
+	return args
+}
+func (c *baseTrafficShapingController) extractArgs(ctx *base.EntryContext) interface{} {
+	args := ctx.Input.Args
+	idx := c.BoundParamIndex()
+	if idx < 0 {
+		idx = len(args) + idx
+	}
+	if idx < 0 {
+		if logging.DebugEnabled() {
+			logging.Debug("[extractArgs] The param index of hotspot traffic shaping controller is invalid",
+				"args", args, "paramIndex", c.BoundParamIndex())
+		}
+		return nil
+	}
+	if idx >= len(args) {
+		if logging.DebugEnabled() {
+			logging.Debug("[extractArgs] The argument in index doesn't exist",
+				"args", args, "paramIndex", c.BoundParamIndex())
+		}
+		return nil
+	}
+	return args[idx]
+}
+func (c *baseTrafficShapingController) extractAttachmentArgs(ctx *base.EntryContext) interface{} {
+	attachments := ctx.Input.Attachments
+
+	if attachments == nil {
+		if logging.DebugEnabled() {
+			logging.Debug("[paramKey] The attachments of ctx is nil",
+				"args", attachments, "paramIndex", c.paramKey)
+		}
+		return nil
+	}
+	if c.paramKey == "" {
+		if logging.DebugEnabled() {
+			logging.Debug("[paramKey] The param key is nil",
+				"args", attachments, "paramIndex", c.paramKey)
+		}
+		return nil
+	}
+	arg, ok := attachments[c.paramKey]
+	if !ok {
+		if logging.DebugEnabled() {
+			logging.Debug("[paramKey] extracted data does not exist",
+				"args", attachments, "paramIndex", c.paramKey)
+		}
+	}
+
+	return arg
+}
+
 func (c *rejectTrafficShapingController) PerformChecking(arg interface{}, batchCount int64) *base.TokenResult {
 	metric := c.metric
 	if metric == nil {
diff --git a/example/hotspot_param_flow/concurrency/hotspot_params_concurrency_example.go b/example/hotspot_param_flow/concurrency/hotspot_params_concurrency_example.go
index f9ea4e9..64e8d00 100644
--- a/example/hotspot_param_flow/concurrency/hotspot_params_concurrency_example.go
+++ b/example/hotspot_param_flow/concurrency/hotspot_params_concurrency_example.go
@@ -35,6 +35,7 @@ func main() {
 	if err != nil {
 		log.Fatal(err)
 	}
+	testKey := "testKey"
 
 	// the max concurrency is 8
 	_, err = hotspot.LoadRules([]*hotspot.Rule{
@@ -42,6 +43,7 @@ func main() {
 			Resource:      "abc",
 			MetricType:    hotspot.Concurrency,
 			ParamIndex:    0,
+			ParamKey:      testKey,
 			Threshold:     8,
 			DurationInSec: 0,
 		},
@@ -79,7 +81,12 @@ func main() {
 	}
 
 	for {
-		e, b := sentinel.Entry("abc", sentinel.WithArgs(false, uint32(9)))
+		e, b := sentinel.Entry("abc",
+			sentinel.WithArgs(false, uint32(9),
+				sentinel.WithAttachments(map[interface{}]interface{}{
+					testKey: rand.Uint64() % 10,
+				}),
+			))
 		if b != nil {
 			// Blocked. We could get the block reason from the BlockError.
 			time.Sleep(time.Duration(rand.Uint64()%10) * time.Millisecond)
diff --git a/example/hotspot_param_flow/qps_reject/hotspot_params_qps_reject_example.go b/example/hotspot_param_flow/qps_reject/hotspot_params_qps_reject_example.go
index 0e38e24..30e1d63 100644
--- a/example/hotspot_param_flow/qps_reject/hotspot_params_qps_reject_example.go
+++ b/example/hotspot_param_flow/qps_reject/hotspot_params_qps_reject_example.go
@@ -15,10 +15,14 @@
 package main
 
 import (
+	"fmt"
 	"log"
 	"math/rand"
 	"time"
 
+	"github.com/alibaba/sentinel-golang/core/base"
+	"github.com/alibaba/sentinel-golang/core/stat"
+
 	sentinel "github.com/alibaba/sentinel-golang/api"
 	"github.com/alibaba/sentinel-golang/core/config"
 	"github.com/alibaba/sentinel-golang/core/hotspot"
@@ -37,6 +41,8 @@ func main() {
 	if err != nil {
 		log.Fatal(err)
 	}
+	rand.Seed(time.Now().UnixNano())
+	testKey := "testKey"
 
 	_, err = hotspot.LoadRules([]*hotspot.Rule{
 		{
@@ -57,11 +63,36 @@ func main() {
 			BurstCount:      0,
 			DurationInSec:   1,
 		},
+		{
+			Resource:        "efg",
+			MetricType:      hotspot.QPS,
+			ControlBehavior: hotspot.Reject,
+			ParamKey:        testKey,
+			Threshold:       50,
+			BurstCount:      0,
+			DurationInSec:   1,
+		},
 	})
 	if err != nil {
 		log.Fatalf("Unexpected error: %+v", err)
 		return
 	}
+	for _, resource := range []string{"abc", "def", "efg"} {
+		go func(name string) {
+			node := stat.GetOrCreateResourceNode(name, base.ResTypeCommon)
+			for {
+				logging.Info("[HotSpot QPS] "+name,
+					"pass", node.GetQPS(base.MetricEventPass),
+					"block", node.GetQPS(base.MetricEventBlock),
+					"complete", node.GetQPS(base.MetricEventComplete),
+					"error", node.GetQPS(base.MetricEventError),
+					"rt", node.GetQPS(base.MetricEventRt),
+					//"\n total", node.GetQPS(base.MetricEventTotal),
+				)
+				time.Sleep(time.Duration(1000) * time.Millisecond)
+			}
+		}(resource)
+	}
 
 	logging.Info("[HotSpot Reject] Sentinel Go hot-spot param flow control demo is running. You may see the pass/block metric in the metric log.")
 	for i := 0; i < 10; i++ {
@@ -81,15 +112,33 @@ func main() {
 			}
 		}()
 	}
+	go func() {
+		for {
+			e, b := sentinel.Entry("def", sentinel.WithArgs(false, 9, "ahas", fooStruct{rand.Int63() % 5}))
+			if b != nil {
+				// Blocked. We could get the block reason from the BlockError.
+				time.Sleep(time.Duration(rand.Uint64()%10) * time.Millisecond)
+			} else {
+				// Passed, wrap the logic here.
+				time.Sleep(time.Duration(rand.Uint64()%10) * time.Millisecond)
+				// Be sure the entry is exited finally.
+				e.Exit()
+			}
+		}
+	}()
 
 	for {
-		e, b := sentinel.Entry("def", sentinel.WithArgs(false, 9, "ahas", fooStruct{rand.Int63() % 5}))
+		val := fmt.Sprintf("test%v", rand.Int31()%10)
+		e, b := sentinel.Entry("efg",
+			sentinel.WithAttachments(map[interface{}]interface{}{
+				testKey: val,
+			}))
 		if b != nil {
 			// Blocked. We could get the block reason from the BlockError.
 			time.Sleep(time.Duration(rand.Uint64()%10) * time.Millisecond)
 		} else {
 			// Passed, wrap the logic here.
-			time.Sleep(time.Duration(rand.Uint64()%10) * time.Millisecond)
+			time.Sleep(time.Duration(rand.Uint64()%2) * time.Millisecond)
 			// Be sure the entry is exited finally.
 			e.Exit()
 		}
@@ -97,4 +146,5 @@ func main() {
 
 	// The QPS of abc is about: 1500
 	// The QPS of def is about: 50
+	// The QPS of efg is about: 500
 }
diff --git a/go.mod b/go.mod
index 607860b..766695a 100644
--- a/go.mod
+++ b/go.mod
@@ -9,5 +9,7 @@ require (
 	github.com/shirou/gopsutil/v3 v3.20.11
 	github.com/stretchr/testify v1.6.1
 	go.uber.org/multierr v1.5.0
+	golang.org/x/sync v0.0.0-20190423024810-112230192c58
+	gopkg.in/errgo.v2 v2.1.0
 	gopkg.in/yaml.v2 v2.2.8
 )
diff --git a/go.sum b/go.sum
index 509236d..3f524fc 100644
--- a/go.sum
+++ b/go.sum
@@ -44,6 +44,7 @@ golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKG
 golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/sync v0.0.0-20190423024810-112230192c58 h1:8gQV6CLnAEikrhgkHFbMAEhagSSnXWGV915qUMm9mrU=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d h1:+R4KGOnez64A81RvjARKc4UT5/tI9ujCIVX+P5KiHuI=
@@ -60,6 +61,7 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/errgo.v2 v2.1.0 h1:0vLT13EuvQ0hNvakwLuFZ/jYrLp5F3kcWHXdRggjCE8=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
 gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=


Test patch:
diff --git a/core/hotspot/slot_test.go b/core/hotspot/slot_test.go
index 717d52a..cc5b338 100644
--- a/core/hotspot/slot_test.go
+++ b/core/hotspot/slot_test.go
@@ -15,11 +15,8 @@
 package hotspot
 
 import (
-	"reflect"
-	"testing"
-
 	"github.com/alibaba/sentinel-golang/core/base"
-	"github.com/stretchr/testify/assert"
+	//"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 )
 
@@ -52,53 +49,8 @@ func (m *TrafficShapingControllerMock) Replace(r *Rule) {
 	return
 }
 
-func Test_matchArg(t *testing.T) {
-	t.Run("Test_matchArg", func(t *testing.T) {
-		args := make([]interface{}, 10)
-		args[0] = 1
-		args[1] = 2
-		args[2] = 3
-		args[3] = 4
-		args[4] = 5
-		args[5] = 6
-		args[6] = 7
-		args[7] = 8
-		args[8] = 9
-		args[9] = 10
-
-		tcMock := &TrafficShapingControllerMock{}
-		tcMock.On("BoundParamIndex").Return(0)
-		ret0 := matchArg(tcMock, args)
-		assert.True(t, reflect.DeepEqual(ret0, 1))
-
-		tcMock1 := &TrafficShapingControllerMock{}
-		tcMock1.On("BoundParamIndex").Return(5)
-		ret1 := matchArg(tcMock1, args)
-		assert.True(t, reflect.DeepEqual(ret1, 6))
-
-		tcMock2 := &TrafficShapingControllerMock{}
-		tcMock2.On("BoundParamIndex").Return(9)
-		ret2 := matchArg(tcMock2, args)
-		assert.True(t, reflect.DeepEqual(ret2, 10))
-
-		tcMock3 := &TrafficShapingControllerMock{}
-		tcMock3.On("BoundParamIndex").Return(-1)
-		ret3 := matchArg(tcMock3, args)
-		assert.True(t, reflect.DeepEqual(ret3, 10))
-
-		tcMock4 := &TrafficShapingControllerMock{}
-		tcMock4.On("BoundParamIndex").Return(-10)
-		ret4 := matchArg(tcMock4, args)
-		assert.True(t, reflect.DeepEqual(ret4, 1))
-
-		tcMock5 := &TrafficShapingControllerMock{}
-		tcMock5.On("BoundParamIndex").Return(10)
-		ret5 := matchArg(tcMock5, args)
-		assert.True(t, ret5 == nil)
-
-		tcMock6 := &TrafficShapingControllerMock{}
-		tcMock6.On("BoundParamIndex").Return(-11)
-		ret6 := matchArg(tcMock6, args)
-		assert.True(t, ret6 == nil)
-	})
+func (m *TrafficShapingControllerMock) ExtractArgs(ctx *base.EntryContext) []interface{} {
+	_ = m.Called()
+	ret := []interface{}{ctx.Input.Args[m.BoundParamIndex()]}
+	return ret
 }
diff --git a/core/hotspot/traffic_shaping_test.go b/core/hotspot/traffic_shaping_test.go
index 8ab99fa..64b3636 100644
--- a/core/hotspot/traffic_shaping_test.go
+++ b/core/hotspot/traffic_shaping_test.go
@@ -15,10 +15,12 @@
 package hotspot
 
 import (
+	"reflect"
 	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/alibaba/sentinel-golang/core/base"
 	"github.com/alibaba/sentinel-golang/util"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -319,3 +321,53 @@ func Test_newBaseTrafficShapingController(t *testing.T) {
 		assert.True(t, tc.metric.RuleTokenCounter.Len() == ParamsMaxCapacity)
 	})
 }
+
+func Test_baseTrafficShapingController_ExtractArgs(t *testing.T) {
+	t.Run("Test_baseTrafficShapingController_ExtractArgs", func(t *testing.T) {
+
+		c := &baseTrafficShapingController{}
+
+		args := make([]interface{}, 10)
+		attachments := make(map[interface{}]interface{})
+		ctx := base.NewEmptyEntryContext()
+		ctx.Input = &base.SentinelInput{
+			BatchCount:  1,
+			Flag:        0,
+			Args:        args,
+			Attachments: attachments,
+		}
+		// no data
+		ret := c.ExtractArgs(ctx)
+		assert.True(t, reflect.DeepEqual(ret, []interface{}{}), ret)
+
+		// set data
+		args[0] = 1
+		args[1] = 2
+		attachments["test1"] = "v1"
+
+		// set index or key
+		// exist
+		c.paramIndex = 0
+		c.paramKey = "test1"
+		ret = c.ExtractArgs(ctx)
+		assert.True(t, reflect.DeepEqual(ret, []interface{}{1, "v1"}), ret)
+
+		// part exist 1
+		c.paramIndex = 10
+		c.paramKey = "test1"
+		ret = c.ExtractArgs(ctx)
+		assert.True(t, reflect.DeepEqual(ret, []interface{}{"v1"}), ret)
+
+		// part exist 2
+		c.paramIndex = 1
+		c.paramKey = "test2"
+		ret = c.ExtractArgs(ctx)
+		assert.True(t, reflect.DeepEqual(ret, []interface{}{2}), ret)
+
+		// not exist
+		c.paramIndex = 3
+		c.paramKey = "test2"
+		ret = c.ExtractArgs(ctx)
+		assert.True(t, reflect.DeepEqual(ret, []interface{}{}), ret)
+	})
+}