ceremonyclient/go-libp2p/p2p/transport/quicreuse/reuse_test.go
Cassie Heart ebb31fd792 v1.4.19
2024-06-08 11:32:45 +00:00

241 lines
6.5 KiB
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, 1, conn.GetCount())
}
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, 1, tr.GetCount())
}
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, 1, conn.GetCount())
laddr := conn.LocalAddr().(*net.UDPAddr)
require.Equal(t, "0.0.0.0", laddr.IP.String())
require.NotEqual(t, 0, laddr.Port)
}
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, 1, lconn.GetCount())
// 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, 2, conn.GetCount())
}
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, 2, lconn.GetCount())
require.Equal(t, 2, tr.GetCount())
}
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, lTr, tr)
require.Equal(t, 2, tr.GetCount())
// 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, rTr, lconn2)
require.Equal(t, 2, lconn2.GetCount())
}
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, 1, lconn.GetCount())
// dial
conn, err := reuse.TransportForDial("udp4", raddr)
require.NoError(t, err)
require.Equal(t, 1, conn.GetCount())
}
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, 1, dTr.GetCount())
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, 1, lTr.GetCount())
closeTime := time.Now()
lTr.DecreaseCount()
dTr.DecreaseCount()
for {
num := numGlobals()
if closeTime.Add(maxUnusedDuration).Before(time.Now()) {
break
}
require.Equal(t, 2, num)
time.Sleep(2 * time.Millisecond)
}
require.Eventually(t, func() bool { return numGlobals() == 0 }, 4*garbageCollectInterval, 10*time.Millisecond)
}