Skip to content

Commit 6b43fd2

Browse files
cafralouyuting
cafra
and
louyuting
authored
hotspot module Rule support paramKey to match EntryContext attachment map (#376)
 Co-authored-by: louyuting <1849491904@qq.com>
1 parent aafce3f commit 6b43fd2

12 files changed

+282
-110
lines changed

‎core/hotspot/concurrency_stat_slot.go

+2-4
Original file line numberDiff line numberDiff line change
@@ -44,13 +44,12 @@ func (s *ConcurrencyStatSlot) Order() uint32 {
4444

4545
func (c *ConcurrencyStatSlot) OnEntryPassed(ctx *base.EntryContext) {
4646
res := ctx.Resource.Name()
47-
args := ctx.Input.Args
4847
tcs := getTrafficControllersFor(res)
4948
for _, tc := range tcs {
5049
if tc.BoundRule().MetricType != Concurrency {
5150
continue
5251
}
53-
arg := matchArg(tc, args)
52+
arg := tc.ExtractArgs(ctx)
5453
if arg == nil {
5554
continue
5655
}
@@ -72,13 +71,12 @@ func (c *ConcurrencyStatSlot) OnEntryBlocked(ctx *base.EntryContext, blockError
7271

7372
func (c *ConcurrencyStatSlot) OnCompleted(ctx *base.EntryContext) {
7473
res := ctx.Resource.Name()
75-
args := ctx.Input.Args
7674
tcs := getTrafficControllersFor(res)
7775
for _, tc := range tcs {
7876
if tc.BoundRule().MetricType != Concurrency {
7977
continue
8078
}
81-
arg := matchArg(tc, args)
79+
arg := tc.ExtractArgs(ctx)
8280
if arg == nil {
8381
continue
8482
}

‎core/hotspot/rule.go

+7-3
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,10 @@ type Rule struct {
8080
// if ParamIndex is great than or equals to zero, ParamIndex means the <ParamIndex>-th parameter
8181
// if ParamIndex is the negative, ParamIndex means the reversed <ParamIndex>-th parameter
8282
ParamIndex int `json:"paramIndex"`
83+
// ParamKey is the key in EntryContext.Input.Attachments map.
84+
// ParamKey can be used as a supplement to ParamIndex to facilitate rules to quickly obtain parameter from a large number of parameters
85+
// ParamKey is mutually exclusive with ParamIndex, ParamKey has the higher priority than ParamIndex
86+
ParamKey string `json:"paramKey"`
8387
// Threshold is the threshold to trigger rejection
8488
Threshold int64 `json:"threshold"`
8589
// MaxQueueingTimeMs only takes effect when ControlBehavior is Throttling and MetricType is QPS
@@ -100,8 +104,8 @@ func (r *Rule) String() string {
100104
b, err := json.Marshal(r)
101105
if err != nil {
102106
// Return the fallback string
103-
return fmt.Sprintf("{Id:%s, Resource:%s, MetricType:%+v, ControlBehavior:%+v, ParamIndex:%d, Threshold:%d, MaxQueueingTimeMs:%d, BurstCount:%d, DurationInSec:%d, ParamsMaxCapacity:%d, SpecificItems:%+v}",
104-
r.ID, r.Resource, r.MetricType, r.ControlBehavior, r.ParamIndex, r.Threshold, r.MaxQueueingTimeMs, r.BurstCount, r.DurationInSec, r.ParamsMaxCapacity, r.SpecificItems)
107+
return fmt.Sprintf("{Id:%s, Resource:%s, MetricType:%+v, ControlBehavior:%+v, ParamIndex:%d, ParamKey:%s, Threshold:%d, MaxQueueingTimeMs:%d, BurstCount:%d, DurationInSec:%d, ParamsMaxCapacity:%d, SpecificItems:%+v}",
108+
r.ID, r.Resource, r.MetricType, r.ControlBehavior, r.ParamIndex, r.ParamKey, r.Threshold, r.MaxQueueingTimeMs, r.BurstCount, r.DurationInSec, r.ParamsMaxCapacity, r.SpecificItems)
105109
}
106110
return string(b)
107111
}
@@ -116,7 +120,7 @@ func (r *Rule) IsStatReusable(newRule *Rule) bool {
116120

117121
// Equals checks whether current rule is consistent with the given rule.
118122
func (r *Rule) Equals(newRule *Rule) bool {
119-
baseCheck := r.Resource == newRule.Resource && r.MetricType == newRule.MetricType && r.ControlBehavior == newRule.ControlBehavior && r.ParamsMaxCapacity == newRule.ParamsMaxCapacity && r.ParamIndex == newRule.ParamIndex && r.Threshold == newRule.Threshold && r.DurationInSec == newRule.DurationInSec && reflect.DeepEqual(r.SpecificItems, newRule.SpecificItems)
123+
baseCheck := r.Resource == newRule.Resource && r.MetricType == newRule.MetricType && r.ControlBehavior == newRule.ControlBehavior && r.ParamsMaxCapacity == newRule.ParamsMaxCapacity && r.ParamIndex == newRule.ParamIndex && r.ParamKey == newRule.ParamKey && r.Threshold == newRule.Threshold && r.DurationInSec == newRule.DurationInSec && reflect.DeepEqual(r.SpecificItems, newRule.SpecificItems)
120124
if !baseCheck {
121125
return false
122126
}

‎core/hotspot/rule_manager.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -425,12 +425,12 @@ func IsValidRule(rule *Rule) error {
425425
if rule.ControlBehavior < 0 {
426426
return errors.New("invalid control strategy")
427427
}
428-
if rule.ParamIndex < 0 {
429-
return errors.New("invalid param index")
430-
}
431428
if rule.MetricType == QPS && rule.DurationInSec <= 0 {
432429
return errors.New("invalid duration")
433430
}
431+
if rule.ParamIndex > 0 && rule.ParamKey != "" {
432+
return errors.New("invalid param index and param key are mutually exclusive")
433+
}
434434
return checkControlBehaviorField(rule)
435435
}
436436

‎core/hotspot/rule_test.go

+4-1
Original file line numberDiff line numberDiff line change
@@ -39,14 +39,15 @@ func Test_Rule_String(t *testing.T) {
3939
MetricType: Concurrency,
4040
ControlBehavior: Reject,
4141
ParamIndex: 0,
42+
ParamKey: "key",
4243
Threshold: 110.0,
4344
MaxQueueingTimeMs: 5,
4445
BurstCount: 10,
4546
DurationInSec: 1,
4647
ParamsMaxCapacity: 10000,
4748
SpecificItems: specific,
4849
}
49-
assert.True(t, fmt.Sprintf("%+v", []*Rule{r}) == "[{Id:abc, Resource:abc, MetricType:Concurrency, ControlBehavior:Reject, ParamIndex:0, Threshold:110, MaxQueueingTimeMs:5, BurstCount:10, DurationInSec:1, ParamsMaxCapacity:10000, SpecificItems:map[1123:3 sss:1]}]")
50+
assert.True(t, fmt.Sprintf("%+v", []*Rule{r}) == "[{Id:abc, Resource:abc, MetricType:Concurrency, ControlBehavior:Reject, ParamIndex:0, ParamKey:key, Threshold:110, MaxQueueingTimeMs:5, BurstCount:10, DurationInSec:1, ParamsMaxCapacity:10000, SpecificItems:map[1123:3 sss:1]}]")
5051
})
5152
}
5253

@@ -61,6 +62,7 @@ func Test_Rule_Equals(t *testing.T) {
6162
MetricType: Concurrency,
6263
ControlBehavior: Reject,
6364
ParamIndex: 0,
65+
ParamKey: "testKey",
6466
Threshold: 110.0,
6567
MaxQueueingTimeMs: 5,
6668
BurstCount: 10,
@@ -78,6 +80,7 @@ func Test_Rule_Equals(t *testing.T) {
7880
MetricType: Concurrency,
7981
ControlBehavior: Reject,
8082
ParamIndex: 0,
83+
ParamKey: "testKey",
8184
Threshold: 110.0,
8285
MaxQueueingTimeMs: 5,
8386
BurstCount: 10,

‎core/hotspot/slot.go

+1-28
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ package hotspot
1616

1717
import (
1818
"github.com/alibaba/sentinel-golang/core/base"
19-
"github.com/alibaba/sentinel-golang/logging"
2019
"github.com/alibaba/sentinel-golang/util"
2120
)
2221

@@ -40,40 +39,14 @@ func (s *Slot) Order() uint32 {
4039
return RuleCheckSlotOrder
4140
}
4241

43-
// matchArg matches the arg from args based on TrafficShapingController
44-
// return nil if match failed.
45-
func matchArg(tc TrafficShapingController, args []interface{}) interface{} {
46-
if tc == nil {
47-
return nil
48-
}
49-
idx := tc.BoundParamIndex()
50-
if idx < 0 {
51-
idx = len(args) + idx
52-
}
53-
if idx < 0 {
54-
if logging.DebugEnabled() {
55-
logging.Debug("[Slot matchArg] The param index of hotspot traffic shaping controller is invalid", "args", args, "paramIndex", tc.BoundParamIndex())
56-
}
57-
return nil
58-
}
59-
if idx >= len(args) {
60-
if logging.DebugEnabled() {
61-
logging.Debug("[Slot matchArg] The argument in index doesn't exist", "args", args, "paramIndex", tc.BoundParamIndex())
62-
}
63-
return nil
64-
}
65-
return args[idx]
66-
}
67-
6842
func (s *Slot) Check(ctx *base.EntryContext) *base.TokenResult {
6943
res := ctx.Resource.Name()
70-
args := ctx.Input.Args
7144
batch := int64(ctx.Input.BatchCount)
7245

7346
result := ctx.RuleCheckResult
7447
tcs := getTrafficControllersFor(res)
7548
for _, tc := range tcs {
76-
arg := matchArg(tc, args)
49+
arg := tc.ExtractArgs(ctx)
7750
if arg == nil {
7851
continue
7952
}

‎core/hotspot/slot_test.go

+4-53
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,7 @@
1515
package hotspot
1616

1717
import (
18-
"reflect"
19-
"testing"
20-
2118
"github.com/alibaba/sentinel-golang/core/base"
22-
"github.com/stretchr/testify/assert"
2319
"github.com/stretchr/testify/mock"
2420
)
2521

@@ -52,53 +48,8 @@ func (m *TrafficShapingControllerMock) Replace(r *Rule) {
5248
return
5349
}
5450

55-
func Test_matchArg(t *testing.T) {
56-
t.Run("Test_matchArg", func(t *testing.T) {
57-
args := make([]interface{}, 10)
58-
args[0] = 1
59-
args[1] = 2
60-
args[2] = 3
61-
args[3] = 4
62-
args[4] = 5
63-
args[5] = 6
64-
args[6] = 7
65-
args[7] = 8
66-
args[8] = 9
67-
args[9] = 10
68-
69-
tcMock := &TrafficShapingControllerMock{}
70-
tcMock.On("BoundParamIndex").Return(0)
71-
ret0 := matchArg(tcMock, args)
72-
assert.True(t, reflect.DeepEqual(ret0, 1))
73-
74-
tcMock1 := &TrafficShapingControllerMock{}
75-
tcMock1.On("BoundParamIndex").Return(5)
76-
ret1 := matchArg(tcMock1, args)
77-
assert.True(t, reflect.DeepEqual(ret1, 6))
78-
79-
tcMock2 := &TrafficShapingControllerMock{}
80-
tcMock2.On("BoundParamIndex").Return(9)
81-
ret2 := matchArg(tcMock2, args)
82-
assert.True(t, reflect.DeepEqual(ret2, 10))
83-
84-
tcMock3 := &TrafficShapingControllerMock{}
85-
tcMock3.On("BoundParamIndex").Return(-1)
86-
ret3 := matchArg(tcMock3, args)
87-
assert.True(t, reflect.DeepEqual(ret3, 10))
88-
89-
tcMock4 := &TrafficShapingControllerMock{}
90-
tcMock4.On("BoundParamIndex").Return(-10)
91-
ret4 := matchArg(tcMock4, args)
92-
assert.True(t, reflect.DeepEqual(ret4, 1))
93-
94-
tcMock5 := &TrafficShapingControllerMock{}
95-
tcMock5.On("BoundParamIndex").Return(10)
96-
ret5 := matchArg(tcMock5, args)
97-
assert.True(t, ret5 == nil)
98-
99-
tcMock6 := &TrafficShapingControllerMock{}
100-
tcMock6.On("BoundParamIndex").Return(-11)
101-
ret6 := matchArg(tcMock6, args)
102-
assert.True(t, ret6 == nil)
103-
})
51+
func (m *TrafficShapingControllerMock) ExtractArgs(ctx *base.EntryContext) []interface{} {
52+
_ = m.Called()
53+
ret := []interface{}{ctx.Input.Args[m.BoundParamIndex()]}
54+
return ret
10455
}

‎core/hotspot/traffic_shaping.go

+70
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@ type TrafficShapingController interface {
3333

3434
BoundParamIndex() int
3535

36+
ExtractArgs(ctx *base.EntryContext) interface{}
37+
3638
BoundMetric() *ParamsMetric
3739

3840
BoundRule() *Rule
@@ -44,6 +46,7 @@ type baseTrafficShapingController struct {
4446
res string
4547
metricType MetricType
4648
paramIndex int
49+
paramKey string
4750
threshold int64
4851
specificItems map[interface{}]int64
4952
durationInSec int64
@@ -60,6 +63,7 @@ func newBaseTrafficShapingControllerWithMetric(r *Rule, metric *ParamsMetric) *b
6063
res: r.Resource,
6164
metricType: r.MetricType,
6265
paramIndex: r.ParamIndex,
66+
paramKey: r.ParamKey,
6367
threshold: r.Threshold,
6468
specificItems: r.SpecificItems,
6569
durationInSec: r.DurationInSec,
@@ -155,6 +159,72 @@ func (c *baseTrafficShapingController) BoundParamIndex() int {
155159
return c.paramIndex
156160
}
157161

162+
// ExtractArgs matches the arg from ctx based on TrafficShapingController
163+
// return nil if match failed.
164+
func (c *baseTrafficShapingController) ExtractArgs(ctx *base.EntryContext) (value interface{}) {
165+
if c == nil {
166+
return nil
167+
}
168+
value = c.extractAttachmentArgs(ctx)
169+
if value != nil {
170+
return
171+
}
172+
value = c.extractArgs(ctx)
173+
if value != nil {
174+
return
175+
}
176+
return
177+
}
178+
func (c *baseTrafficShapingController) extractArgs(ctx *base.EntryContext) interface{} {
179+
args := ctx.Input.Args
180+
idx := c.BoundParamIndex()
181+
if idx < 0 {
182+
idx = len(args) + idx
183+
}
184+
if idx < 0 {
185+
if logging.DebugEnabled() {
186+
logging.Debug("[extractArgs] The param index of hotspot traffic shaping controller is invalid",
187+
"args", args, "paramIndex", c.BoundParamIndex())
188+
}
189+
return nil
190+
}
191+
if idx >= len(args) {
192+
if logging.DebugEnabled() {
193+
logging.Debug("[extractArgs] The argument in index doesn't exist",
194+
"args", args, "paramIndex", c.BoundParamIndex())
195+
}
196+
return nil
197+
}
198+
return args[idx]
199+
}
200+
func (c *baseTrafficShapingController) extractAttachmentArgs(ctx *base.EntryContext) interface{} {
201+
attachments := ctx.Input.Attachments
202+
203+
if attachments == nil {
204+
if logging.DebugEnabled() {
205+
logging.Debug("[paramKey] The attachments of ctx is nil",
206+
"args", attachments, "paramKey", c.paramKey)
207+
}
208+
return nil
209+
}
210+
if c.paramKey == "" {
211+
if logging.DebugEnabled() {
212+
logging.Debug("[paramKey] The param key is nil",
213+
"args", attachments, "paramKey", c.paramKey)
214+
}
215+
return nil
216+
}
217+
arg, ok := attachments[c.paramKey]
218+
if !ok {
219+
if logging.DebugEnabled() {
220+
logging.Debug("[paramKey] extracted data does not exist",
221+
"args", attachments, "paramKey", c.paramKey)
222+
}
223+
}
224+
225+
return arg
226+
}
227+
158228
func (c *rejectTrafficShapingController) PerformChecking(arg interface{}, batchCount int64) *base.TokenResult {
159229
metric := c.metric
160230
if metric == nil {

‎core/hotspot/traffic_shaping_test.go

+53
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,12 @@
1515
package hotspot
1616

1717
import (
18+
"reflect"
1819
"sync/atomic"
1920
"testing"
2021
"time"
2122

23+
"github.com/alibaba/sentinel-golang/core/base"
2224
"github.com/alibaba/sentinel-golang/util"
2325
"github.com/stretchr/testify/assert"
2426
"github.com/stretchr/testify/mock"
@@ -321,3 +323,54 @@ func Test_newBaseTrafficShapingController(t *testing.T) {
321323
assert.True(t, tc.metric.RuleTokenCounter.Len() == ParamsMaxCapacity)
322324
})
323325
}
326+
327+
func Test_baseTrafficShapingController_ExtractArgs(t *testing.T) {
328+
t.Run("Test_baseTrafficShapingController_ExtractArgs", func(t *testing.T) {
329+
330+
c := &baseTrafficShapingController{}
331+
332+
args := make([]interface{}, 10)
333+
attachments := make(map[interface{}]interface{})
334+
ctx := base.NewEmptyEntryContext()
335+
ctx.Input = &base.SentinelInput{
336+
BatchCount: 1,
337+
Flag: 0,
338+
Args: args,
339+
Attachments: attachments,
340+
}
341+
// no data
342+
ret := c.ExtractArgs(ctx)
343+
assert.Nil(t, ret)
344+
345+
// set data
346+
args[0] = 1
347+
args[1] = 2
348+
value1 := "v1"
349+
attachments["test1"] = value1
350+
351+
// set index or key
352+
// exist
353+
c.paramIndex = 0
354+
c.paramKey = "test1"
355+
ret = c.ExtractArgs(ctx)
356+
assert.True(t, reflect.DeepEqual(ret, value1), ret)
357+
358+
// part exist 1
359+
c.paramIndex = 10
360+
c.paramKey = "test1"
361+
ret = c.ExtractArgs(ctx)
362+
assert.True(t, reflect.DeepEqual(ret, value1), ret)
363+
364+
// part exist 2
365+
c.paramIndex = 1
366+
c.paramKey = "test2"
367+
ret = c.ExtractArgs(ctx)
368+
assert.True(t, reflect.DeepEqual(ret, 2), ret)
369+
370+
// not exist
371+
c.paramIndex = 3
372+
c.paramKey = "test2"
373+
ret = c.ExtractArgs(ctx)
374+
assert.Nil(t, ret)
375+
})
376+
}

0 commit comments

Comments
 (0)