fix: retry container cb test (#41)

Author: Felix021

thriftrpc/retrycall/retrycall.go

@@ -42,8 +42,13 @@ func (h *STServiceHandler) TestSTReq(ctx context.Context, req *stability.STReque
 		Mp:      req.StringMap,
 		FlagMsg: req.FlagMsg,
 	}
-	if v := atomic.AddInt32(&testSTReqCount, 1); v%10 == 0 {
-		time.Sleep(50 * time.Millisecond)
+	if !skipCounterSleep(ctx) { // better not rely on this since it's a global variable
+		if v := atomic.AddInt32(&testSTReqCount, 1); v%10 == 0 {
+			time.Sleep(50 * time.Millisecond)
+		}
+	}
+	if sleepTime := getSleepTimeMS(ctx); sleepTime > 0 {
+		time.Sleep(sleepTime)
 	}
 	return resp, nil
 }

thriftrpc/retrycall/retrycall_mockresult.go

@@ -27,10 +27,39 @@ import (
 )
 
 const (
-	retryMsg       = "retry"
-	sleepTimeMsKey = "TIME_SLEEP_TIME_MS"
+	retryMsg            = "retry"
+	sleepTimeMsKey      = "TIME_SLEEP_TIME_MS"
+	skipCounterSleepKey = "TIME_SKIP_COUNTER_SLEEP"
 )
 
+func setSkipCounterSleep(ctx context.Context) context.Context {
+	return metainfo.WithPersistentValue(ctx, skipCounterSleepKey, "1")
+}
+
+func skipCounterSleep(ctx context.Context) bool {
+	if _, exist := metainfo.GetValue(ctx, skipCounterSleepKey); exist {
+		return true
+	}
+	if _, exist := metainfo.GetPersistentValue(ctx, skipCounterSleepKey); exist {
+		return true
+	}
+	return false
+}
+
+func getSleepTimeMS(ctx context.Context) time.Duration {
+	if value, exist := metainfo.GetPersistentValue(ctx, sleepTimeMsKey); exist {
+		if sleepTimeMS, err := strconv.Atoi(value); err == nil && sleepTimeMS > 0 {
+			return time.Duration(sleepTimeMS) * time.Millisecond
+		}
+	}
+	if value, exist := metainfo.GetValue(ctx, sleepTimeMsKey); exist {
+		if sleepTimeMS, err := strconv.Atoi(value); err == nil && sleepTimeMS > 0 {
+			return time.Duration(sleepTimeMS) * time.Millisecond
+		}
+	}
+	return 0
+}
+
 var _ stability.STService = &STServiceMockResultHandler{}
 
 // STServiceMockResultHandler .
@@ -91,10 +120,8 @@ func (h *STServiceMockResultHandler) CircuitBreakTest(ctx context.Context, req *
 	// use ttheader
 	if _, exist := metainfo.GetPersistentValue(ctx, retry.TransitKey); !exist {
 		sleepTime := 200 * time.Millisecond
-		if value, exist := metainfo.GetPersistentValue(ctx, sleepTimeMsKey); exist {
-			if sleepTimeMS, err := strconv.Atoi(value); err == nil && sleepTimeMS > 0 {
-				sleepTime = time.Duration(sleepTimeMS) * time.Millisecond
-			}
+		if v := getSleepTimeMS(ctx); v > 0 {
+			sleepTime = v
 		}
 		time.Sleep(sleepTime)
 	}

thriftrpc/retrycall/retrycall_test.go

@@ -24,16 +24,15 @@ import (
 	"time"
 
 	"github.com/bytedance/gopkg/cloud/metainfo"
-	"github.com/cloudwego/kitex/pkg/endpoint"
-	"github.com/cloudwego/kitex/pkg/kerrors"
-
 	"github.com/cloudwego/kitex-tests/kitex_gen/thrift/stability"
 	"github.com/cloudwego/kitex-tests/kitex_gen/thrift/stability/stservice"
 	"github.com/cloudwego/kitex-tests/pkg/test"
 	"github.com/cloudwego/kitex-tests/thriftrpc"
 	"github.com/cloudwego/kitex/client"
 	"github.com/cloudwego/kitex/client/callopt"
 	"github.com/cloudwego/kitex/pkg/circuitbreak"
+	"github.com/cloudwego/kitex/pkg/endpoint"
+	"github.com/cloudwego/kitex/pkg/kerrors"
 	"github.com/cloudwego/kitex/pkg/remote"
 	"github.com/cloudwego/kitex/pkg/retry"
 	"github.com/cloudwego/kitex/pkg/rpcinfo"
@@ -79,31 +78,69 @@ func genCBKey(ri rpcinfo.RPCInfo) string {
 }
 
 func TestRetryCB(t *testing.T) {
-	atomic.StoreInt32(&testSTReqCount, -1)
-	// retry config
-	fp := retry.NewFailurePolicy()
-
+	reqCount := int32(0)
 	cli = getKitexClient(
-		transport.PurePayload,
-		client.WithFailureRetry(fp),
+		transport.TTHeader,
+		client.WithRetryContainer(retry.NewRetryContainer()), // use default circuit breaker
+		client.WithFailureRetry(retry.NewFailurePolicy()),
 		client.WithRPCTimeout(20*time.Millisecond),
+		client.WithMiddleware(func(next endpoint.Endpoint) endpoint.Endpoint {
+			return func(ctx context.Context, req, resp interface{}) (err error) {
+				count := atomic.AddInt32(&reqCount, 1)
+				if count%10 == 0 {
+					time.Sleep(50 * time.Millisecond)
+				}
+				return nil // no need to send the real request
+			}
+		}),
 	)
 
 	ctx, stReq := thriftrpc.CreateSTRequest(context.Background())
 	cbCount := 0
 	for i := 0; i < 300; i++ {
-		stResp, err := cli.TestSTReq(ctx, stReq)
-		if err != nil {
-			test.Assert(t, strings.Contains(err.Error(), "retry circuit break"), err, i)
+		_, err := cli.TestSTReq(ctx, stReq)
+		if err != nil && strings.Contains(err.Error(), "retry circuit break") {
 			cbCount++
-		} else {
-			test.Assert(t, err == nil, err, i)
-			test.Assert(t, stReq.Str == stResp.Str)
 		}
 	}
+	test.Assert(t, reqCount == 300, reqCount)
 	test.Assert(t, cbCount == 30, cbCount)
 }
 
+func TestRetryCBPercentageLimit(t *testing.T) {
+	reqCount := int32(0)
+	cli := getKitexClient(
+		transport.TTHeaderFramed,
+		client.WithRetryContainer(retry.NewRetryContainerWithPercentageLimit()),
+		client.WithFailureRetry(retry.NewFailurePolicy()), // cb threshold = 10%
+		client.WithRPCTimeout(20*time.Millisecond),
+		client.WithMiddleware(func(next endpoint.Endpoint) endpoint.Endpoint {
+			return func(ctx context.Context, req, resp interface{}) (err error) {
+				atomic.AddInt32(&reqCount, 1)
+				if tm := getSleepTimeMS(ctx); tm > 0 {
+					time.Sleep(tm)
+				}
+				return nil // no need to send a real request
+			}
+		}),
+	)
+	ctx, stReq := thriftrpc.CreateSTRequest(context.Background())
+	ctx = setSkipCounterSleep(ctx)
+	cbCount := 0
+	for i := 0; i < 300; i++ {
+		reqCtx := ctx
+		if i%5 == 0 {
+			reqCtx = metainfo.WithPersistentValue(ctx, sleepTimeMsKey, "50")
+		}
+		_, err := cli.TestSTReq(reqCtx, stReq)
+		if err != nil && strings.Contains(err.Error(), "retry circuit break") {
+			cbCount++
+		}
+	}
+	test.Assert(t, reqCount == 333, reqCount)
+	test.Assert(t, cbCount == 58, cbCount)
+}
+
 func TestNoCB(t *testing.T) {
 	atomic.StoreInt32(&testSTReqCount, -1)
 	// retry config