client/service_inline_test.go

/*
 * Copyright 2023 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"
	"runtime"
	"runtime/debug"
	"sync/atomic"
	"testing"
	"time"

	"github.com/golang/mock/gomock"

	"github.com/cloudwego/kitex/client/callopt"
	"github.com/cloudwego/kitex/internal/client"
	"github.com/cloudwego/kitex/internal/mocks"
	internal_server "github.com/cloudwego/kitex/internal/server"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/consts"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
)

type serverInitialInfoImpl struct {
	EndpointsFunc func(ctx context.Context, req, resp interface{}) (err error)
}

func (s serverInitialInfoImpl) Endpoints() endpoint.Endpoint {
	if s.EndpointsFunc != nil {
		return s.EndpointsFunc
	}
	return func(ctx context.Context, req, resp interface{}) (err error) {
		return nil
	}
}

func (s serverInitialInfoImpl) Option() *internal_server.Options {
	return internal_server.NewOptions(nil)
}

func (s serverInitialInfoImpl) GetServiceInfos() map[string]*serviceinfo.ServiceInfo {
	return nil
}

func newMockServerInitialInfo() ServerInitialInfo {
	return &serverInitialInfoImpl{}
}

func newMockServiceInlineClient(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 := NewServiceInlineClient(svcInfo, newMockServerInitialInfo(), opts...)
	test.Assert(tb, err == nil)

	return cli
}

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

	mtd := mocks.MockMethod
	cli := newMockServiceInlineClient(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 TestServiceInlineTagOptions(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 := newMockServiceInlineClient(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 = newMockServiceInlineClient(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 TestServiceInlineTagOptionLocks0(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 := newMockServiceInlineClient(t, ctrl, options...)
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestServiceInlineTagOptionLocks1(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 := newMockServiceInlineClient(t, ctrl, options...)
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestServiceInlineTagOptionLocks2(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 := newMockServiceInlineClient(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 TestServiceInlineTimeoutOptions(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 := newMockServiceInlineClient(t, ctrl, options...)
	err := cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil)
}

func TestServiceInlineClientFinalizer(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 TestServiceInlineMethodKeyCall(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	mtd := mocks.MockMethod
	opts := []Option{
		WithTransHandlerFactory(newMockCliTransHandlerFactory(ctrl)),
		WithResolver(resolver404(ctrl)),
		WithDialer(newDialer(ctrl)),
		WithDestService("destService"),
	}
	svcInfo := mocks.ServiceInfo()
	s := serverInitialInfoImpl{}
	s.EndpointsFunc = func(ctx context.Context, req, resp interface{}) (err error) {
		test.Assert(t, ctx.Value(consts.CtxKeyMethod) == mtd)
		return nil
	}
	cli, err := NewServiceInlineClient(svcInfo, s, opts...)
	test.Assert(t, err == nil)
	ctx := context.Background()
	req := new(MockTStruct)
	res := new(MockTStruct)
	err = cli.Call(ctx, mtd, req, res)
	test.Assert(t, err == nil, err)
}