github.com/cloudwego/kitex@v0.9.0/server/option_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 server

import (
	"context"
	"math/rand"
	"net"
	"testing"
	"time"

	"github.com/cloudwego/kitex/internal/mocks"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/diagnosis"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/codec/protobuf"
	"github.com/cloudwego/kitex/pkg/remote/trans/detection"
	"github.com/cloudwego/kitex/pkg/remote/trans/netpoll"
	"github.com/cloudwego/kitex/pkg/remote/trans/netpollmux"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/grpc"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/stats"
	"github.com/cloudwego/kitex/pkg/utils"
)

// TestOptionDebugInfo tests the creation of a server with DebugService option
func TestOptionDebugInfo(t *testing.T) {
	var opts []Option
	md := newMockDiagnosis()
	opts = append(opts, WithDiagnosisService(md))
	buildMw := 0
	opts = append(opts, WithMiddleware(func(next endpoint.Endpoint) endpoint.Endpoint {
		buildMw++
		return func(ctx context.Context, req, resp interface{}) (err error) {
			return next(ctx, req, resp)
		}
	}))
	opts = append(opts, WithMiddlewareBuilder(func(ctx context.Context) endpoint.Middleware {
		return func(next endpoint.Endpoint) endpoint.Endpoint {
			buildMw++
			return next
		}
	}))

	svr := NewServer(opts...)
	// check middleware build
	test.Assert(t, buildMw == 2)

	// check probe result
	pp := md.ProbePairs()
	test.Assert(t, pp != nil)

	optFunc := pp[diagnosis.OptionsKey]
	test.Assert(t, optFunc != nil)

	is := svr.(*server)
	debugInfo := is.opt.DebugInfo
	optDebugInfo := optFunc().(utils.Slice)
	test.Assert(t, len(debugInfo) == len(optDebugInfo))
	test.Assert(t, len(debugInfo) == 3)

	for i := range debugInfo {
		test.Assert(t, debugInfo[i] == optDebugInfo[i])
	}
}

// TestProxyOption tests the creation of a server with Proxy option
func TestProxyOption(t *testing.T) {
	var opts []Option
	addr, err := net.ResolveTCPAddr("tcp", ":8888")
	test.Assert(t, err == nil, err)
	opts = append(opts, WithServiceAddr(addr))
	opts = append(opts, WithProxy(&proxyMock{}))
	svr := NewServer(opts...)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr.Stop()
		test.Assert(t, err == nil, err)
	})

	err = svr.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil)

	err = svr.Run()
	test.Assert(t, err == nil, err)
	iSvr := svr.(*server)
	test.Assert(t, iSvr.opt.RemoteOpt.Address.String() == "mock", iSvr.opt.RemoteOpt.Address.Network())
}

type proxyMock struct{}

func (p *proxyMock) Replace(addr net.Addr) (nAddr net.Addr, err error) {
	if nAddr, err = net.ResolveUnixAddr("unix", "mock"); err == nil {
		return nAddr, err
	}
	return
}

func newMockDiagnosis() *mockDiagnosis {
	return &mockDiagnosis{probes: make(map[diagnosis.ProbeName]diagnosis.ProbeFunc)}
}

type mockDiagnosis struct {
	probes map[diagnosis.ProbeName]diagnosis.ProbeFunc
}

func (m *mockDiagnosis) RegisterProbeFunc(name diagnosis.ProbeName, probeFunc diagnosis.ProbeFunc) {
	m.probes[name] = probeFunc
}

func (m *mockDiagnosis) ProbePairs() map[diagnosis.ProbeName]diagnosis.ProbeFunc {
	return m.probes
}

// TestExitWaitTimeOption tests the creation of a server with RemoteOpt.ExitWaitTime option
func TestExitWaitTimeOption(t *testing.T) {
	// random timeout value
	testTimeOut := time.Duration(time.Now().Unix()) * time.Microsecond
	svr := NewServer(WithExitWaitTime(testTimeOut))
	err := svr.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr.Run()
	test.Assert(t, err == nil, err)
	iSvr := svr.(*server)
	test.Assert(t, iSvr.opt.RemoteOpt.ExitWaitTime == testTimeOut)
}

// TestMaxConnIdleTimeOption tests the creation of a server with RemoteOpt.MaxConnectionIdleTime option
func TestMaxConnIdleTimeOption(t *testing.T) {
	// random timeout value
	testTimeOut := time.Duration(time.Now().Unix())
	svr := NewServer(WithMaxConnIdleTime(testTimeOut))
	err := svr.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr.Run()
	test.Assert(t, err == nil, err)
	iSvr := svr.(*server)
	test.Assert(t, iSvr.opt.RemoteOpt.MaxConnectionIdleTime == testTimeOut)
}

type myTracer struct{}

func (t *myTracer) Start(ctx context.Context) context.Context {
	return nil
}

func (t *myTracer) Finish(ctx context.Context) {
}

// TestTracerOption tests the creation of a server with TracerCtl option
func TestTracerOption(t *testing.T) {
	svr1 := NewServer()
	err := svr1.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr1.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr1.Run()
	test.Assert(t, err == nil, err)
	iSvr1 := svr1.(*server)
	test.Assert(t, iSvr1.opt.TracerCtl != nil)
	test.Assert(t, iSvr1.opt.TracerCtl.HasTracer() != true)

	tracer := &myTracer{}
	svr2 := NewServer(WithTracer(tracer))
	err = svr2.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr2.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr2.Run()
	test.Assert(t, err == nil, err)
	iSvr2 := svr2.(*server)
	test.Assert(t, iSvr2.opt.TracerCtl != nil)
	test.Assert(t, iSvr2.opt.TracerCtl.HasTracer())
}

// TestStatsLevelOption tests the creation of a server with StatsLevel option
func TestStatsLevelOption(t *testing.T) {
	svr1 := NewServer()
	err := svr1.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr1.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr1.Run()
	test.Assert(t, err == nil, err)
	iSvr1 := svr1.(*server)
	test.Assert(t, iSvr1.opt.StatsLevel != nil)
	test.Assert(t, *iSvr1.opt.StatsLevel == stats.LevelDisabled)

	svr2 := NewServer(WithStatsLevel(stats.LevelDetailed))
	err = svr2.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr2.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr2.Run()
	test.Assert(t, err == nil, err)
	iSvr2 := svr2.(*server)
	test.Assert(t, *iSvr2.opt.StatsLevel == stats.LevelDetailed)
}

type mySuiteOption struct {
	myOpts []Option
}

func (s *mySuiteOption) Options() []Option {
	return s.myOpts
}

// TestSuiteOption tests the creation of a server with SuiteOption option
func TestSuiteOption(t *testing.T) {
	svr1 := NewServer()
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr1.Stop()
		test.Assert(t, err == nil, err)
	})
	err := svr1.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr1.Run()
	test.Assert(t, err == nil, err)
	iSvr1 := svr1.(*server)
	defaultExitWaitTime := iSvr1.opt.RemoteOpt.ExitWaitTime
	defaultConnectionIdleTime := iSvr1.opt.RemoteOpt.MaxConnectionIdleTime
	test.Assert(t, defaultExitWaitTime != 0)
	test.Assert(t, defaultConnectionIdleTime != 0)

	tmpWaitTime := defaultExitWaitTime*2 + time.Duration(rand.Intn(20))
	tmpConnIdleTime := defaultConnectionIdleTime*2 + time.Duration(rand.Intn(20))
	suiteOpt := &mySuiteOption{myOpts: []Option{
		WithExitWaitTime(tmpWaitTime),
		WithMaxConnIdleTime(tmpConnIdleTime),
	}}
	svr2 := NewServer(WithSuite(suiteOpt))
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr2.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr2.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr2.Run()
	test.Assert(t, err == nil, err)
	iSvr2 := svr2.(*server)
	test.Assert(t, iSvr2.opt.RemoteOpt.ExitWaitTime != defaultExitWaitTime, iSvr2.opt.RemoteOpt.ExitWaitTime)
	test.Assert(t, iSvr2.opt.RemoteOpt.MaxConnectionIdleTime != defaultConnectionIdleTime, iSvr2.opt.RemoteOpt.ExitWaitTime)

	test.Assert(t, iSvr2.opt.RemoteOpt.ExitWaitTime == tmpWaitTime, iSvr2.opt.RemoteOpt.ExitWaitTime)
	test.Assert(t, iSvr2.opt.RemoteOpt.MaxConnectionIdleTime == tmpConnIdleTime, iSvr2.opt.RemoteOpt.MaxConnectionIdleTime)
}

// TestMuxTransportOption tests the creation of a server,with netpollmux remote.ServerTransHandlerFactory option,
func TestMuxTransportOption(t *testing.T) {
	svr1 := NewServer()
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr1.Stop()
		test.Assert(t, err == nil, err)
	})
	err := svr1.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr1.Run()
	test.Assert(t, err == nil, err)
	iSvr1 := svr1.(*server)
	test.DeepEqual(t, iSvr1.opt.RemoteOpt.SvrHandlerFactory, detection.NewSvrTransHandlerFactory(netpoll.NewSvrTransHandlerFactory(), nphttp2.NewSvrTransHandlerFactory()))

	svr2 := NewServer(WithMuxTransport())
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr2.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr2.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr2.Run()
	test.Assert(t, err == nil, err)
	iSvr2 := svr2.(*server)
	handlerFactory := iSvr2.opt.RemoteOpt.SvrHandlerFactory
	test.DeepEqual(t, handlerFactory, netpollmux.NewSvrTransHandlerFactory())
}

// TestPayloadCodecOption tests the creation of a server with RemoteOpt.PayloadCodec option
func TestPayloadCodecOption(t *testing.T) {
	svr1 := NewServer()
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr1.Stop()
		test.Assert(t, err == nil, err)
	})
	err := svr1.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr1.Run()
	test.Assert(t, err == nil, err)
	iSvr1 := svr1.(*server)
	test.Assert(t, iSvr1.opt.RemoteOpt.PayloadCodec == nil)

	svr2 := NewServer(WithPayloadCodec(protobuf.NewProtobufCodec()))
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr2.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr2.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr2.Run()
	test.Assert(t, err == nil, err)
	iSvr2 := svr2.(*server)
	test.Assert(t, iSvr2.opt.RemoteOpt.PayloadCodec != nil)
	test.DeepEqual(t, iSvr2.opt.RemoteOpt.PayloadCodec, protobuf.NewProtobufCodec())
}

// TestRemoteOptGRPCCfgUintValueOption tests the creation of a server with RemoteOpt.GRPCCfg option
// and the value of the option is uint32
func TestRemoteOptGRPCCfgUintValueOption(t *testing.T) {
	// random value between 1 and 100
	randUint1 := uint32(rand.Int31n(100)) + 1
	randUint2 := uint32(rand.Int31n(100)) + 1
	randUint3 := uint32(rand.Int31n(100)) + 1
	randUint4 := uint32(rand.Int31n(100)) + 1

	svr1 := NewServer(
		WithGRPCInitialWindowSize(randUint1),
		WithGRPCInitialConnWindowSize(randUint2),
		WithGRPCMaxConcurrentStreams(randUint3),
		WithGRPCMaxHeaderListSize(randUint4),
	)

	time.AfterFunc(100*time.Millisecond, func() {
		err := svr1.Stop()
		test.Assert(t, err == nil, err)
	})
	err := svr1.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr1.Run()
	test.Assert(t, err == nil, err)
	iSvr1 := svr1.(*server)

	test.Assert(t, iSvr1.opt.RemoteOpt.GRPCCfg.InitialWindowSize == randUint1)
	test.Assert(t, iSvr1.opt.RemoteOpt.GRPCCfg.InitialConnWindowSize == randUint2)
	test.Assert(t, iSvr1.opt.RemoteOpt.GRPCCfg.MaxStreams == randUint3)
	test.Assert(t, *iSvr1.opt.RemoteOpt.GRPCCfg.MaxHeaderListSize == randUint4)
}

// TestGRPCKeepaliveEnforcementPolicyOption tests the creation of a server with RemoteOpt.GRPCCfg.KeepaliveEnforcementPolicy option
func TestGRPCKeepaliveEnforcementPolicyOption(t *testing.T) {
	// random value between 1 and 100
	randInt := rand.Int31n(100) + 1
	kep := grpc.EnforcementPolicy{
		MinTime:             time.Duration(randInt) * time.Second,
		PermitWithoutStream: true,
	}
	svr1 := NewServer(
		WithGRPCKeepaliveEnforcementPolicy(kep),
	)

	time.AfterFunc(100*time.Millisecond, func() {
		err := svr1.Stop()
		test.Assert(t, err == nil, err)
	})
	err := svr1.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr1.Run()
	test.Assert(t, err == nil, err)
	iSvr1 := svr1.(*server)
	test.DeepEqual(t, iSvr1.opt.RemoteOpt.GRPCCfg.KeepaliveEnforcementPolicy, kep)
}

// TestGRPCKeepaliveParamsOption tests the creation of a server with RemoteOpt.GRPCCfg.KeepaliveParams option
func TestGRPCKeepaliveParamsOption(t *testing.T) {
	randTimeDuration1 := time.Duration(rand.Int31n(100)+1) * time.Second
	randTimeDuration2 := time.Duration(rand.Int31n(100)+1) * time.Second
	randTimeDuration3 := time.Duration(rand.Int31n(100)+1) * time.Second
	randTimeDuration4 := time.Duration(rand.Int31n(10)+1) * time.Hour
	randTimeDuration5 := time.Duration(rand.Int31n(100)+1) * time.Second

	kp := grpc.ServerKeepalive{
		MaxConnectionIdle:     randTimeDuration1,
		MaxConnectionAge:      randTimeDuration2,
		MaxConnectionAgeGrace: randTimeDuration3,
		Time:                  randTimeDuration4,
		Timeout:               randTimeDuration5,
	}
	svr1 := NewServer(
		WithGRPCKeepaliveParams(kp),
	)

	time.AfterFunc(100*time.Millisecond, func() {
		err := svr1.Stop()
		test.Assert(t, err == nil, err)
	})
	err := svr1.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	err = svr1.Run()
	test.Assert(t, err == nil, err)
	iSvr1 := svr1.(*server)
	test.DeepEqual(t, iSvr1.opt.RemoteOpt.GRPCCfg.KeepaliveParams, kp)
}

func TestWithProfilerMessageTagging(t *testing.T) {
	var msgTagging1 remote.MessageTagging = func(ctx context.Context, msg remote.Message) (context.Context, []string) {
		return context.WithValue(ctx, "ctx1", 1), []string{"a", "1", "b", "1"}
	}
	var msgTagging2 remote.MessageTagging = func(ctx context.Context, msg remote.Message) (context.Context, []string) {
		return context.WithValue(ctx, "ctx2", 2), []string{"b", "2", "c", "2"}
	}
	svr := NewServer(WithProfilerMessageTagging(msgTagging1), WithProfilerMessageTagging(msgTagging2))
	err := svr.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr.Run()
	test.Assert(t, err == nil, err)
	iSvr := svr.(*server)
	from := rpcinfo.NewEndpointInfo("caller", "", nil, nil)
	to := rpcinfo.NewEndpointInfo("callee", "method", nil, nil)
	ri := rpcinfo.NewRPCInfo(from, to, nil, nil, nil)
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	svcInfo := mocks.ServiceInfo()
	svcSearchMap := map[string]*serviceinfo.ServiceInfo{
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockMethod):          svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockExceptionMethod): svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockErrorMethod):     svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockOnewayMethod):    svcInfo,
		mocks.MockMethod:          svcInfo,
		mocks.MockExceptionMethod: svcInfo,
		mocks.MockErrorMethod:     svcInfo,
		mocks.MockOnewayMethod:    svcInfo,
	}
	msg := remote.NewMessageWithNewer(svcInfo, svcSearchMap, ri, remote.Call, remote.Server, false)

	newCtx, tags := iSvr.opt.RemoteOpt.ProfilerMessageTagging(ctx, msg)
	test.Assert(t, len(tags) == 8)
	test.DeepEqual(t, tags, []string{"b", "2", "c", "2", "a", "1", "b", "1"})
	test.Assert(t, newCtx.Value("ctx1").(int) == 1)
	test.Assert(t, newCtx.Value("ctx2").(int) == 2)
}

func TestRefuseTrafficWithoutServiceNamOption(t *testing.T) {
	svr := NewServer(WithRefuseTrafficWithoutServiceName())
	err := svr.RegisterService(mocks.ServiceInfo(), mocks.MyServiceHandler())
	test.Assert(t, err == nil, err)
	time.AfterFunc(100*time.Millisecond, func() {
		err := svr.Stop()
		test.Assert(t, err == nil, err)
	})
	err = svr.Run()
	test.Assert(t, err == nil, err)
	iSvr := svr.(*server)
	test.Assert(t, iSvr.opt.RefuseTrafficWithoutServiceName)
}