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

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/cloudwego/kitex/pkg/connpool"

	"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 (
	mockDestService = "destService"
	mockAddr0       = "127.0.0.1:8000"
	mockAddr1       = "127.0.0.1:8001"
)

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

	var (
		minIdle        = 0
		maxIdle        = 1
		maxIdleTimeout = time.Millisecond
	)

	p := newPool(minIdle, maxIdle, maxIdleTimeout)
	count := make(map[*longConn]bool)

	conn := newLongConnForTest(ctrl, mockAddr0)
	recycled := p.Put(conn)
	test.Assert(t, recycled == true)
	test.Assert(t, p.Len() == 1)

	c, reused, decN := p.Get()
	test.Assert(t, c != nil)
	test.Assert(t, reused == true)
	test.Assert(t, decN == 1)
	count[c] = true
	test.Assert(t, len(count) == 1)
}

// TestPoolGetInactiveConn tests the pool with only inactive connection. Get() should return nil.
func TestPoolGetInactiveConn(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var (
		minIdle        = 0
		maxIdle        = 1
		maxIdleTimeout = time.Millisecond
	)

	p := newPool(minIdle, maxIdle, maxIdleTimeout)

	// inactive conn
	var closed bool
	c := mocksnetpoll.NewMockConnection(ctrl)
	c.EXPECT().IsActive().Return(false).AnyTimes()
	c.EXPECT().Close().DoAndReturn(func() error {
		closed = true
		return nil
	}).AnyTimes()
	conn := &longConn{
		Conn:    c,
		address: mockAddr0,
	}

	recycled := p.Put(conn)
	test.Assert(t, recycled == true)
	// inactive
	conn, reused, decNum := p.Get()
	test.Assert(t, decNum == 1)
	test.Assert(t, conn == nil)
	test.Assert(t, reused == false)
	test.Assert(t, closed == true)
}

// TestPoolGetWithInactiveConn tests the pool with inactive connection. Get() should return the first active one.
func TestPoolGetWithInactiveConn(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var (
		minIdle        = 0
		maxIdle        = 10
		maxIdleTimeout = time.Millisecond
		inactiveNum    = 5
	)

	p := newPool(minIdle, maxIdle, maxIdleTimeout)
	// put active conn
	activeConn := newLongConnForTest(ctrl, mockAddr0)
	recycled := p.Put(activeConn)
	test.Assert(t, recycled == true)

	// put inactive conn
	closed := make([]bool, inactiveNum)
	for i := 0; i < inactiveNum; i++ {
		idx := i
		c := mocksnetpoll.NewMockConnection(ctrl)
		c.EXPECT().IsActive().Return(false).AnyTimes()
		c.EXPECT().Close().DoAndReturn(func() error {
			closed[idx] = true
			return nil
		}).AnyTimes()
		conn := &longConn{
			Conn:    c,
			address: mockAddr0,
		}
		recycled := p.Put(conn)
		test.Assert(t, recycled == true)
	}
	test.Assert(t, p.Len() == inactiveNum+1)

	// decNum should be inactiveNum + 1
	conn, reused, decNum := p.Get()
	test.Assert(t, conn != nil)
	test.Assert(t, reused == true)
	test.Assert(t, decNum == inactiveNum+1)
	// check if all inactive conns have been closed
	for i := 0; i < inactiveNum; i++ {
		test.Assert(t, closed[i] == true)
	}
}

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

	var (
		minIdle        = 0
		maxIdle        = 2
		maxIdleTimeout = time.Millisecond
	)

	p := newPool(minIdle, maxIdle, maxIdleTimeout)
	for i := 0; i < maxIdle+1; i++ {
		recycled := p.Put(newLongConnForTest(ctrl, mockAddr0))
		if i < maxIdle {
			test.Assert(t, recycled == true)
		} else {
			test.Assert(t, recycled == false)
		}
	}
	test.Assert(t, p.Len() == maxIdle)
}

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

	var (
		minIdle        = 1
		maxIdle        = 10
		maxIdleTimeout = time.Millisecond
	)

	p := newPool(minIdle, maxIdle, maxIdleTimeout)
	for i := 0; i < maxIdle+1; i++ {
		p.Put(newLongConnForTest(ctrl, mockAddr0))
	}
	test.Assert(t, p.Len() == maxIdle)

	time.Sleep(maxIdleTimeout)
	p.Evict()
	test.Assert(t, p.Len() == minIdle)
}

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

	var (
		minIdle        = 0
		maxIdle        = 2
		maxIdleTimeout = time.Millisecond
	)

	p := newPool(minIdle, maxIdle, maxIdleTimeout)
	for i := 0; i < maxIdle+1; i++ {
		p.Put(newLongConnForTest(ctrl, mockAddr0))
	}

	n := p.Close()
	test.Assert(t, n == maxIdle)
	test.Assert(t, p.idleList == nil)
	test.Assert(t, p.Len() == 0)
}

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

	var (
		minIdle        = 0
		maxIdle        = 2
		maxIdleTimeout = time.Millisecond
	)

	p := newPool(minIdle, maxIdle, maxIdleTimeout)
	for i := 0; i < maxIdle+1; i++ {
		p.Put(newLongConnForTest(ctrl, mockAddr0))
	}

	dump := p.Dump()
	test.Assert(t, dump.IdleNum == 2)
	test.Assert(t, len(dump.ConnsDeadline) == 2)
}

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

	p := newLongPoolForTest(0, 2, 3, time.Second)
	defer p.Close()

	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().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 TestLongConnPoolReuse(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockReporter := newMockConnReporter()
	SetReporter(mockReporter)
	defer SetReporter(&DummyReporter{}) // reset the reporter to default

	idleTime := time.Millisecond * 100
	p := newLongPoolForTest(0, 2, 3, idleTime)
	defer p.Close()
	p.EnableReporter()

	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().AnyTimes()
		return conn, nil
	}).AnyTimes()

	addr1, addr2 := mockAddr1, "127.0.0.1:8002"
	netAddr1 := utils.NewNetAddr("tcp", addr1)
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

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

	count = make(map[net.Conn]int)
	for i := 0; i < getCnt; 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]++
	}
	// the first Get dial one connection, the following Gets reuse it.
	test.Assert(t, mockReporter.getConnSucceed(mockDestService, netAddr1) == 1)
	test.Assert(t, mockReporter.getReuseSucceed(mockDestService, netAddr1) == getCnt-1)
	test.Assert(t, len(count) == 1)

	count = make(map[net.Conn]int)
	for i := 0; i < getCnt; 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)

	// test exceed idleTime
	mockReporter.reset()
	count = make(map[net.Conn]int)
	getCnt = 3
	for i := 0; i < getCnt; i++ {
		time.Sleep(idleTime * 3)
		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]++
	}
	// Every Get needs dial one connection because old connections cannot be used
	test.Assert(t, mockReporter.getConnSucceed(mockDestService, netAddr1) == getCnt)
	test.Assert(t, len(count) == 3)
}

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

	p := newLongPoolForTest(0, 2, 5, time.Second)
	defer p.Close()

	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().AnyTimes()
		return conn, nil
	}).AnyTimes()

	addr := mockAddr1
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

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

	for _, c := range conns {
		err := p.Put(c)
		test.Assert(t, err == nil)
	}

	for i := 0; i < 10; i++ {
		c, err := p.Get(context.TODO(), "tcp", addr, opt)
		test.Assert(t, err == nil)
		count[c]++
	}
	test.Assert(t, len(count) == 18)
}

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

	p := newLongPoolForTest(0, 2, 3, time.Second)
	defer p.Close()

	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().AnyTimes()
		return conn, nil
	}).AnyTimes()

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

	var conns []net.Conn
	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]++
		conns = append(conns, c)

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

	for _, c := range conns {
		err := p.Put(c)
		test.Assert(t, err == nil)
	}

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

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

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

	p := newLongPoolForTest(0, 2, 5, time.Second)
	defer p.Close()

	var closed bool
	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().DoAndReturn(func() error {
			if closed {
				return errors.New("connection already closed")
			}
			closed = true
			return nil
		}).AnyTimes()
		return conn, nil
	}).AnyTimes()

	addr := mockAddr1
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

	c, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	err = p.Put(c)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	c2, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	test.Assert(t, c == c2)
	test.Assert(t, !closed)

	err = p.Discard(c2)
	test.Assert(t, err == nil)
	test.Assert(t, closed)

	c3, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	test.Assert(t, c3 != c2)
}

func TestLongConnPoolCloseOnError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	p := newLongPoolForTest(0, 2, 5, time.Second)
	defer p.Close()

	var closed, read bool

	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().DoAndReturn(func() error {
			if closed {
				return errors.New("connection already closed")
			}
			closed = true
			return nil
		}).AnyTimes()
		conn.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
			// Error on second time reading
			if read {
				return 0, errors.New("read timeout")
			}
			read = true
			return 0, nil
		}).AnyTimes()
		return conn, nil
	}).AnyTimes()

	addr := mockAddr1
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

	c, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	var buf []byte
	_, err = c.Read(buf)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	err = p.Put(c)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	c2, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	test.Assert(t, c == c2)
	test.Assert(t, !closed)

	_, err = c.Read(buf)
	test.Assert(t, err != nil)
	test.Assert(t, !closed)

	c3, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	test.Assert(t, c2 != c3)
}

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

	idleTime := time.Second
	p := newLongPoolForTest(0, 2, 5, idleTime)
	defer p.Close()

	var closed bool
	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().DoAndReturn(func() error {
			if closed {
				return errors.New("connection already closed")
			}
			closed = true
			return nil
		}).AnyTimes()
		return conn, nil
	}).AnyTimes()

	addr := mockAddr1
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

	c, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	err = p.Put(c)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	time.Sleep(idleTime * 3)

	c2, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	test.Assert(t, c != c2)
	test.Assert(t, closed) // the first connection should be closed
}

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

	p := newLongPoolForTest(0, 2, 5, time.Second)
	defer p.Close()

	var closed bool
	var mu sync.RWMutex // fix data race in test
	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().DoAndReturn(func() error {
			mu.RLock()
			if closed {
				mu.RUnlock()
				return errors.New("connection already closed")
			}
			mu.RUnlock()
			mu.Lock()
			closed = true
			mu.Unlock()
			return nil
		}).AnyTimes()
		return conn, nil
	}).AnyTimes()

	addr := mockAddr1
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}

	c, err := p.Get(context.TODO(), "tcp", addr, opt)
	test.Assert(t, err == nil)
	mu.RLock()
	test.Assert(t, !closed)
	mu.RUnlock()

	err = p.Put(c)
	test.Assert(t, err == nil)
	mu.RLock()
	test.Assert(t, !closed)
	mu.RUnlock()

	p.Clean("tcp", addr)
	time.Sleep(100 * time.Millisecond) // Wait for the clean goroutine to finish
	mu.RLock()
	test.Assert(t, closed)
	mu.RUnlock()
}

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

	p := newLongPoolForTest(0, 2, 5, time.Second)
	defer p.Close()

	var closed bool
	conn := mocksnetpoll.NewMockConnection(ctrl)
	conn.EXPECT().IsActive().Return(true).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) {
		na := utils.NewNetAddr(network, address)
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().DoAndReturn(func() error {
			if closed {
				return errors.New("connection already closed")
			}
			closed = true
			return nil
		}).MaxTimes(1)
		return conn, nil
	}).AnyTimes()

	network, address := "tcp", mockAddr1

	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}
	c, err := p.Get(context.TODO(), network, address, opt)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	lc, ok := c.(*longConn)
	test.Assert(t, ok)
	test.Assert(t, lc.Conn == conn)

	err = p.Discard(lc.Conn)
	test.Assert(t, err == nil)
	test.Assert(t, closed)
}

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

	p := newLongPoolForTest(0, 2, 3, time.Second)

	var closed int
	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().DoAndReturn(func() error {
			closed++
			return nil
		}).AnyTimes()
		return conn, nil
	}).AnyTimes()

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

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

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

	p.Close()
	test.Assert(t, closed == 2, closed)
	test.Assert(t, p.globalIdle.Now() == 0)

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

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

	var (
		poolNum         = 10
		idleTimeoutUnit = 111 * time.Millisecond
		pools           = make([]*LongPool, poolNum)
	)
	// add new pool with different idleTimeout, increasing the number of shared ticker
	for i := 0; i < poolNum; i++ {
		pools[i] = newLongPoolForTest(0, 2, 3, time.Duration(i+1)*idleTimeoutUnit)
	}
	// close
	for i := 0; i < poolNum; i++ {
		pools[i].Close()
		// should be removed from shardTickers
		_, ok := sharedTickers.Load(pools[i])
		test.Assert(t, !ok)
	}
}

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

	p := newLongPoolForTest(0, 2, 5, time.Second)
	defer p.Close()

	var closed bool
	conn := mocksnetpoll.NewMockConnection(ctrl)
	conn.EXPECT().IsActive().Return(true).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) {
		na := utils.NewNetAddr(network, address)
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().DoAndReturn(func() error {
			if closed {
				return errors.New("connection already closed")
			}
			closed = true
			return nil
		}).AnyTimes()
		return conn, nil
	}).AnyTimes()

	network, address := "tcp", mockAddr1
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}
	c, err := p.Get(context.TODO(), network, address, opt)
	test.Assert(t, err == nil)
	test.Assert(t, !closed)

	lc, ok := c.(*longConn)
	test.Assert(t, ok)
	test.Assert(t, lc.Conn == conn)

	err = p.Put(lc.Conn)
	test.Assert(t, err == nil)
	test.Assert(t, closed)
}

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

	var (
		idleTime = time.Millisecond * 100
		minIdle  = 1
		maxIdle  = 5
	)
	p := newLongPoolForTest(minIdle, maxIdle, maxIdle, idleTime)
	defer p.Close()

	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().DoAndReturn(func() error {
			return nil
		}).MaxTimes(1)
		return conn, nil
	}).AnyTimes()

	network, address := "tcp", mockAddr1
	opt := dialer.ConnOption{Dialer: d, ConnectTimeout: time.Second}
	// get a new conn
	var conns []net.Conn
	for i := 0; i < maxIdle; i++ {
		c, err := p.Get(context.TODO(), network, address, opt)
		test.Assert(t, err == nil)
		conns = append(conns, c)
	}
	for i := 0; i < maxIdle; i++ {
		// put conn back to the pool
		err := p.Put(conns[i])
		test.Assert(t, err == nil)
	}

	// only `minIdle` of connections should be kept in the pool
	// 3 times of idleTime to make sure the eviction goroutine can be done
	time.Sleep(3 * idleTime)
	p.peerMap.Range(func(key, value interface{}) bool {
		v := value.(*peer)
		test.Assert(t, v.Len() == minIdle)
		return true
	})
	// globalIdle should also be decreased when evicting
	test.Assert(t, int(p.globalIdle.Now()) == minIdle, p.globalIdle.Now())
	// get after eviction
	conns = []net.Conn{}
	for i := 0; i < maxIdle; i++ {
		c, err := p.Get(context.TODO(), network, address, opt)
		test.Assert(t, err == nil)
		conns = append(conns, c)
	}
	for i := 0; i < maxIdle; i++ {
		// put conn back to the pool
		err := p.Put(conns[i])
		test.Assert(t, err == nil)
	}
}

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

	idleTime := time.Second
	p := newLongPoolForTest(0, 2, 3, idleTime)
	defer p.Close()

	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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().AnyTimes()
		return conn, nil
	}).AnyTimes()

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

	// put conn back to the pool
	err = p.Put(conn)
	test.Assert(t, err == nil, err)

	// test Dump() to get conn pool info
	data := p.Dump().(map[string]interface{})
	val := data[mockAddr0]
	test.Assert(t, val != nil)

	length := len(val.(PoolDump).ConnsDeadline)
	test.Assert(t, length == 1)
}

func BenchmarkLongPoolGetOne(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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().AnyTimes()
		return conn, nil
	}).AnyTimes()

	p := newLongPoolForTest(0, 100000, 1<<20, time.Second)
	defer p.Close()
	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 BenchmarkLongPoolGetRand2000(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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().AnyTimes()
		return conn, nil
	}).AnyTimes()
	p := newLongPoolForTest(0, 50, 50, time.Second)
	defer p.Close()
	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 BenchmarkLongPoolGetRand2000Mesh(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().IsActive().Return(true).AnyTimes()
		conn.EXPECT().RemoteAddr().Return(na).AnyTimes()
		conn.EXPECT().Close().AnyTimes()
		return conn, nil
	}).AnyTimes()
	p := newLongPoolForTest(0, 100000, 1<<20, time.Second)
	defer p.Close()
	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)
		}
	})
}

// fakeNewLongPool creates a LongPool object and modifies it to fit tests.
func newLongPoolForTest(minPeerAddr, maxPeerAddr, global int, timeout time.Duration) *LongPool {
	cfg := connpool.IdleConfig{
		MinIdlePerAddress: minPeerAddr,
		MaxIdlePerAddress: maxPeerAddr,
		MaxIdleGlobal:     global,
		MaxIdleTimeout:    timeout,
	}
	return NewLongPool(mockDestService, cfg)
}

func newLongConnForTest(ctrl *gomock.Controller, addr string) *longConn {
	conn := mocksnetpoll.NewMockConnection(ctrl)
	conn.EXPECT().IsActive().Return(true).AnyTimes()
	conn.EXPECT().Close().AnyTimes()

	return &longConn{
		Conn:    conn,
		address: addr,
	}
}

type reportStat struct {
	reuseSucceed int
	connSucceed  int
	connFailed   int
}

func newMockConnReporter() *mockConnReporter {
	return &mockConnReporter{stats: make(map[string]*reportStat)}
}

type mockConnReporter struct {
	stats map[string]*reportStat // serviceName:addr -> reportStat
	sync.RWMutex
}

var _ Reporter = &mockConnReporter{}

func (m *mockConnReporter) ConnSucceed(poolType ConnectionPoolType, serviceName string, addr net.Addr) {
	m.Lock()
	defer m.Unlock()
	key := getKey(serviceName, addr)
	s, ok := m.stats[key]
	if !ok {
		s = &reportStat{}
	}
	s.connSucceed++
	m.stats[key] = s
}

func (m *mockConnReporter) ConnFailed(poolType ConnectionPoolType, serviceName string, addr net.Addr) {
	m.Lock()
	defer m.Unlock()
	key := getKey(serviceName, addr)
	s, ok := m.stats[key]
	if !ok {
		s = &reportStat{}
	}
	s.connFailed++
	m.stats[key] = s
}

func (m *mockConnReporter) ReuseSucceed(poolType ConnectionPoolType, serviceName string, addr net.Addr) {
	m.Lock()
	defer m.Unlock()
	key := getKey(serviceName, addr)
	s, ok := m.stats[key]
	if !ok {
		s = &reportStat{}
	}
	s.reuseSucceed++
	m.stats[key] = s
}

func (m *mockConnReporter) getConnSucceed(serviceName string, addr net.Addr) int {
	m.RLock()
	defer m.RUnlock()
	key := getKey(serviceName, addr)
	s, ok := m.stats[key]
	if !ok {
		return 0
	}
	return s.connSucceed
}

func (m *mockConnReporter) getReuseSucceed(serviceName string, addr net.Addr) int {
	m.RLock()
	defer m.RUnlock()
	key := getKey(serviceName, addr)
	s, ok := m.stats[key]
	if !ok {
		return 0
	}
	return s.reuseSucceed
}

func (m *mockConnReporter) reset() {
	m.Lock()
	defer m.Unlock()
	m.stats = make(map[string]*reportStat)
}

func getKey(serviceName string, addr net.Addr) string {
	if addr != nil {
		return fmt.Sprintf("%s:%s", serviceName, addr.String())
	}
	return serviceName
}