github.com/cloudwego/kitex@v0.9.0/pkg/remote/remotecli/conn_wrapper_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 remotecli

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/golang/mock/gomock"

	mocksnetpoll "github.com/cloudwego/kitex/internal/mocks/netpoll"
	mocksremote "github.com/cloudwego/kitex/internal/mocks/remote"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"

	"github.com/cloudwego/kitex/internal/test"
	connpool2 "github.com/cloudwego/kitex/pkg/connpool"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/remote/connpool"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

var poolCfg = connpool2.IdleConfig{MaxIdlePerAddress: 100, MaxIdleGlobal: 100, MaxIdleTimeout: time.Second}

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

	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
	conf := rpcinfo.NewRPCConfig()
	ri := rpcinfo.NewRPCInfo(nil, to, rpcinfo.NewInvocation("", ""), conf, rpcinfo.NewRPCStats())
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)

	connW := NewConnWrapper(connpool.NewLongPool("destService", poolCfg))
	_, err := connW.GetConn(ctx, mocksremote.NewMockDialer(ctrl), ri)
	test.Assert(t, err != nil)
	test.Assert(t, errors.Is(err, kerrors.ErrNoDestAddress))
}

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

	addr := utils.NewNetAddr("tcp", "to")
	conn := mocksnetpoll.NewMockConnection(ctrl)
	dialer := mocksremote.NewMockDialer(ctrl)
	dialer.EXPECT().DialTimeout(addr.Network(), addr.String(), gomock.Any()).Return(conn, nil)
	from := rpcinfo.NewEndpointInfo("from", "method", nil, nil)
	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
	to.SetInstance(discovery.NewInstance(addr.Network(), addr.String(), discovery.DefaultWeight, nil))
	conf := rpcinfo.NewRPCConfig()
	ri := rpcinfo.NewRPCInfo(from, to, rpcinfo.NewInvocation("", ""), conf, rpcinfo.NewRPCStats())

	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	connW := NewConnWrapper(nil)
	conn2, err := connW.GetConn(ctx, dialer, ri)
	test.Assert(t, err == nil, err)
	test.Assert(t, conn == conn2)
}

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

	addr := utils.NewNetAddr("tcp", "to")
	conn := mocksnetpoll.NewMockConnection(ctrl)
	conn.EXPECT().RemoteAddr().Return(addr).AnyTimes()
	conn.EXPECT().IsActive().Return(true).AnyTimes()
	conn.EXPECT().Close().AnyTimes()
	dialer := mocksremote.NewMockDialer(ctrl)
	dialer.EXPECT().DialTimeout(addr.Network(), addr.String(), gomock.Any()).Return(conn, nil).Times(1)
	from := rpcinfo.NewEndpointInfo("from", "method", nil, nil)
	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
	to.SetInstance(discovery.NewInstance(addr.Network(), addr.String(), discovery.DefaultWeight, nil))

	conf := rpcinfo.NewRPCConfig()
	ri := rpcinfo.NewRPCInfo(from, to, rpcinfo.NewInvocation("", ""), conf, rpcinfo.NewRPCStats())
	connPool := connpool.NewLongPool("destService", poolCfg)
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	// 释放连接, 连接复用
	for i := 0; i < 10; i++ {
		connW := NewConnWrapper(connPool)
		conn2, err := connW.GetConn(ctx, dialer, ri)
		test.Assert(t, err == nil, err)
		test.Assert(t, conn == conn2)
		connW.ReleaseConn(nil, ri)

	}

	dialer.EXPECT().DialTimeout(addr.Network(), addr.String(), gomock.Any()).Return(conn, nil).Times(10)
	connPool.Clean(addr.Network(), addr.String())
	// 未释放连接, 连接重建
	for i := 0; i < 10; i++ {
		connW := NewConnWrapper(connPool)
		conn2, err := connW.GetConn(ctx, dialer, ri)
		test.Assert(t, err == nil, err)
		test.Assert(t, conn == conn2)
	}
}

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

	addr := utils.NewNetAddr("tcp", "to")
	conn := mocksnetpoll.NewMockConnection(ctrl)
	conn.EXPECT().RemoteAddr().Return(addr).AnyTimes()
	conn.EXPECT().IsActive().Return(true).AnyTimes()
	conn.EXPECT().Close().AnyTimes()

	shortConnDialer := mocksremote.NewMockDialer(ctrl)
	shortConnDialer.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).Return(conn, nil).Times(b.N)

	longConnDialer := mocksremote.NewMockDialer(ctrl)
	longConnDialer.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).Return(conn, nil).Times(1)

	from := rpcinfo.NewEndpointInfo("from", "method", nil, nil)
	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
	to.SetInstance(discovery.NewInstance(addr.Network(), addr.String(), discovery.DefaultWeight, nil))

	conf := rpcinfo.NewRPCConfig()
	ri := rpcinfo.NewRPCInfo(from, to, rpcinfo.NewInvocation("", ""), conf, rpcinfo.NewRPCStats())
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	connPool := connpool.NewLongPool("destService", poolCfg)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		connW := NewConnWrapper(connPool)
		conn2, err := connW.GetConn(ctx, longConnDialer, ri)
		test.Assert(b, err == nil, err)
		test.Assert(b, conn == conn2)
		connW.ReleaseConn(nil, ri)

		connW2 := NewConnWrapper(nil)
		conn2, err = connW2.GetConn(ctx, shortConnDialer, ri)
		test.Assert(b, err == nil, err)
		test.Assert(b, conn == conn2)
		connW.ReleaseConn(nil, ri)
	}
}

// TestReleaseConnUseNilConn test release conn without before GetConn
func TestReleaseConnUseNilConn(t *testing.T) {
	addr := utils.NewNetAddr("tcp", "to")
	ri := newMockRPCInfo(addr)
	connPool := connpool.NewLongPool("destService", poolCfg)
	connW := NewConnWrapper(connPool)

	connW.ReleaseConn(nil, ri)
}

// TestReleaseConn test release conn
func TestReleaseConn(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	addr := utils.NewNetAddr("tcp", "to")
	ri := newMockRPCInfo(addr)

	conn := mocksnetpoll.NewMockConnection(ctrl)
	conn.EXPECT().RemoteAddr().Return(addr).AnyTimes()
	conn.EXPECT().Close().Return(nil).MinTimes(1)

	dialer := mocksremote.NewMockDialer(ctrl)
	dialer.EXPECT().DialTimeout(addr.Network(), addr.String(), gomock.Any()).Return(conn, nil).AnyTimes()

	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)

	connW := NewConnWrapper(nil)
	conn2, err := connW.GetConn(ctx, dialer, ri)
	test.Assert(t, err == nil, err)
	test.Assert(t, conn == conn2)
	connW.ReleaseConn(nil, ri)
}