traefik/pkg/proxy/fast/connpool_test.go

package fast

import (
	"net"
	"runtime"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestConnPool_ConnReuse(t *testing.T) {
	testCases := []struct {
		desc     string
		poolFn   func(pool *connPool)
		expected int
	}{
		{
			desc: "One connection",
			poolFn: func(pool *connPool) {
				c1, _ := pool.AcquireConn()
				pool.ReleaseConn(c1)
			},
			expected: 1,
		},
		{
			desc: "Two connections with release",
			poolFn: func(pool *connPool) {
				c1, _ := pool.AcquireConn()
				pool.ReleaseConn(c1)

				c2, _ := pool.AcquireConn()
				pool.ReleaseConn(c2)
			},
			expected: 1,
		},
		{
			desc: "Two concurrent connections",
			poolFn: func(pool *connPool) {
				c1, _ := pool.AcquireConn()
				c2, _ := pool.AcquireConn()

				pool.ReleaseConn(c1)
				pool.ReleaseConn(c2)
			},
			expected: 2,
		},
	}

	for _, test := range testCases {
		t.Run(test.desc, func(t *testing.T) {
			t.Parallel()

			var connAlloc int
			dialer := func() (net.Conn, error) {
				connAlloc++
				return &net.TCPConn{}, nil
			}

			pool := newConnPool(2, 0, 0, dialer)
			test.poolFn(pool)

			assert.Equal(t, test.expected, connAlloc)
		})
	}
}

func TestConnPool_MaxIdleConn(t *testing.T) {
	testCases := []struct {
		desc        string
		poolFn      func(pool *connPool)
		maxIdleConn int
		expected    int
	}{
		{
			desc: "One connection",
			poolFn: func(pool *connPool) {
				c1, _ := pool.AcquireConn()
				pool.ReleaseConn(c1)
			},
			maxIdleConn: 1,
			expected:    1,
		},
		{
			desc: "Multiple connections with defered release",
			poolFn: func(pool *connPool) {
				for range 7 {
					c, _ := pool.AcquireConn()
					defer pool.ReleaseConn(c)
				}
			},
			maxIdleConn: 5,
			expected:    5,
		},
	}

	for _, test := range testCases {
		t.Run(test.desc, func(t *testing.T) {
			t.Parallel()

			var keepOpenedConn int
			dialer := func() (net.Conn, error) {
				keepOpenedConn++
				return &mockConn{
					doneCh: make(chan struct{}),
					closeFn: func() error {
						keepOpenedConn--
						return nil
					},
				}, nil
			}

			pool := newConnPool(test.maxIdleConn, 0, 0, dialer)
			test.poolFn(pool)

			assert.Equal(t, test.expected, keepOpenedConn)
		})
	}
}

func TestGC(t *testing.T) {
	// TODO: make the test stable if possible.
	t.Skip("This test is flaky")

	var isDestroyed bool
	pools := map[string]*connPool{}
	dialer := func() (net.Conn, error) {
		c := &mockConn{closeFn: func() error {
			return nil
		}}
		return c, nil
	}

	pools["test"] = newConnPool(10, 1*time.Second, 0, dialer)
	runtime.SetFinalizer(pools["test"], func(p *connPool) {
		isDestroyed = true
	})
	c, err := pools["test"].AcquireConn()
	require.NoError(t, err)

	pools["test"].ReleaseConn(c)

	pools["test"].Close()

	delete(pools, "test")

	runtime.GC()

	require.True(t, isDestroyed)
}

type mockConn struct {
	closeFn func() error
	doneCh  chan struct{} // makes sure that the readLoop is blocking avoiding close.
}

func (m *mockConn) Read(_ []byte) (n int, err error) {
	<-m.doneCh
	return 0, nil
}

func (m *mockConn) Write(_ []byte) (n int, err error) {
	panic("implement me")
}

func (m *mockConn) Close() error {
	defer close(m.doneCh)
	if m.closeFn != nil {
		return m.closeFn()
	}
	return nil
}

func (m *mockConn) LocalAddr() net.Addr {
	panic("implement me")
}

func (m *mockConn) RemoteAddr() net.Addr {
	panic("implement me")
}

func (m *mockConn) SetDeadline(_ time.Time) error {
	panic("implement me")
}

func (m *mockConn) SetReadDeadline(_ time.Time) error {
	panic("implement me")
}

func (m *mockConn) SetWriteDeadline(_ time.Time) error {
	panic("implement me")
}