From 321db0288833648acfc2fd9d49690aace96ba534 Mon Sep 17 00:00:00 2001
From: ansiz <zhiyu.zxh@antfin.com>
Date: Mon, 20 Mar 2023 19:35:38 +0800
Subject: [PATCH 1/2] feat: add context.Context to SentinelInput (#519)

* feat:  add biz-defined error counter

* feat: add context.Context to sentinel entry context input

* fix: reset context in reset()

* chore: change field name
---
 api/api.go           | 13 +++++++++++++
 core/base/context.go |  9 ++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/api/api.go b/api/api.go
index ea3e445b8..124c01a3e 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"
@@ -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)
diff --git a/core/base/context.go b/core/base/context.go
index e78b98a4a..6ca279205 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
 	i.Args = i.Args[:0]
 	if len(i.Attachments) != 0 {
 		i.Attachments = make(map[interface{}]interface{})

From 930a63230ca6dd83eedb9749b667c3821ec2fcad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=AD=A2=E8=AF=AD?= <zhiyu.zxh@antfin.com>
Date: Mon, 20 Mar 2023 20:18:02 +0800
Subject: [PATCH 2/2] chore: add more test code

---
 api/api_test.go | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/api/api_test.go b/api/api_test.go
index cb38e2912..7dd48f404 100644
--- a/api/api_test.go
+++ b/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)
+}