github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpollmux/mux_pool_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 netpollmux

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"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/internal/test"
	dialer "github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/utils"
)

var (
	mockAddr0 = "127.0.0.1:8000"
	mockAddr1 = "127.0.0.1:8001"
)

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

	p := NewMuxConnPool(1)

	d := mocksremote.NewMockDialer(ctrl)
	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
		connectCost := time.Millisecond * 10
		if timeout < connectCost {
			return nil, errors.New("connect timeout")
		}
		na := utils.NewNetAddr(network, address)
		conn := mocksnetpoll.NewMockConnection(ctrl)
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().SetOnRequest(gomock.Any())
		conn.EXPECT().AddCloseCallback(gomock.Any())
		conn.EXPECT().IsActive().Return(false).AnyTimes()
		return conn, nil
	}).AnyTimes()

	var err error

	_, err = p.Get(context.TODO(), "tcp", mockAddr0, dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second})
	test.Assert(t, err == nil)

	_, err = p.Get(context.TODO(), "tcp", mockAddr0, dialer.ConnOption{Dialer: d, ConnectTimeout: time.Millisecond})
	test.Assert(t, err != nil)
}

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

	p := NewMuxConnPool(1)

	d := mocksremote.NewMockDialer(ctrl)
	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
		na := utils.NewNetAddr(network, address)
		conn := mocksnetpoll.NewMockConnection(ctrl)
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().SetOnRequest(gomock.Any())
		conn.EXPECT().AddCloseCallback(gomock.Any())
		conn.EXPECT().IsActive().Return(true).AnyTimes()
		return conn, nil
	}).AnyTimes()

	addr1, addr2 := mockAddr1, "127.0.0.1:8002"
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

	count := make(map[net.Conn]int)
	for i := 0; i < 10; i++ {
		c, err := p.Get(context.TODO(), "tcp", addr1, opt)
		test.Assert(t, err == nil)
		count[c]++
	}
	test.Assert(t, len(count) == 1)

	count = make(map[net.Conn]int)
	for i := 0; i < 10; i++ {
		c, err := p.Get(context.TODO(), "tcp", addr1, opt)
		test.Assert(t, err == nil)
		err = p.Put(c)
		test.Assert(t, err == nil)
		count[c]++
	}
	test.Assert(t, len(count) == 1)

	count = make(map[net.Conn]int)
	for i := 0; i < 10; i++ {
		c, err := p.Get(context.TODO(), "tcp", addr1, opt)
		test.Assert(t, err == nil)
		err = p.Put(c)
		test.Assert(t, err == nil)
		count[c]++

		c, err = p.Get(context.TODO(), "tcp", addr2, opt)
		test.Assert(t, err == nil)
		err = p.Put(c)
		test.Assert(t, err == nil)
		count[c]++
	}
	test.Assert(t, len(count) == 2)
}

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

	size := 4
	conn := mocksnetpoll.NewMockConnection(ctrl)
	d := mocksremote.NewMockDialer(ctrl)
	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
		conn.EXPECT().Close().Return(nil).Times(0)
		conn.EXPECT().IsActive().Return(true).AnyTimes()
		conn.EXPECT().SetOnRequest(gomock.Any())
		conn.EXPECT().AddCloseCallback(gomock.Any())
		return conn, nil
	}).AnyTimes()

	p := NewMuxConnPool(size)

	network, address := "tcp", mockAddr0
	var conns []net.Conn
	for i := 0; i < 1024; i++ {
		conn, err := p.Get(context.TODO(), network, address, dialer.ConnOption{Dialer: d})
		test.Assert(t, err == nil)
		conns = append(conns, conn)
	}
	for i := 0; i < 128; i++ {
		p.Discard(conns[i])
	}
	conn.EXPECT().Close().Return(nil).Times(size)
	p.Clean(network, address)
}

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

	conn := mocksnetpoll.NewMockConnection(ctrl)
	conn.EXPECT().Close().Return(nil).Times(2)
	conn.EXPECT().IsActive().Return(true).AnyTimes()
	conn.EXPECT().SetOnRequest(gomock.Any()).AnyTimes()
	conn.EXPECT().AddCloseCallback(gomock.Any()).AnyTimes()
	d := mocksremote.NewMockDialer(ctrl)
	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
		return conn, nil
	}).AnyTimes()

	p := NewMuxConnPool(1)

	network, address := "tcp", mockAddr0
	_, err := p.Get(context.TODO(), network, address, dialer.ConnOption{Dialer: d})
	test.Assert(t, err == nil)

	network, address = "tcp", mockAddr1
	_, err = p.Get(context.TODO(), network, address, dialer.ConnOption{Dialer: d})
	test.Assert(t, err == nil)

	connCount := 0
	p.connMap.Range(func(key, value interface{}) bool {
		connCount++
		return true
	})
	test.Assert(t, connCount == 2)

	p.Close()

	connCount = 0
	p.connMap.Range(func(key, value interface{}) bool {
		connCount++
		return true
	})
	test.Assert(t, connCount == 0)
}

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

	d := mocksremote.NewMockDialer(ctrl)
	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
		na := utils.NewNetAddr(network, address)
		conn := mocksnetpoll.NewMockConnection(ctrl)
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().SetOnRequest(gomock.Any()).AnyTimes()
		conn.EXPECT().AddCloseCallback(gomock.Any()).AnyTimes()
		conn.EXPECT().IsActive().Return(false).AnyTimes()
		conn.EXPECT().Close().Return(nil).AnyTimes()
		return conn, nil
	}).AnyTimes()
	p := NewMuxConnPool(1)
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

	b.ResetTimer()
	b.ReportAllocs()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			p.Get(context.TODO(), "tcp", mockAddr1, opt)
		}
	})
}

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

	d := mocksremote.NewMockDialer(ctrl)
	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
		na := utils.NewNetAddr(network, address)
		conn := mocksnetpoll.NewMockConnection(ctrl)
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().SetOnRequest(gomock.Any()).AnyTimes()
		conn.EXPECT().AddCloseCallback(gomock.Any()).AnyTimes()
		conn.EXPECT().IsActive().Return(false).AnyTimes()
		conn.EXPECT().Close().Return(nil).AnyTimes()
		return conn, nil
	}).AnyTimes()
	p := NewMuxConnPool(1)
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

	var addrs []string
	for i := 0; i < 2000; i++ {
		addrs = append(addrs, fmt.Sprintf("127.0.0.1:%d", 8000+rand.Intn(10000)))
	}
	b.ResetTimer()
	b.ReportAllocs()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			p.Get(context.TODO(), "tcp", addrs[rand.Intn(2000)], opt)
		}
	})
}

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

	d := mocksremote.NewMockDialer(ctrl)
	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
		na := utils.NewNetAddr(network, address)
		conn := mocksnetpoll.NewMockConnection(ctrl)
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().SetOnRequest(gomock.Any()).AnyTimes()
		conn.EXPECT().AddCloseCallback(gomock.Any()).AnyTimes()
		conn.EXPECT().IsActive().Return(false).AnyTimes()
		conn.EXPECT().Close().Return(nil).AnyTimes()
		return conn, nil
	}).AnyTimes()
	p := NewMuxConnPool(1)
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

	var addrs []string
	for i := 0; i < 2000; i++ {
		addrs = append(addrs, fmt.Sprintf("127.0.0.1:%d", 8000+rand.Intn(10000)))
	}
	b.ResetTimer()
	b.ReportAllocs()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			p.Get(context.TODO(), "tcp", addrs[rand.Intn(2000)], opt)
		}
	})
}