github.com/cloudwego/kitex@v0.9.0/client/client_test.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package client

import (
	"context"
	"errors"
	"net"
	"runtime"
	"runtime/debug"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/bytedance/gopkg/cloud/metainfo"
	"github.com/cloudwego/localsession"
	"github.com/cloudwego/localsession/backup"
	"github.com/cloudwego/netpoll"
	"github.com/golang/mock/gomock"

	"github.com/cloudwego/kitex/client/callopt"
	"github.com/cloudwego/kitex/internal/client"
	"github.com/cloudwego/kitex/internal/mocks"
	mocksnetpoll "github.com/cloudwego/kitex/internal/mocks/netpoll"
	mocksremote "github.com/cloudwego/kitex/internal/mocks/remote"
	"github.com/cloudwego/kitex/internal/mocks/rpc_info"
	mocksstats "github.com/cloudwego/kitex/internal/mocks/stats"
	mockthrift "github.com/cloudwego/kitex/internal/mocks/thrift"
	internal_stream "github.com/cloudwego/kitex/internal/stream"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/fallback"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/retry"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
	"github.com/cloudwego/kitex/pkg/rpctimeout"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/stats"
	"github.com/cloudwego/kitex/pkg/utils"
	"github.com/cloudwego/kitex/pkg/warmup"
	"github.com/cloudwego/kitex/transport"
)

var (
	// method="mock" typeID=thrift.REPLY seqID=1
	bs               = []byte{128, 1, 0, 2, 0, 0, 0, 4, 109, 111, 99, 107, 0, 0, 0, 1}
	mockWarmupOption = &warmup.ClientOption{
		ErrorHandling: warmup.ErrorLog,
		ResolverOption: &warmup.ResolverOption{
			Dests: []*rpcinfo.EndpointBasicInfo{
				{
					ServiceName: "mock_service",
					Method:      "mock_method",
				},
			},
		},
		PoolOption: &warmup.PoolOption{
			ConnNum:  128,
			Parallel: 128,
		},
	}
)

func newMockConn(ctrl *gomock.Controller) netpoll.Connection {
	conn := mocksnetpoll.NewMockConnection(ctrl)
	conn.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
		return copy(b, bs), nil
	}).AnyTimes()
	conn.EXPECT().Write(gomock.Any()).DoAndReturn(func(b []byte) (n int, err error) {
		return len(b), nil
	}).AnyTimes()
	conn.EXPECT().RemoteAddr().Return(utils.NewNetAddr("tcp", "mock")).AnyTimes()
	conn.EXPECT().IsActive().Return(true).AnyTimes()
	conn.EXPECT().Close().Return(nil).AnyTimes()
	return conn
}

func newDialer(ctrl *gomock.Controller) remote.Dialer {
	conn := newMockConn(ctrl)
	dialer := mocksremote.NewMockDialer(ctrl)
	dialer.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).Return(conn, nil).AnyTimes()
	return dialer
}

func newMockCliTransHandlerFactory(ctrl *gomock.Controller) remote.ClientTransHandlerFactory {
	handler := mocksremote.NewMockClientTransHandler(ctrl)
	handler.EXPECT().OnMessage(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, args, result remote.Message) (context.Context, error) {
		return ctx, nil
	}).AnyTimes()
	handler.EXPECT().Read(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, conn net.Conn, msg remote.Message) (context.Context, error) {
		return ctx, nil
	}).AnyTimes()
	handler.EXPECT().Write(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, conn net.Conn, msg remote.Message) (context.Context, error) {
		return ctx, nil
	}).AnyTimes()
	handler.EXPECT().SetPipeline(gomock.Any()).AnyTimes()
	handler.EXPECT().OnError(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()
	handler.EXPECT().OnInactive(gomock.Any(), gomock.Any()).AnyTimes()
	factory := mocksremote.NewMockClientTransHandlerFactory(ctrl)
	factory.EXPECT().NewTransHandler(gomock.Any()).DoAndReturn(func(opt *remote.ClientOption) (remote.ClientTransHandler, error) {
		return handler, nil
	}).AnyTimes()
	return factory
}

func newMockClient(tb testing.TB, ctrl *gomock.Controller, extra ...Option) Client {
	opts := []Option{
		WithTransHandlerFactory(newMockCliTransHandlerFactory(ctrl)),
		WithResolver(resolver404(ctrl)),
		WithDialer(newDialer(ctrl)),
		WithDestService("destService"),
	}
	opts = append(opts, extra...)

	svcInfo := mocks.ServiceInfo()
	cli, err := NewClient(svcInfo, opts...)
	test.Assert(tb, err == nil)

	return cli
}

func TestCall(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mtd := mocks.MockMethod
	cli := newMockClient(t, ctrl)
	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil, err)
}

func TestCallWithContextBackup(t *testing.T) {
	localsession.InitDefaultManager(localsession.DefaultManagerOptions())
	d, dd := "d", "dd"
	ctx := context.WithValue(context.Background(), "a", "a")
	ctx = context.WithValue(ctx, d, dd)
	ctx = metainfo.WithPersistentValues(ctx, "b", "b", "c", "c")
	ctx = metainfo.WithValues(ctx, "e", "e")
	backup.BackupCtx(ctx)

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mtd := mocks.MockMethod
	cli := newMockClient(t, ctrl, WithContextBackup(func(prev, cur context.Context) (ctx context.Context, backup bool) {
		if v := prev.Value(d); v != nil {
			cur = context.WithValue(cur, d, v)
		}
		return cur, true
	}), WithMiddleware(func(e endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, res interface{}) error {
			test.Assert(t, ctx.Value("a") == "aa")
			b, _ := metainfo.GetPersistentValue(ctx, "b")
			test.Assert(t, b == "bb")
			c, _ := metainfo.GetPersistentValue(ctx, "c")
			test.Assert(t, c == "c")
			test.Assert(t, ctx.Value(d) == dd)
			_, ok := metainfo.GetPersistentValue(ctx, "e")
			test.Assert(t, !ok)
			f, _ := metainfo.GetValue(ctx, "f")
			test.Assert(t, f == "f")
			return nil
		}
	}))
	req := new(MockTStruct)
	res := new(MockTStruct)

	ctx2 := context.WithValue(context.Background(), "a", "aa")
	ctx2 = metainfo.WithPersistentValue(ctx2, "b", "bb")
	ctx2 = metainfo.WithValue(ctx2, "f", "f")
	err := cli.Call(ctx2, mtd, req, res)
	test.Assert(t, err == nil, err)
}

func TestWithRetryOption(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockRetryContainer := retry.NewRetryContainer()
	cli := newMockClient(t, ctrl, WithRetryContainer(mockRetryContainer))

	test.Assert(t, cli.(*kcFinalizerClient).opt.RetryContainer == mockRetryContainer)
}

func BenchmarkCall(b *testing.B) {
	ctrl := gomock.NewController(b)
	defer ctrl.Finish()

	mtd := mocks.MockExceptionMethod
	cli := newMockClient(b, ctrl)
	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		cli.Call(ctx, mtd, req, res)
	}
}

func BenchmarkCallParallel(b *testing.B) {
	ctrl := gomock.NewController(b)
	defer ctrl.Finish()

	mtd := mocks.MockOnewayMethod
	cli := newMockClient(b, ctrl)
	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			cli.Call(ctx, mtd, req, res)
		}
	})
}

func TestTagOptions(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mtd := mocks.MockMethod
	tgs := map[string]string{}
	cls := func(m map[string]string) {
		for k := range m {
			delete(m, k)
		}
	}
	md := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, res interface{}) error {
			ri := rpcinfo.GetRPCInfo(ctx)
			test.Assert(t, ri != nil)

			to := ri.To()
			test.Assert(t, to != nil)

			re := remoteinfo.AsRemoteInfo(to)
			test.Assert(t, re != nil)

			for k, v := range tgs {
				val, ok := re.Tag(k)
				test.Assertf(t, ok, "expected tag not found: %s", k)
				test.Assertf(t, v == val, "values of tag '%s' not equal: '%s' != '%s'", k, v, val)
			}
			return nil
		}
	}
	var options []client.Option
	options = append(options, WithMiddleware(md))

	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	cli := newMockClient(t, ctrl, options...)
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)

	cls(tgs)
	tgs["cluster"] = "client cluster"
	tgs["idc"] = "client idc"
	cli = newMockClient(t, ctrl, WithTag("cluster", "client cluster"), WithTag("idc", "client idc"))
	err = cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)

	cls(tgs)
	tgs["cluster"] = "call cluster"
	tgs["idc"] = "call idc"
	ctx = NewCtxWithCallOptions(ctx, []callopt.Option{
		callopt.WithTag("cluster", "cluster"),
		callopt.WithTag("idc", "idc"),
	})
	err = cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestTagOptionLocks0(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mtd := mocks.MockMethod
	md := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, res interface{}) error {
			re := remoteinfo.AsRemoteInfo(rpcinfo.GetRPCInfo(ctx).To())

			var err error
			err = re.SetTag("cluster", "clusterx")
			test.Assert(t, err == nil)
			test.Assert(t, re.DefaultTag("cluster", "") == "clusterx")

			err = re.SetTag("idc", "idcx")
			test.Assert(t, err == nil)
			test.Assert(t, re.DefaultTag("idc", "") == "idcx")
			return nil
		}
	}
	var options []client.Option
	options = append(options, WithMiddleware(md))

	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	cli := newMockClient(t, ctrl, options...)
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestTagOptionLocks1(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mtd := mocks.MockMethod
	md := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, res interface{}) error {
			re := remoteinfo.AsRemoteInfo(rpcinfo.GetRPCInfo(ctx).To())

			var err error
			err = re.SetTag("cluster", "whatever")
			test.Assert(t, errors.Is(err, kerrors.ErrNotSupported))

			err = re.SetTag("idc", "whatever")
			test.Assert(t, errors.Is(err, kerrors.ErrNotSupported))
			return nil
		}
	}
	var options []client.Option
	options = append(options, WithMiddleware(md))
	options = append(options, WithTag("cluster", "client cluster"))
	options = append(options, WithTag("idc", "client idc"))

	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	cli := newMockClient(t, ctrl, options...)
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestTagOptionLocks2(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mtd := mocks.MockOnewayMethod
	md := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, res interface{}) error {
			re := remoteinfo.AsRemoteInfo(rpcinfo.GetRPCInfo(ctx).To())

			var err error
			err = re.SetTag("cluster", "whatever")
			test.Assert(t, errors.Is(err, kerrors.ErrNotSupported))

			err = re.SetTag("idc", "whatever")
			test.Assert(t, errors.Is(err, kerrors.ErrNotSupported))
			return nil
		}
	}
	var options []client.Option
	options = append(options, WithMiddleware(md))

	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	cli := newMockClient(t, ctrl, options...)
	ctx = NewCtxWithCallOptions(ctx, []callopt.Option{
		callopt.WithTag("cluster", "cluster"),
		callopt.WithTag("idc", "idc"),
	})
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestWarmingUpOption(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var options []client.Option
	options = append(options, WithWarmingUp(mockWarmupOption))

	cli := newMockClient(t, ctrl, options...)
	test.Assert(t, cli.(*kcFinalizerClient).opt.WarmUpOption == mockWarmupOption)
}

func TestTimeoutOptions(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mtd := mocks.MockMethod
	md := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, res interface{}) error {
			ri := rpcinfo.GetRPCInfo(ctx)
			test.Assert(t, ri != nil)

			cfgs := ri.Config()
			test.Assert(t, cfgs != nil)

			mcfg := rpcinfo.AsMutableRPCConfig(cfgs)
			test.Assert(t, mcfg != nil)

			var err error
			err = mcfg.SetRPCTimeout(time.Hour)
			test.Assert(t, err == nil)
			test.Assert(t, cfgs.RPCTimeout() == time.Hour)

			err = mcfg.SetConnectTimeout(time.Hour * 2)
			test.Assert(t, err == nil)
			test.Assert(t, cfgs.ConnectTimeout() == time.Hour*2)
			return nil
		}
	}
	var options []client.Option
	options = append(options, WithMiddleware(md))

	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	cli := newMockClient(t, ctrl, options...)
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestTimeoutCtxCall(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mtd := mocks.MockMethod
	var accessed int32
	md := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, resp interface{}) (err error) {
			time.Sleep(time.Millisecond * 100)
			atomic.StoreInt32(&accessed, 1)
			return next(ctx, req, resp)
		}
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*10)
	req, res := new(MockTStruct), new(MockTStruct)
	cli := newMockClient(t, ctrl, WithMiddleware(md))
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout))
	test.Assert(t, atomic.LoadInt32(&accessed) == 0)
	cancel()

	err = cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, err == nil)
	test.Assert(t, atomic.LoadInt32(&accessed) == 1)
}

func TestTimeoutOptionsLock0(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	tos := []time.Duration{time.Hour / 2, time.Hour / 3}
	mtd := mocks.MockMethod
	md := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, res interface{}) error {
			ri := rpcinfo.GetRPCInfo(ctx)
			test.Assert(t, ri != nil)

			cfgs := ri.Config()
			test.Assert(t, cfgs != nil)

			mcfg := rpcinfo.AsMutableRPCConfig(cfgs)
			test.Assert(t, mcfg != nil)

			var err error
			err = mcfg.SetRPCTimeout(time.Hour)
			test.Assert(t, errors.Is(err, kerrors.ErrNotSupported))
			test.Assert(t, cfgs.RPCTimeout() == tos[0])

			err = mcfg.SetConnectTimeout(time.Hour * 2)
			test.Assert(t, errors.Is(err, kerrors.ErrNotSupported))
			test.Assert(t, cfgs.ConnectTimeout() == tos[1])
			return nil
		}
	}
	var options []client.Option
	options = append(options, WithMiddleware(md))
	options = append(options, WithRPCTimeout(tos[0]))
	options = append(options, WithConnectTimeout(tos[1]))

	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)

	cli := newMockClient(t, ctrl, options...)
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)

	tos[0], tos[1] = time.Minute, time.Minute*2
	ctx = NewCtxWithCallOptions(ctx, []callopt.Option{
		callopt.WithRPCTimeout(tos[0]),
		callopt.WithConnectTimeout(tos[1]),
	})
	err = cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestAdjustTimeout(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	adjustTimeoutMWB := func(moreTimeout time.Duration) endpoint.MiddlewareBuilder {
		return func(mwCtx context.Context) endpoint.Middleware {
			if p, ok := mwCtx.Value(rpctimeout.TimeoutAdjustKey).(*time.Duration); ok {
				*p = moreTimeout
			}
			return endpoint.DummyMiddleware
		}
	}
	mtd := mocks.MockMethod
	md := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, resp interface{}) (err error) {
			time.Sleep(time.Second)
			return nil
		}
	}
	req, res := new(MockTStruct), new(MockTStruct)

	// should timeout
	cli := newMockClient(t, ctrl,
		WithMiddlewareBuilder(adjustTimeoutMWB(0)),
		WithMiddleware(md),
		WithRPCTimeout(time.Millisecond*500))
	err := cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout))
	// shouldn't timeout
	cli = newMockClient(t, ctrl,
		WithMiddlewareBuilder(adjustTimeoutMWB(time.Second)),
		WithMiddleware(md),
		WithRPCTimeout(time.Millisecond*500))
	err = cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, err == nil, err)
}

func TestRetry(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var count int32
	tp := rpc_info.NewMockTimeoutProvider(ctrl)
	tp.EXPECT().Timeouts(gomock.Any()).DoAndReturn(func(ri rpcinfo.RPCInfo) rpcinfo.Timeouts {
		retryTimeStr, ok := ri.To().Tag(rpcinfo.RetryTag)
		if atomic.AddInt32(&count, 1) == 1 {
			// first request
			test.Assert(t, !ok)
		} else {
			// retry request
			test.Assert(t, retryTimeStr == "1", retryTimeStr)
		}
		return ri.Config()
	}).AnyTimes()
	mockTimeoutMW := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, resp interface{}) (err error) {
			ri := rpcinfo.GetRPCInfo(ctx)
			if atomic.LoadInt32(&count) == 1 {
				return kerrors.ErrRPCTimeout
			} else {
				retryTimeStr := ri.To().DefaultTag(rpcinfo.RetryTag, "0")
				test.Assert(t, retryTimeStr == "1", retryTimeStr)
				r, _ := metainfo.GetPersistentValue(ctx, retry.TransitKey)
				test.Assert(t, r == "1", r)
			}
			return nil
		}
	}

	// case1: return timeout
	cli := newMockClient(t, ctrl,
		WithTimeoutProvider(tp),
		WithMiddleware(mockTimeoutMW))
	mtd := mocks.MockMethod
	req, res := new(MockTStruct), new(MockTStruct)
	err := cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout))
	test.Assert(t, count == 1, count)

	// case 2: have timeout retry, return success
	count = 0
	cli = newMockClient(t, ctrl,
		WithTimeoutProvider(tp),
		WithMiddleware(mockTimeoutMW),
		WithFailureRetry(&retry.FailurePolicy{
			StopPolicy: retry.StopPolicy{
				MaxRetryTimes: 3,
				CBPolicy: retry.CBPolicy{
					ErrorRate: 0.1,
				},
			},
			RetrySameNode: true,
		}))
	err = cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, err == nil, err)
	test.Assert(t, count == 2, count)
}

func TestRetryWithDifferentTimeout(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var count int32
	tp := rpc_info.NewMockTimeoutProvider(ctrl)
	tp.EXPECT().Timeouts(gomock.Any()).DoAndReturn(func(ri rpcinfo.RPCInfo) rpcinfo.Timeouts {
		cfg := rpcinfo.AsMutableRPCConfig(ri.Config()).Clone()
		retryTimeStr, ok := ri.To().Tag(rpcinfo.RetryTag)
		if atomic.AddInt32(&count, 1) == 1 {
			// first request
			test.Assert(t, !ok)
			cfg.SetRPCTimeout(100 * time.Millisecond)
		} else {
			// retry request
			test.Assert(t, retryTimeStr == "1", retryTimeStr)
			// reset rpc timeout for retry request
			cfg.SetRPCTimeout(500 * time.Millisecond)
		}
		return cfg.ImmutableView()
	}).AnyTimes()
	mockTimeoutMW := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, resp interface{}) (err error) {
			ri := rpcinfo.GetRPCInfo(ctx)
			if atomic.LoadInt32(&count) == 1 {
				test.Assert(t, ri.Config().RPCTimeout().Milliseconds() == 100, ri.Config().RPCTimeout().Milliseconds())
			} else {
				retryTimeStr := ri.To().DefaultTag(rpcinfo.RetryTag, "0")
				test.Assert(t, retryTimeStr == "1", retryTimeStr)
				r, _ := metainfo.GetPersistentValue(ctx, retry.TransitKey)
				test.Assert(t, r == "1", r)
				test.Assert(t, ri.Config().RPCTimeout().Milliseconds() == 500, ri.Config().RPCTimeout().Milliseconds())
			}
			time.Sleep(200 * time.Millisecond)
			return nil
		}
	}

	// case: sleep 200ms
	//   first request timeout is 100 then the timeout happens
	//   retry request timeout is 500 then success
	cli := newMockClient(t, ctrl,
		WithTimeoutProvider(tp),
		WithMiddleware(mockTimeoutMW),
		WithFailureRetry(&retry.FailurePolicy{
			StopPolicy: retry.StopPolicy{
				MaxRetryTimes: 3,
				CBPolicy: retry.CBPolicy{
					ErrorRate: 0.1,
				},
			},
			RetrySameNode: true,
		}))
	mtd := mocks.MockMethod
	req, res := new(MockTStruct), new(MockTStruct)
	err := cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, err == nil, err)
	test.Assert(t, count == 2, count)
}

func TestRetryWithResultRetry(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockErr := errors.New("mock")
	retryWithMockErr := false
	errMW := func(next endpoint.Endpoint) endpoint.Endpoint {
		var count int32
		return func(ctx context.Context, req, resp interface{}) (err error) {
			if atomic.CompareAndSwapInt32(&count, 0, 1) {
				return mockErr
			}
			return nil
		}
	}
	mockTimeoutMW := func(next endpoint.Endpoint) endpoint.Endpoint {
		var count int32
		return func(ctx context.Context, req, resp interface{}) (err error) {
			if atomic.CompareAndSwapInt32(&count, 0, 1) {
				time.Sleep(300 * time.Millisecond)
			}
			return nil
		}
	}
	errRetryFunc := func(err error, ri rpcinfo.RPCInfo) bool {
		if errors.Is(err, mockErr) {
			retryWithMockErr = true
			return true
		}
		return false
	}

	// case 1: retry for mockErr
	cli := newMockClient(t, ctrl,
		WithMiddleware(errMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFailureRetry(&retry.FailurePolicy{
			StopPolicy: retry.StopPolicy{
				MaxRetryTimes: 3,
				CBPolicy: retry.CBPolicy{
					ErrorRate: 0.1,
				},
			},
			RetrySameNode:     true,
			ShouldResultRetry: &retry.ShouldResultRetry{ErrorRetry: errRetryFunc},
		}))
	mtd := mocks.MockMethod
	req, res := new(MockTStruct), new(MockTStruct)
	err := cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, err == nil, err)
	test.Assert(t, retryWithMockErr)

	// case 1: retry for timeout
	cli = newMockClient(t, ctrl,
		WithMiddleware(mockTimeoutMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFailureRetry(&retry.FailurePolicy{
			StopPolicy: retry.StopPolicy{
				MaxRetryTimes: 3,
				CBPolicy: retry.CBPolicy{
					ErrorRate: 0.1,
				},
			},
			RetrySameNode:     true,
			ShouldResultRetry: &retry.ShouldResultRetry{ErrorRetry: errRetryFunc},
		}))
	err = cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, err == nil, err)

	// case 2: set NotRetryForTimeout as true, won't do retry for timeout
	cli = newMockClient(t, ctrl,
		WithMiddleware(mockTimeoutMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFailureRetry(&retry.FailurePolicy{
			StopPolicy: retry.StopPolicy{
				MaxRetryTimes: 3,
				CBPolicy: retry.CBPolicy{
					ErrorRate: 0.1,
				},
			},
			RetrySameNode:     true,
			ShouldResultRetry: &retry.ShouldResultRetry{ErrorRetry: errRetryFunc, NotRetryForTimeout: true},
		}))
	err = cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, errors.Is(err, kerrors.ErrRPCTimeout))
}

func TestFallbackForError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	isTimeout := false
	isMockErr := false
	isBizErr := false
	errForReportIsNil := false
	var rpcResult interface{}

	// prepare mock data
	retStr := "success"
	sleepMW := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, resp interface{}) (err error) {
			time.Sleep(300 * time.Millisecond)
			return nil
		}
	}
	mockErr := errors.New("mock")
	mockErrMW := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, resp interface{}) (err error) {
			return mockErr
		}
	}
	bizErrMW := func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, resp interface{}) (err error) {
			if setter, ok := rpcinfo.GetRPCInfo(ctx).Invocation().(rpcinfo.InvocationSetter); ok {
				setter.SetBizStatusErr(kerrors.NewBizStatusError(100, mockErr.Error()))
			}
			return nil
		}
	}

	mockTracer := mocksstats.NewMockTracer(ctrl)
	mockTracer.EXPECT().Start(gomock.Any()).DoAndReturn(func(ctx context.Context) context.Context {
		return ctx
	}).AnyTimes()
	mockTracer.EXPECT().Finish(gomock.Any()).DoAndReturn(func(ctx context.Context) {
		if rpcinfo.GetRPCInfo(ctx).Stats().Error() == nil {
			errForReportIsNil = true
		} else {
			errForReportIsNil = false
		}
	}).AnyTimes()
	newFallback := func(realReqResp bool, testCase string) *fallback.Policy {
		var fallbackFunc fallback.Func
		checkErr := func(err error) {
			isTimeout, isMockErr, isBizErr = false, false, false
			if errors.Is(err, kerrors.ErrRPCTimeout) {
				isTimeout = true
			}
			if errors.Is(err, mockErr) {
				isMockErr = true
			}
			if _, ok := kerrors.FromBizStatusError(err); ok {
				isBizErr = true
			}
		}
		if realReqResp {
			fallbackFunc = fallback.UnwrapHelper(func(ctx context.Context, req, resp interface{}, err error) (fbResp interface{}, fbErr error) {
				checkErr(err)
				_, ok := req.(*mockthrift.MockReq)
				test.Assert(t, ok, testCase)
				return &retStr, nil
			})
		} else {
			fallbackFunc = func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
				checkErr(err)
				_, ok := args.(*mockthrift.MockTestArgs)
				test.Assert(t, ok, testCase)
				result.SetSuccess(&retStr)
				return nil
			}
		}
		return fallback.NewFallbackPolicy(fallbackFunc)
	}

	// case 1: fallback for timeout, return nil err, but report original err
	cli := newMockClient(t, ctrl,
		WithMiddleware(sleepMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFallback(newFallback(false, "case 1")),
		WithTracer(mockTracer),
	)
	mtd := mocks.MockMethod
	rpcResult = mockthrift.NewMockTestResult()
	err := cli.Call(context.Background(), mtd, mockthrift.NewMockTestArgs(), rpcResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, *rpcResult.(utils.KitexResult).GetResult().(*string) == retStr)
	test.Assert(t, isTimeout, err)
	test.Assert(t, !errForReportIsNil)

	// case 2: fallback for mock error, but report original err
	cli = newMockClient(t, ctrl,
		WithMiddleware(mockErrMW),
		WithFallback(newFallback(false, "case 2")))
	rpcResult = mockthrift.NewMockTestResult()
	err = cli.Call(context.Background(), mtd, mockthrift.NewMockTestArgs(), rpcResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, *rpcResult.(utils.KitexResult).GetResult().(*string) == retStr)
	test.Assert(t, isMockErr, err)
	test.Assert(t, !errForReportIsNil)

	// case 3: fallback for timeout, return nil err, and report nil err as enable reportAsFallback
	cli = newMockClient(t, ctrl,
		WithMiddleware(sleepMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFallback(newFallback(true, "case 3").EnableReportAsFallback()),
		WithTracer(mockTracer),
	)
	// reset
	isTimeout = false
	errForReportIsNil = false
	rpcResult = mockthrift.NewMockTestResult()
	err = cli.Call(context.Background(), mtd, mockthrift.NewMockTestArgs(), rpcResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, *rpcResult.(utils.KitexResult).GetResult().(*string) == retStr)
	test.Assert(t, isTimeout, err)
	test.Assert(t, errForReportIsNil)

	// case 4: no fallback, return biz err, but report nil err
	cli = newMockClient(t, ctrl,
		WithMiddleware(bizErrMW),
		WithRPCTimeout(100*time.Millisecond),
		WithTracer(mockTracer),
	)
	// reset
	errForReportIsNil = false
	rpcResult = mockthrift.NewMockTestResult()
	err = cli.Call(context.Background(), mtd, mockthrift.NewMockTestArgs(), rpcResult)
	_, ok := kerrors.FromBizStatusError(err)
	test.Assert(t, ok, err)
	test.Assert(t, errForReportIsNil)

	// case 5: fallback for biz error, return nil err,
	// and report nil err even if don't enable reportAsFallback as biz error won't report failure by design
	cli = newMockClient(t, ctrl,
		WithMiddleware(bizErrMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFallback(newFallback(true, "case 5")),
		WithTracer(mockTracer),
	)
	// reset
	isBizErr = false
	errForReportIsNil = false
	rpcResult = mockthrift.NewMockTestResult()
	err = cli.Call(context.Background(), mtd, mockthrift.NewMockTestArgs(), rpcResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, *rpcResult.(utils.KitexResult).GetResult().(*string) == retStr)
	test.Assert(t, isBizErr, err)
	test.Assert(t, errForReportIsNil)

	// case 6: fallback for timeout error, return nil err
	cli = newMockClient(t, ctrl,
		WithMiddleware(sleepMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFallback(fallback.TimeoutAndCBFallback(fallback.UnwrapHelper(func(ctx context.Context, req, resp interface{}, err error) (fbResp interface{}, fbErr error) {
			if errors.Is(err, kerrors.ErrRPCTimeout) {
				isTimeout = true
			}
			return &retStr, nil
		}))),
		WithTracer(mockTracer),
	)
	// reset
	isTimeout = false
	errForReportIsNil = true
	rpcResult = mockthrift.NewMockTestResult()
	err = cli.Call(context.Background(), mtd, mockthrift.NewMockTestArgs(), rpcResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, *rpcResult.(utils.KitexResult).GetResult().(*string) == retStr, rpcResult)
	test.Assert(t, isTimeout, err)
	test.Assert(t, !errForReportIsNil)

	// case 6: use TimeoutAndCBFallback, non-timeout and non-CB error cannot do fallback
	var fallbackExecuted bool
	cli = newMockClient(t, ctrl,
		WithMiddleware(mockErrMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFallback(fallback.TimeoutAndCBFallback(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
			fallbackExecuted = true
			return
		})),
		WithTracer(mockTracer),
	)
	// reset
	errForReportIsNil = true
	rpcResult = mockthrift.NewMockTestResult()
	err = cli.Call(context.Background(), mtd, mockthrift.NewMockTestArgs(), rpcResult)
	test.Assert(t, err != nil, err)
	test.Assert(t, !fallbackExecuted)

	// case 7: fallback return nil-resp and nil-err, framework will return the original resp and err
	cli = newMockClient(t, ctrl,
		WithMiddleware(mockErrMW),
		WithRPCTimeout(100*time.Millisecond),
		WithFallback(fallback.NewFallbackPolicy(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
			return
		})),
		WithTracer(mockTracer),
	)
	mtd = mocks.MockMethod
	rpcResult = mockthrift.NewMockTestResult()
	err = cli.Call(context.Background(), mtd, mockthrift.NewMockTestArgs(), rpcResult)
	test.Assert(t, err == mockErr)
	test.Assert(t, !errForReportIsNil)
}

func TestClientFinalizer(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	debug.SetGCPercent(-1)
	defer debug.SetGCPercent(100)

	var ms runtime.MemStats
	runtime.ReadMemStats(&ms)
	t.Logf("Before new clients, allocation: %f Mb, Number of allocation: %d\n", mb(ms.HeapAlloc), ms.HeapObjects)

	var (
		closeCalledCnt int32
		succeedCnt     = 10000
		failedCnt      = 10000
		cliCnt         = succeedCnt + failedCnt
	)
	clis := make([]Client, cliCnt)
	// clients that init successfully.
	for i := 0; i < succeedCnt; i++ {
		svcInfo := mocks.ServiceInfo()
		mockClient, err := NewClient(svcInfo, WithDestService("destService"), WithShortConnection(),
			WithCloseCallbacks(func() error {
				atomic.AddInt32(&closeCalledCnt, 1)
				return nil
			}))
		test.Assert(t, err == nil, err)
		clis[i] = mockClient
	}
	// clients that init failed, closeCallback should be called
	for i := succeedCnt; i < cliCnt; i++ {
		mockClient, err := NewClient(svcInfo, WithDestService(""), WithShortConnection(),
			WithCloseCallbacks(func() error {
				atomic.AddInt32(&closeCalledCnt, 1)
				return nil
			}))
		test.Assert(t, err != nil, err)
		clis[i] = mockClient
	}

	runtime.ReadMemStats(&ms)
	t.Logf("After new clients, allocation: %f Mb, Number of allocation: %d\n", mb(ms.HeapAlloc), ms.HeapObjects)

	runtime.GC()
	runtime.ReadMemStats(&ms)
	firstGCHeapAlloc, firstGCHeapObjects := mb(ms.HeapAlloc), ms.HeapObjects
	t.Logf("After first GC, allocation: %f Mb, Number of allocation: %d\n", firstGCHeapAlloc, firstGCHeapObjects)
	time.Sleep(200 * time.Millisecond)                                 // ensure the finalizer be executed
	test.Assert(t, atomic.LoadInt32(&closeCalledCnt) == int32(cliCnt)) // ensure CloseCallback of client has been called

	runtime.GC()
	runtime.ReadMemStats(&ms)
	secondGCHeapAlloc, secondGCHeapObjects := mb(ms.HeapAlloc), ms.HeapObjects
	t.Logf("After second GC, allocation: %f Mb, Number of allocation: %d\n", secondGCHeapAlloc, secondGCHeapObjects)
	test.Assert(t, secondGCHeapAlloc < firstGCHeapAlloc/2 && secondGCHeapObjects < firstGCHeapObjects/2)
}

func TestPanicInMiddleware(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	reportPanic := false
	mockTracer := mocksstats.NewMockTracer(ctrl)
	mockTracer.EXPECT().Start(gomock.Any()).DoAndReturn(func(ctx context.Context) context.Context {
		return ctx
	}).AnyTimes()
	mockTracer.EXPECT().Finish(gomock.Any()).DoAndReturn(func(ctx context.Context) {
		if ok, _ := rpcinfo.GetRPCInfo(ctx).Stats().Panicked(); ok {
			reportPanic = true
		}
	}).AnyTimes()
	panicMW := func(endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request, response interface{}) (err error) {
			panic(panicMsg)
		}
	}

	// panic with timeout, the panic is recovered in timeout mw
	cli := newMockClient(t, ctrl,
		WithMiddleware(panicMW),
		WithRPCTimeout(100*time.Millisecond),
		WithTracer(mockTracer))
	mtd := mocks.MockMethod
	req, res := new(MockTStruct), new(MockTStruct)
	err := cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, err != nil)
	test.Assert(t, strings.Contains(err.Error(), panicMsg), err)
	test.Assert(t, reportPanic, err)

	// panic without timeout, the panic is recovered in client.Call
	reportPanic = false
	cli = newMockClient(t, ctrl,
		WithMiddleware(panicMW),
		WithRPCTimeout(0),
		WithTracer(mockTracer))
	req, res = new(MockTStruct), new(MockTStruct)
	err = cli.Call(context.Background(), mtd, req, res)
	test.Assert(t, err != nil)
	test.Assert(t, strings.Contains(err.Error(), panicMsg), err)
	test.Assert(t, reportPanic, err)
}

func TestMethodInfoIsNil(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	cli := newMockClient(t, ctrl)
	req, res := new(MockTStruct), new(MockTStruct)
	err := cli.Call(context.Background(), "notExist", req, res)
	test.Assert(t, err != nil)
	test.Assert(t, strings.Contains(err.Error(), "method info is nil"))
}

func mb(byteSize uint64) float32 {
	return float32(byteSize) / float32(1024*1024)
}

func Test_initTransportProtocol(t *testing.T) {
	t.Run("kitex_pb", func(t *testing.T) {
		svcInfo := &serviceinfo.ServiceInfo{
			PayloadCodec: serviceinfo.Protobuf,
		}
		rpcConfig := rpcinfo.NewRPCConfig()

		initTransportProtocol(svcInfo, rpcConfig)
		test.Assert(t, rpcConfig.PayloadCodec() == serviceinfo.Protobuf, rpcConfig.PayloadCodec())
		test.Assert(t, rpcConfig.TransportProtocol() == transport.TTHeaderFramed, rpcConfig.TransportProtocol())
	})

	t.Run("grpc_pb", func(t *testing.T) {
		svcInfo := &serviceinfo.ServiceInfo{
			PayloadCodec: serviceinfo.Protobuf,
		}
		rpcConfig := rpcinfo.NewRPCConfig()
		rpcinfo.AsMutableRPCConfig(rpcConfig).SetTransportProtocol(transport.GRPC)

		initTransportProtocol(svcInfo, rpcConfig)
		test.Assert(t, rpcConfig.PayloadCodec() == serviceinfo.Protobuf, rpcConfig.PayloadCodec())
		test.Assert(t, rpcConfig.TransportProtocol() == transport.GRPC, rpcConfig.TransportProtocol())
	})

	t.Run("grpc_thrift", func(t *testing.T) {
		svcInfo := &serviceinfo.ServiceInfo{
			PayloadCodec: serviceinfo.Thrift,
		}
		rpcConfig := rpcinfo.NewRPCConfig()
		rpcinfo.AsMutableRPCConfig(rpcConfig).SetTransportProtocol(transport.GRPC)

		initTransportProtocol(svcInfo, rpcConfig)
		test.Assert(t, rpcConfig.PayloadCodec() == serviceinfo.Thrift, rpcConfig.PayloadCodec())
		test.Assert(t, rpcConfig.TransportProtocol() == transport.GRPC, rpcConfig.TransportProtocol())
	})
}

type mockTracerForInitMW struct {
	stats.Tracer
	rpcinfo.StreamEventReporter
}

func Test_kClient_initStreamMiddlewares(t *testing.T) {
	ctl := &rpcinfo.TraceController{}
	ctl.Append(mockTracerForInitMW{})
	c := &kClient{
		opt: &client.Options{
			TracerCtl: ctl,
			Streaming: internal_stream.StreamingConfig{},
		},
	}
	c.initStreamMiddlewares(context.Background())

	test.Assert(t, len(c.opt.Streaming.RecvMiddlewares) > 0, "init middlewares failed")
	test.Assert(t, len(c.opt.Streaming.SendMiddlewares) > 0, "init middlewares failed")
}