ceremonyclient/go-libp2p/p2p/transport/quicreuse/reuse_test.go

package quicreuse

import (
	"bytes"
	"net"
	"os"
	"runtime/pprof"
	"strings"
	"testing"
	"time"

	"github.com/libp2p/go-netroute"
	"github.com/stretchr/testify/require"
)

func (c *refcountedTransport) GetCount() int {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	return c.refCount
}

func closeAllConns(reuse *reuse) {
	reuse.mutex.Lock()
	for _, tr := range reuse.globalListeners {
		for tr.GetCount() > 0 {
			tr.DecreaseCount()
		}
	}
	for _, tr := range reuse.globalDialers {
		for tr.GetCount() > 0 {
			tr.DecreaseCount()
		}
	}
	for _, trs := range reuse.unicast {
		for _, tr := range trs {
			for tr.GetCount() > 0 {
				tr.DecreaseCount()
			}
		}
	}
	reuse.mutex.Unlock()
}

func platformHasRoutingTables() bool {
	_, err := netroute.New()
	return err == nil
}

func isGarbageCollectorRunning() bool {
	var b bytes.Buffer
	pprof.Lookup("goroutine").WriteTo(&b, 1)
	return strings.Contains(b.String(), "quicreuse.(*reuse).gc")
}

func cleanup(t *testing.T, reuse *reuse) {
	t.Cleanup(func() {
		closeAllConns(reuse)
		reuse.Close()
		require.False(t, isGarbageCollectorRunning(), "reuse gc still running")
	})
}

func TestReuseListenOnAllIPv4(t *testing.T) {
	reuse := newReuse(nil, nil)
	require.Eventually(t, isGarbageCollectorRunning, 500*time.Millisecond, 50*time.Millisecond, "expected garbage collector to be running")
	cleanup(t, reuse)

	addr, err := net.ResolveUDPAddr("udp4", "0.0.0.0:0")
	require.NoError(t, err)
	conn, err := reuse.TransportForListen("udp4", addr)
	require.NoError(t, err)
	require.Equal(t, conn.GetCount(), 1)
}

func TestReuseListenOnAllIPv6(t *testing.T) {
	reuse := newReuse(nil, nil)
	require.Eventually(t, isGarbageCollectorRunning, 500*time.Millisecond, 50*time.Millisecond, "expected garbage collector to be running")
	cleanup(t, reuse)

	addr, err := net.ResolveUDPAddr("udp6", "[::]:1234")
	require.NoError(t, err)
	tr, err := reuse.TransportForListen("udp6", addr)
	require.NoError(t, err)
	defer tr.Close()
	require.Equal(t, tr.GetCount(), 1)
}

func TestReuseCreateNewGlobalConnOnDial(t *testing.T) {
	reuse := newReuse(nil, nil)
	cleanup(t, reuse)

	addr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
	require.NoError(t, err)
	conn, err := reuse.TransportForDial("udp4", addr)
	require.NoError(t, err)
	require.Equal(t, conn.GetCount(), 1)
	laddr := conn.LocalAddr().(*net.UDPAddr)
	require.Equal(t, laddr.IP.String(), "0.0.0.0")
	require.NotEqual(t, laddr.Port, 0)
}

func TestReuseConnectionWhenDialing(t *testing.T) {
	reuse := newReuse(nil, nil)
	cleanup(t, reuse)

	addr, err := net.ResolveUDPAddr("udp4", "0.0.0.0:0")
	require.NoError(t, err)
	lconn, err := reuse.TransportForListen("udp4", addr)
	require.NoError(t, err)
	require.Equal(t, lconn.GetCount(), 1)
	// dial
	raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
	require.NoError(t, err)
	conn, err := reuse.TransportForDial("udp4", raddr)
	require.NoError(t, err)
	require.Equal(t, conn.GetCount(), 2)
}

func TestReuseConnectionWhenListening(t *testing.T) {
	reuse := newReuse(nil, nil)
	cleanup(t, reuse)

	raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
	require.NoError(t, err)
	tr, err := reuse.TransportForDial("udp4", raddr)
	require.NoError(t, err)
	laddr := &net.UDPAddr{IP: net.IPv4zero, Port: tr.LocalAddr().(*net.UDPAddr).Port}
	lconn, err := reuse.TransportForListen("udp4", laddr)
	require.NoError(t, err)
	require.Equal(t, lconn.GetCount(), 2)
	require.Equal(t, tr.GetCount(), 2)
}

func TestReuseConnectionWhenDialBeforeListen(t *testing.T) {
	reuse := newReuse(nil, nil)
	cleanup(t, reuse)

	// dial any address
	raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
	require.NoError(t, err)
	rTr, err := reuse.TransportForDial("udp4", raddr)
	require.NoError(t, err)

	// open a listener
	laddr := &net.UDPAddr{IP: net.IPv4zero, Port: 1234}
	lTr, err := reuse.TransportForListen("udp4", laddr)
	require.NoError(t, err)

	// new dials should go via the listener connection
	raddr, err = net.ResolveUDPAddr("udp4", "1.1.1.1:1235")
	require.NoError(t, err)
	tr, err := reuse.TransportForDial("udp4", raddr)
	require.NoError(t, err)
	require.Equal(t, tr, lTr)
	require.Equal(t, tr.GetCount(), 2)

	// a listener on an unspecified port should reuse the dialer
	laddr2 := &net.UDPAddr{IP: net.IPv4zero, Port: 0}
	lconn2, err := reuse.TransportForListen("udp4", laddr2)
	require.NoError(t, err)
	require.Equal(t, lconn2, rTr)
	require.Equal(t, lconn2.GetCount(), 2)
}

func TestReuseListenOnSpecificInterface(t *testing.T) {
	if platformHasRoutingTables() {
		t.Skip("this test only works on platforms that support routing tables")
	}
	reuse := newReuse(nil, nil)
	cleanup(t, reuse)

	router, err := netroute.New()
	require.NoError(t, err)

	raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
	require.NoError(t, err)
	_, _, ip, err := router.Route(raddr.IP)
	require.NoError(t, err)
	// listen
	addr, err := net.ResolveUDPAddr("udp4", ip.String()+":0")
	require.NoError(t, err)
	lconn, err := reuse.TransportForListen("udp4", addr)
	require.NoError(t, err)
	require.Equal(t, lconn.GetCount(), 1)
	// dial
	conn, err := reuse.TransportForDial("udp4", raddr)
	require.NoError(t, err)
	require.Equal(t, conn.GetCount(), 1)
}

func TestReuseGarbageCollect(t *testing.T) {
	maxUnusedDurationOrig := maxUnusedDuration
	garbageCollectIntervalOrig := garbageCollectInterval
	t.Cleanup(func() {
		maxUnusedDuration = maxUnusedDurationOrig
		garbageCollectInterval = garbageCollectIntervalOrig
	})
	garbageCollectInterval = 50 * time.Millisecond
	maxUnusedDuration = 100 * time.Millisecond
	if os.Getenv("CI") != "" {
		// Increase these timeouts if in CI
		garbageCollectInterval = 10 * garbageCollectInterval
		maxUnusedDuration = 10 * maxUnusedDuration
	}

	reuse := newReuse(nil, nil)
	cleanup(t, reuse)

	numGlobals := func() int {
		reuse.mutex.Lock()
		defer reuse.mutex.Unlock()
		return len(reuse.globalListeners) + len(reuse.globalDialers)
	}

	raddr, err := net.ResolveUDPAddr("udp4", "1.2.3.4:1234")
	require.NoError(t, err)
	dTr, err := reuse.TransportForDial("udp4", raddr)
	require.NoError(t, err)
	require.Equal(t, dTr.GetCount(), 1)

	addr, err := net.ResolveUDPAddr("udp4", "0.0.0.0:1234")
	require.NoError(t, err)
	lTr, err := reuse.TransportForListen("udp4", addr)
	require.NoError(t, err)
	require.Equal(t, lTr.GetCount(), 1)

	closeTime := time.Now()
	lTr.DecreaseCount()
	dTr.DecreaseCount()

	for {
		num := numGlobals()
		if closeTime.Add(maxUnusedDuration).Before(time.Now()) {
			break
		}
		require.Equal(t, num, 2)
		time.Sleep(2 * time.Millisecond)
	}
	require.Eventually(t, func() bool { return numGlobals() == 0 }, 4*garbageCollectInterval, 10*time.Millisecond)
}