feat: add context.Context to SentinelInput

api/api.go

@@ -15,6 +15,7 @@
 package api
 
 import (
+	"context"
 	"sync"
 
 	"github.com/alibaba/sentinel-golang/core/base"
@@ -36,6 +37,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
@@ -53,6 +55,7 @@ func (o *EntryOptions) Reset() {
 	o.slotChain = nil
 	o.args = o.args[:0]
 	o.attachments = nil
+	o.ctx = nil
 }
 
 type EntryOption func(*EntryOptions)
@@ -129,6 +132,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)
@@ -164,6 +174,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)

api/api_test.go

@@ -15,6 +15,7 @@
 package api
 
 import (
+	"context"
 	"testing"
 
 	"github.com/alibaba/sentinel-golang/core/base"
@@ -157,3 +158,11 @@ func Test_entryWithArgsAndChainBlock(t *testing.T) {
 	ssm.AssertNumberOfCalls(t, "OnEntryBlocked", 1)
 	ssm.AssertNumberOfCalls(t, "OnCompleted", 0)
 }
+
+func TestWithContext(t *testing.T) {
+	entry, _ := Entry("testing", WithContext(context.Background()))
+	assert.NotNil(t, entry.Context().Input.Context)
+
+	entry.Exit()
+	assert.Nil(t, entry.Context().Input.Context)
+}

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
 	i.Args = i.Args[:0]
 	if len(i.Attachments) != 0 {
 		i.Attachments = make(map[interface{}]interface{})