add ctx parameter to the retryKey func

pkg/retry/policy_test.go

@@ -17,6 +17,7 @@
 package retry
 
 import (
+	"context"
 	"reflect"
 	"testing"
 
@@ -192,6 +193,7 @@ func TestRetryPolicyBothNotNil(t *testing.T) {
 		FailurePolicy: NewFailurePolicy(),
 		BackupPolicy:  NewBackupPolicy(20),
 	}
+	ctx := context.Background()
 	jsonRet, err := jsoni.MarshalToString(p)
 	test.Assert(t, err == nil, err)
 
@@ -205,7 +207,7 @@ func TestRetryPolicyBothNotNil(t *testing.T) {
 	rc := Container{}
 	rc.NotifyPolicyChange(ri.To().Method(), p2)
 
-	r := rc.getRetryer(ri)
+	r := rc.getRetryer(ctx, ri)
 	fr, ok := r.(*failureRetryer)
 	test.Assert(t, ok)
 	test.Assert(t, fr.enable)
@@ -231,7 +233,7 @@ func TestRetryPolicyBothNil(t *testing.T) {
 	rc := Container{}
 	rc.NotifyPolicyChange(ri.To().Method(), p2)
 
-	r := rc.getRetryer(ri)
+	r := rc.getRetryer(context.Background(), ri)
 	test.Assert(t, r != nil, r)
 }
 
@@ -252,7 +254,7 @@ func TestRetryPolicyFailure(t *testing.T) {
 	ri := genRPCInfo()
 	rc := Container{}
 	rc.NotifyPolicyChange(ri.To().Method(), p2)
-	r := rc.getRetryer(ri)
+	r := rc.getRetryer(context.Background(), ri)
 	fr, ok := r.(*failureRetryer)
 	test.Assert(t, ok)
 	test.Assert(t, fr.policy.Equals(p.FailurePolicy))
@@ -277,7 +279,7 @@ func TestRetryPolicyFailure(t *testing.T) {
 
 	// 更新配置
 	rc.NotifyPolicyChange(ri.To().Method(), p3)
-	r = rc.getRetryer(ri)
+	r = rc.getRetryer(context.Background(), ri)
 	fr, ok = r.(*failureRetryer)
 	test.Assert(t, ok)
 	test.Assert(t, fr.policy.Equals(p.FailurePolicy))

pkg/retry/retryer.go

@@ -35,7 +35,7 @@ type RPCCallFunc func(context.Context, Retryer) (rpcinfo rpcinfo.RPCInfo, resp i
 
 // GenRetryKeyFunc to generate retry key through rpcinfo.
 // You can customize the config key according to your config center.
-type GenRetryKeyFunc func(ri rpcinfo.RPCInfo) string
+type GenRetryKeyFunc func(ctx context.Context, ri rpcinfo.RPCInfo) string
 
 // Retryer is the interface for Retry implements
 type Retryer interface {
@@ -186,7 +186,7 @@ func NewRetryContainer(opts ...ContainerOption) *Container {
 	return rc
 }
 
-func defaultGenRetryKey(rpcInfo rpcinfo.RPCInfo) string {
+func defaultGenRetryKey(_ context.Context, rpcInfo rpcinfo.RPCInfo) string {
 	return rpcInfo.To().Method()
 }
 
@@ -339,7 +339,7 @@ func (rc *Container) WithRetryIfNeeded(ctx context.Context, callOptRetry *Policy
 			klog.Warnf("KITEX: new callopt retryer[%s] failed, err=%w", callOptRetry.Type, err)
 		}
 	} else {
-		retryer = rc.getRetryer(ri)
+		retryer = rc.getRetryer(ctx, ri)
 	}
 
 	// case 1(default, fast path): no retry policy
@@ -377,13 +377,13 @@ func NewRetryer(p Policy, r *ShouldResultRetry, cbC *cbContainer) (retryer Retry
 	return
 }
 
-func (rc *Container) getRetryer(ri rpcinfo.RPCInfo) Retryer {
+func (rc *Container) getRetryer(ctx context.Context, ri rpcinfo.RPCInfo) Retryer {
 	keyFunc := defaultGenRetryKey
 	if rc.genRetryKey != nil {
 		keyFunc = rc.genRetryKey
 	}
 	// the priority of specific method is high
-	r, ok := rc.retryerMap.Load(keyFunc(ri))
+	r, ok := rc.retryerMap.Load(keyFunc(ctx, ri))
 	if ok {
 		return r.(Retryer)
 	}

pkg/retry/retryer_test.go

@@ -48,7 +48,7 @@ func TestNewRetryContainer(t *testing.T) {
 		BackupPolicy:  NewBackupPolicy(10),
 		FailurePolicy: NewFailurePolicy(),
 	})
-	r := rc.getRetryer(genRPCInfo())
+	r := rc.getRetryer(context.Background(), genRPCInfo())
 	_, ok := r.(*failureRetryer)
 	test.Assert(t, ok)
 
@@ -67,7 +67,7 @@ func TestNewRetryContainer(t *testing.T) {
 		BackupPolicy:  NewBackupPolicy(10),
 		FailurePolicy: NewFailurePolicy(),
 	})
-	r = rc.getRetryer(genRPCInfo())
+	r = rc.getRetryer(context.Background(), genRPCInfo())
 	_, ok = r.(*failureRetryer)
 	test.Assert(t, ok)
 	_, allow = r.AllowRetry(context.Background())
@@ -78,7 +78,7 @@ func TestNewRetryContainer(t *testing.T) {
 		Type:         1,
 		BackupPolicy: NewBackupPolicy(20),
 	})
-	r = rc.getRetryer(genRPCInfo())
+	r = rc.getRetryer(context.Background(), genRPCInfo())
 	_, ok = r.(*backupRetryer)
 	test.Assert(t, ok)
 	_, allow = r.AllowRetry(context.Background())
@@ -89,7 +89,7 @@ func TestNewRetryContainer(t *testing.T) {
 		Type:         1,
 		BackupPolicy: NewBackupPolicy(20),
 	})
-	r = rc.getRetryer(genRPCInfo())
+	r = rc.getRetryer(context.Background(), genRPCInfo())
 	_, ok = r.(*backupRetryer)
 	test.Assert(t, ok)
 	_, allow = r.AllowRetry(context.Background())
@@ -271,7 +271,7 @@ func TestNewRetryContainer(t *testing.T) {
 	test.Assert(t, err != nil, err)
 
 	rc.DeletePolicy(method)
-	r = rc.getRetryer(genRPCInfo())
+	r = rc.getRetryer(context.Background(), genRPCInfo())
 	test.Assert(t, r == nil)
 }
 
@@ -946,14 +946,14 @@ func TestResultRetryWithPolicyChange(t *testing.T) {
 
 	// case 1: first time trigger NotifyPolicyChange, the `initRetryer` will be executed, check if the ShouldResultRetry is not nil
 	rc.NotifyPolicyChange(Wildcard, BuildFailurePolicy(NewFailurePolicy()))
-	r := rc.getRetryer(genRPCInfo())
+	r := rc.getRetryer(context.Background(), genRPCInfo())
 	fr, ok := r.(*failureRetryer)
 	test.Assert(t, ok)
 	test.Assert(t, fr.policy.ShouldResultRetry == shouldResultRetry)
 
 	// case 2: second time trigger NotifyPolicyChange, the `UpdatePolicy` will be executed, check if the ShouldResultRetry is not nil
 	rc.NotifyPolicyChange(Wildcard, BuildFailurePolicy(NewFailurePolicy()))
-	r = rc.getRetryer(genRPCInfo())
+	r = rc.getRetryer(context.Background(), genRPCInfo())
 	fr, ok = r.(*failureRetryer)
 	test.Assert(t, ok)
 	test.Assert(t, fr.policy.ShouldResultRetry == shouldResultRetry)
@@ -974,14 +974,14 @@ func TestResultRetryWithCtxWhenPolicyChange(t *testing.T) {
 
 	// case 1: first time trigger NotifyPolicyChange, the `initRetryer` will be executed, check if the ShouldResultRetry is not nil
 	rc.NotifyPolicyChange(Wildcard, BuildFailurePolicy(NewFailurePolicy()))
-	r := rc.getRetryer(genRPCInfo())
+	r := rc.getRetryer(context.Background(), genRPCInfo())
 	fr, ok := r.(*failureRetryer)
 	test.Assert(t, ok)
 	test.Assert(t, fr.policy.ShouldResultRetry == shouldResultRetry)
 
 	// case 2: second time trigger NotifyPolicyChange, the `UpdatePolicy` will be executed, check if the ShouldResultRetry is not nil
 	rc.NotifyPolicyChange(Wildcard, BuildFailurePolicy(NewFailurePolicy()))
-	r = rc.getRetryer(genRPCInfo())
+	r = rc.getRetryer(context.Background(), genRPCInfo())
 	fr, ok = r.(*failureRetryer)
 	test.Assert(t, ok)
 	test.Assert(t, fr.policy.ShouldResultRetry == shouldResultRetry)