Skip to content

Commit f7991df

Browse files
cafracafra
committed
feat: hotspot rule add paramKey;
1 parent aafce3f commit f7991df

File tree

11 files changed

+274
-102
lines changed

11 files changed

+274
-102
lines changed

‎core/hotspot/concurrency_stat_slot.go‎

Lines changed: 2 additions & 4 deletions
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‎

Lines changed: 6 additions & 2 deletions
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 context 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+
// Mutually exclusive with 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:%v, 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
}

‎core/hotspot/rule_manager.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,6 +431,9 @@ func IsValidRule(rule *Rule) error {
431431
if rule.MetricType == QPS && rule.DurationInSec <= 0 {
432432
return errors.New("invalid duration")
433433
}
434+
if rule.ParamIndex > 0 && rule.ParamKey != "" {
435+
return errors.New("invalid param index and param key are mutually exclusive")
436+
}
434437
return checkControlBehaviorField(rule)
435438
}
436439

‎core/hotspot/rule_test.go‎

Lines changed: 2 additions & 1 deletion
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

‎core/hotspot/slot.go‎

Lines changed: 1 addition & 28 deletions
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‎

Lines changed: 4 additions & 53 deletions
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‎

Lines changed: 70 additions & 0 deletions
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.extractArgs(ctx)
169+
if value != nil {
170+
return
171+
}
172+
value = c.extractAttachmentArgs(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‎

Lines changed: 52 additions & 0 deletions
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,53 @@ 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+
attachments["test1"] = "v1"
349+
350+
// set index or key
351+
// exist
352+
c.paramIndex = 0
353+
c.paramKey = "test1"
354+
ret = c.ExtractArgs(ctx)
355+
assert.True(t, reflect.DeepEqual(ret, 1), ret)
356+
357+
// part exist 1
358+
c.paramIndex = 10
359+
c.paramKey = "test1"
360+
ret = c.ExtractArgs(ctx)
361+
assert.True(t, reflect.DeepEqual(ret, "v1"), ret)
362+
363+
// part exist 2
364+
c.paramIndex = 1
365+
c.paramKey = "test2"
366+
ret = c.ExtractArgs(ctx)
367+
assert.True(t, reflect.DeepEqual(ret, 2), ret)
368+
369+
// not exist
370+
c.paramIndex = 3
371+
c.paramKey = "test2"
372+
ret = c.ExtractArgs(ctx)
373+
assert.Nil(t, ret)
374+
})
375+
}

‎example/hotspot_param_flow/concurrency/hotspot_params_concurrency_example.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,15 @@ func main() {
3535
if err != nil {
3636
log.Fatal(err)
3737
}
38+
testKey := "testKey"
3839

3940
// the max concurrency is 8
4041
_, err = hotspot.LoadRules([]*hotspot.Rule{
4142
{
4243
Resource: "abc",
4344
MetricType: hotspot.Concurrency,
4445
ParamIndex: 0,
46+
ParamKey: testKey,
4547
Threshold: 8,
4648
DurationInSec: 0,
4749
},
@@ -79,7 +81,12 @@ func main() {
7981
}
8082

8183
for {
82-
e, b := sentinel.Entry("abc", sentinel.WithArgs(false, uint32(9)))
84+
e, b := sentinel.Entry("abc",
85+
sentinel.WithArgs(false, uint32(9),
86+
sentinel.WithAttachments(map[interface{}]interface{}{
87+
testKey: rand.Uint64() % 10,
88+
}),
89+
))
8390
if b != nil {
8491
// Blocked. We could get the block reason from the BlockError.
8592
time.Sleep(time.Duration(rand.Uint64()%10) * time.Millisecond)

0 commit comments

Comments
 (0)