Repository: alibaba/sentinel-golang
Pull request #507 (2023-02-03 11:15:59)

Sentinel rule add block type hotspot set first monitor node
<!--  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
sentinel rule interface add blockType.
Because this is an attribute of the rule, it is convenient for the state slot to obtain the rule type of the current limit through the rule.

### 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
by UT
### Special notes for reviews


Patch:
diff --git a/api/api.go b/api/api.go
index 3e97032..19e04de 100644
--- a/api/api.go
+++ b/api/api.go
@@ -15,6 +15,7 @@
 package api
 
 import (
+	"context"
 	"sync"
 
 	"github.com/alibaba/sentinel-golang/core/base"
@@ -37,6 +38,7 @@ var entryOptsPool = sync.Pool{
 
 // EntryOptions represents the options of a Sentinel resource entry.
 type EntryOptions struct {
+	ctx          context.Context
 	resourceType base.ResourceType
 	entryType    base.TrafficType
 	batchCount   uint32
@@ -54,6 +56,7 @@ func (o *EntryOptions) Reset() {
 	o.slotChain = nil
 	o.args = nil
 	o.attachments = nil
+	o.ctx = nil
 }
 
 type EntryOption func(*EntryOptions)
@@ -130,6 +133,13 @@ func WithAttachments(data map[interface{}]interface{}) EntryOption {
 	}
 }
 
+// WithContext sets the resource entry with the given context.
+func WithContext(ctx context.Context) EntryOption {
+	return func(opts *EntryOptions) {
+		opts.ctx = ctx
+	}
+}
+
 // Entry is the basic API of Sentinel.
 func Entry(resource string, opts ...EntryOption) (*base.SentinelEntry, *base.BlockError) {
 	options := entryOptsPool.Get().(*EntryOptions)
@@ -168,6 +178,9 @@ func entry(resource string, options *EntryOptions) (*base.SentinelEntry, *base.B
 	if len(options.attachments) != 0 {
 		ctx.Input.Attachments = options.attachments
 	}
+	if options.ctx != nil {
+		ctx.Input.Context = options.ctx
+	}
 	e := base.NewSentinelEntry(ctx, rw, sc)
 	ctx.SetEntry(e)
 	r := sc.Entry(ctx)
diff --git a/core/base/context.go b/core/base/context.go
index 8635ddd..e052dde 100644
--- a/core/base/context.go
+++ b/core/base/context.go
@@ -14,7 +14,11 @@
 
 package base
 
-import "github.com/alibaba/sentinel-golang/util"
+import (
+	"context"
+
+	"github.com/alibaba/sentinel-golang/util"
+)
 
 type EntryContext struct {
 	entry *SentinelEntry
@@ -86,11 +90,14 @@ type SentinelInput struct {
 	Args       []interface{}
 	// store some values in this context when calling context in slot.
 	Attachments map[interface{}]interface{}
+	// store the request context
+	Context context.Context
 }
 
 func (i *SentinelInput) reset() {
 	i.BatchCount = 1
 	i.Flag = 0
+	i.Context = nil
 	if len(i.Args) != 0 {
 		i.Args = make([]interface{}, 0)
 	}
diff --git a/core/base/mock_rule.go b/core/base/mock_rule.go
new file mode 100644
index 0000000..556903c
--- /dev/null
+++ b/core/base/mock_rule.go
@@ -0,0 +1,21 @@
+package base
+
+type MockRule struct {
+	Id string `json:"id"`
+}
+
+func (m *MockRule) BlockType() BlockType {
+	return BlockTypeFlow
+}
+
+func (m *MockRule) String() string {
+	return "mock rule"
+}
+
+func (m *MockRule) ResourceName() string {
+	return "mock resource"
+}
+
+func (m *MockRule) RuleID() string {
+	return m.Id
+}
diff --git a/core/base/rule.go b/core/base/rule.go
index 981f8a9..46e74b6 100644
--- a/core/base/rule.go
+++ b/core/base/rule.go
@@ -22,4 +22,6 @@ type SentinelRule interface {
 	ResourceName() string
 
 	RuleID() string
+
+	BlockType() BlockType
 }
diff --git a/core/circuitbreaker/rule.go b/core/circuitbreaker/rule.go
index d8f7f80..68005b6 100644
--- a/core/circuitbreaker/rule.go
+++ b/core/circuitbreaker/rule.go
@@ -17,6 +17,8 @@ package circuitbreaker
 import (
 	"fmt"
 
+	"github.com/alibaba/sentinel-golang/core/base"
+
 	"github.com/alibaba/sentinel-golang/util"
 )
 
@@ -80,6 +82,10 @@ type Rule struct {
 	Threshold float64 `json:"threshold"`
 }
 
+func (r *Rule) BlockType() base.BlockType {
+	return base.BlockTypeCircuitBreaking
+}
+
 func (r *Rule) String() string {
 	// fallback string
 	return fmt.Sprintf("{id=%s, resource=%s, strategy=%s, RetryTimeoutMs=%d, MinRequestAmount=%d, StatIntervalMs=%d, StatSlidingWindowBucketCount=%d, MaxAllowedRtMs=%d, Threshold=%f}",
diff --git a/core/flow/rule.go b/core/flow/rule.go
index d627da8..4769dac 100644
--- a/core/flow/rule.go
+++ b/core/flow/rule.go
@@ -18,6 +18,8 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/alibaba/sentinel-golang/core/base"
+
 	"github.com/alibaba/sentinel-golang/util"
 )
 
@@ -149,3 +151,7 @@ func (r *Rule) ResourceName() string {
 func (r *Rule) RuleID() string {
 	return r.ID
 }
+
+func (r *Rule) BlockType() base.BlockType {
+	return base.BlockTypeFlow
+}
diff --git a/core/hotspot/rule.go b/core/hotspot/rule.go
index d39e09d..5c48e89 100644
--- a/core/hotspot/rule.go
+++ b/core/hotspot/rule.go
@@ -19,6 +19,8 @@ import (
 	"fmt"
 	"reflect"
 	"strconv"
+
+	"github.com/alibaba/sentinel-golang/core/base"
 )
 
 // ControlBehavior indicates the traffic shaping behaviour.
@@ -129,6 +131,10 @@ func (r *Rule) RuleID() string {
 	return r.ID
 }
 
+func (r *Rule) BlockType() base.BlockType {
+	return base.BlockTypeHotSpotParamFlow
+}
+
 // IsStatReusable checks whether current rule is "statistically" equal to the given rule.
 func (r *Rule) IsStatReusable(newRule *Rule) bool {
 	return r.Resource == newRule.Resource && r.ControlBehavior == newRule.ControlBehavior && r.ParamsMaxCapacity == newRule.ParamsMaxCapacity && r.DurationInSec == newRule.DurationInSec && r.MetricType == newRule.MetricType
diff --git a/core/hotspot/slot.go b/core/hotspot/slot.go
index 653541f..77e4215 100644
--- a/core/hotspot/slot.go
+++ b/core/hotspot/slot.go
@@ -25,11 +25,6 @@ import (
 const (
 	RuleCheckSlotName  = "sentinel-core-hotspot-rule-check-slot"
 	RuleCheckSlotOrder = 4000
-
-	keyMonitorBlockNodes = "monitorBlockNodes"
-	keyControlBlockNode  = "controlBlockNode"
-	keyIsMonitorBlocked  = "isMonitorBlocked"
-	keyChildNodes        = "childNodes"
 )
 
 var (
@@ -51,7 +46,7 @@ func (s *Slot) Check(ctx *base.EntryContext) *base.TokenResult {
 	res := ctx.Resource.Name()
 	batch := int64(ctx.Input.BatchCount)
 	result := ctx.RuleCheckResult
-	nodes := getAllNodes(ctx)
+	nodes := stat.GetAllNodes(ctx)
 	for _, node := range nodes {
 		tcs := getTrafficControllersFor(node.ResourceName())
 		for _, tc := range tcs {
@@ -69,10 +64,10 @@ func (s *Slot) Check(ctx *base.EntryContext) *base.TokenResult {
 
 				if r.Status() == base.ResultStatusBlocked {
 					if tc.BoundRule().Mode == MONITOR {
-						appendMonitorBlockNode(ctx, node)
+						stat.SetFirstMonitorNodeAndRule(ctx, node, tc.BoundRule())
 						continue
 					}
-					setBlockNode(ctx, node)
+					stat.SetBlockNode(ctx, node)
 					r.ResetToBlockedWith(
 						base.WithBlockResource(res),
 						base.WithBlockType(base.BlockTypeHotSpotParamFlow),
@@ -101,49 +96,3 @@ func canPassCheck(tc TrafficShapingController, arg interface{}, batch int64) *ba
 func canPassLocalCheck(tc TrafficShapingController, arg interface{}, batch int64) *base.TokenResult {
 	return tc.PerformChecking(arg, batch)
 }
-
-func PutOutputAttachment(ctx *base.EntryContext, key interface{}, value interface{}) {
-	if ctx.Data == nil {
-		ctx.Data = make(map[interface{}]interface{})
-	}
-	ctx.Data[key] = value
-}
-
-func getOutputAttachment(ctx *base.EntryContext, key interface{}) interface{} {
-	if ctx.Data == nil {
-		return nil
-	}
-	return ctx.Data[key]
-}
-
-// getAllNodes gets all child node and parent node list.
-func getAllNodes(ctx *base.EntryContext) []*stat.ResourceNode {
-	childNodes, childOk := getOutputAttachment(ctx, keyChildNodes).([]*stat.ResourceNode)
-	var allNodes []*stat.ResourceNode
-	if childOk {
-		allNodes = append(allNodes, childNodes...)
-	}
-	parentResNode, ok := ctx.StatNode.(*stat.ResourceNode)
-	if ok {
-		allNodes = append(allNodes, parentResNode)
-	}
-	return allNodes
-}
-
-func appendMonitorBlockNode(ctx *base.EntryContext, node *stat.ResourceNode) {
-	if ctx == nil || node == nil {
-		return
-	}
-	nodes, ok := getOutputAttachment(ctx, keyMonitorBlockNodes).([]*stat.ResourceNode)
-	if ok && nodes != nil {
-		nodes = append(nodes, node)
-	} else {
-		nodes = []*stat.ResourceNode{node}
-	}
-	PutOutputAttachment(ctx, keyIsMonitorBlocked, true)
-	PutOutputAttachment(ctx, keyMonitorBlockNodes, nodes)
-}
-
-func setBlockNode(ctx *base.EntryContext, node *stat.ResourceNode) {
-	PutOutputAttachment(ctx, keyControlBlockNode, node)
-}
diff --git a/core/isolation/rule.go b/core/isolation/rule.go
index d363826..c713d2d 100644
--- a/core/isolation/rule.go
+++ b/core/isolation/rule.go
@@ -17,6 +17,8 @@ package isolation
 import (
 	"encoding/json"
 	"fmt"
+
+	"github.com/alibaba/sentinel-golang/core/base"
 )
 
 // MetricType represents the target metric type.
@@ -61,3 +63,7 @@ func (r *Rule) ResourceName() string {
 func (r *Rule) RuleID() string {
 	return r.ID
 }
+
+func (r *Rule) BlockType() base.BlockType {
+	return base.BlockTypeIsolation
+}
diff --git a/core/stat/common.go b/core/stat/common.go
new file mode 100644
index 0000000..a53b541
--- /dev/null
+++ b/core/stat/common.go
@@ -0,0 +1,78 @@
+package stat
+
+import "github.com/alibaba/sentinel-golang/core/base"
+
+const (
+	KeyMonitorBlockNode = "monitorBlockNode"
+	KeyControlBlockNode = "controlBlockNode"
+	KeyIsMonitorBlocked = "isMonitorBlocked"
+	KeyMonitorBlockRule = "monitorBlockRule"
+	KeyChildNodes       = "childNodes"
+)
+
+// SetFirstMonitorNodeAndRule set monitor block node. only record the first.
+func SetFirstMonitorNodeAndRule(ctx *base.EntryContext, node *ResourceNode, rule base.SentinelRule) {
+	if ctx == nil || node == nil {
+		return
+	}
+	if IsMonitorBlocked(ctx) {
+		return
+	}
+
+	PutOutputAttachment(ctx, KeyIsMonitorBlocked, true)
+	PutOutputAttachment(ctx, KeyMonitorBlockNode, node)
+	PutOutputAttachment(ctx, KeyMonitorBlockRule, rule)
+}
+
+func GetMonitorBlockInfo(ctx *base.EntryContext) (node *ResourceNode, rule base.SentinelRule, isMonitorBlocked bool) {
+	if !IsMonitorBlocked(ctx) {
+		return nil, nil, false
+	}
+	node, _ = GetOutputAttachment(ctx, KeyMonitorBlockNode).(*ResourceNode)
+	rule, _ = GetOutputAttachment(ctx, KeyMonitorBlockRule).(base.SentinelRule)
+	return node, rule, true
+}
+
+// IsMonitorBlocked reports whether current request was blocked in monitor mode.
+func IsMonitorBlocked(ctx *base.EntryContext) bool {
+	if ctx == nil {
+		return false
+	}
+	isMonitorBlocked, ok := GetOutputAttachment(ctx, KeyIsMonitorBlocked).(bool)
+	if !ok {
+		return false
+	}
+	return isMonitorBlocked
+}
+
+func GetOutputAttachment(ctx *base.EntryContext, key interface{}) interface{} {
+	if ctx.Data == nil {
+		return nil
+	}
+	return ctx.Data[key]
+}
+
+func PutOutputAttachment(ctx *base.EntryContext, key interface{}, value interface{}) {
+	if ctx.Data == nil {
+		ctx.Data = make(map[interface{}]interface{})
+	}
+	ctx.Data[key] = value
+}
+
+// GetAllNodes gets all child node and parent node list.
+func GetAllNodes(ctx *base.EntryContext) []*ResourceNode {
+	childNodes, childOk := GetOutputAttachment(ctx, KeyChildNodes).([]*ResourceNode)
+	var allNodes []*ResourceNode
+	if childOk {
+		allNodes = append(allNodes, childNodes...)
+	}
+	parentResNode, ok := ctx.StatNode.(*ResourceNode)
+	if ok {
+		allNodes = append(allNodes, parentResNode)
+	}
+	return allNodes
+}
+
+func SetBlockNode(ctx *base.EntryContext, node *ResourceNode) {
+	PutOutputAttachment(ctx, KeyControlBlockNode, node)
+}
diff --git a/core/system/rule.go b/core/system/rule.go
index 5657bc5..9947249 100644
--- a/core/system/rule.go
+++ b/core/system/rule.go
@@ -17,6 +17,8 @@ package system
 import (
 	"encoding/json"
 	"fmt"
+
+	"github.com/alibaba/sentinel-golang/core/base"
 )
 
 type MetricType uint32
@@ -94,3 +96,7 @@ func (r *Rule) ResourceName() string {
 func (r *Rule) RuleID() string {
 	return r.ID
 }
+
+func (r *Rule) BlockType() base.BlockType {
+	return base.BlockTypeSystemFlow
+}
diff --git a/go.mod b/go.mod
index 766695a..3547a5a 100644
--- a/go.mod
+++ b/go.mod
@@ -10,6 +10,5 @@ require (
 	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
 )


Test patch:
diff --git a/core/base/block_error_test.go b/core/base/block_error_test.go
index d22d3af..107e95f 100644
--- a/core/base/block_error_test.go
+++ b/core/base/block_error_test.go
@@ -36,19 +36,3 @@ func TestNewBlockError(t *testing.T) {
 		})
 	}
 }
-
-type MockRule struct {
-	Id string `json:"id"`
-}
-
-func (m *MockRule) String() string {
-	return "mock rule"
-}
-
-func (m *MockRule) ResourceName() string {
-	return "mock resource"
-}
-
-func (m *MockRule) RuleID() string {
-	return m.Id
-}
diff --git a/core/stat/common_test.go b/core/stat/common_test.go
new file mode 100644
index 0000000..e907ec9
--- /dev/null
+++ b/core/stat/common_test.go
@@ -0,0 +1,71 @@
+package stat
+
+import (
+	"testing"
+
+	"github.com/alibaba/sentinel-golang/core/base"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetAllNodes(t *testing.T) {
+	ctx := base.NewEmptyEntryContext()
+	ctx.Input = &base.SentinelInput{}
+	ctx.Resource = &base.ResourceWrapper{}
+	data := GetAllNodes(ctx)
+	assert.Nil(t, data)
+}
+
+func TestIsMonitorBlocked(t *testing.T) {
+	assert.False(t, IsMonitorBlocked(nil))
+
+	ctx := base.NewEmptyEntryContext()
+	ctx.Input = &base.SentinelInput{}
+	ctx.Data = make(map[interface{}]interface{})
+	assert.False(t, IsMonitorBlocked(ctx))
+
+	PutOutputAttachment(ctx, KeyIsMonitorBlocked, true)
+	assert.True(t, IsMonitorBlocked(ctx))
+
+	PutOutputAttachment(ctx, KeyIsMonitorBlocked, false)
+	assert.False(t, IsMonitorBlocked(ctx))
+
+	PutOutputAttachment(ctx, KeyIsMonitorBlocked, 1)
+	assert.False(t, IsMonitorBlocked(ctx))
+}
+
+func mockCtx() *base.EntryContext {
+	ctx := base.NewEmptyEntryContext()
+	ctx.Input = &base.SentinelInput{}
+	ctx.Data = make(map[interface{}]interface{})
+	return ctx
+}
+
+func TestSetFirstMonitorNode(t *testing.T) {
+	//nil
+	SetFirstMonitorNodeAndRule(nil, nil, nil)
+	_, _, isBlocked := GetMonitorBlockInfo(nil)
+	assert.False(t, isBlocked)
+
+	// first set
+	ctx := mockCtx()
+	resourceNode := NewResourceNode("test", base.ResTypeRPC)
+	rule := &base.MockRule{Id: "1"}
+	SetFirstMonitorNodeAndRule(ctx, resourceNode, rule)
+	gotNode, gotRule, isBlocked2 := GetMonitorBlockInfo(ctx)
+
+	assert.True(t, isBlocked2)
+	assert.Equal(t, gotNode, resourceNode)
+	assert.Equal(t, gotRule, rule)
+
+	// second set
+
+	resourceNode2 := NewResourceNode("test2", base.ResTypeRPC)
+	rule2 := &base.MockRule{Id: "2"}
+	SetFirstMonitorNodeAndRule(ctx, resourceNode2, rule2)
+	gotNode2, gotRule2, isBlock2 := GetMonitorBlockInfo(ctx)
+
+	assert.True(t, isBlock2)
+	assert.Equal(t, gotNode2, resourceNode)
+	assert.Equal(t, gotRule2, rule)
+}