ceremonyclient/go-libp2p/p2p/test/reconnects/reconnect_test.go
2023-08-20 23:07:43 -05:00

120 lines
2.6 KiB
Go

package reconnect
import (
"context"
"io"
"math/rand"
"runtime"
"sync"
"testing"
"time"
"github.com/libp2p/go-libp2p/core/host"
"github.com/libp2p/go-libp2p/core/network"
"github.com/libp2p/go-libp2p/core/protocol"
bhost "github.com/libp2p/go-libp2p/p2p/host/basic"
swarmt "github.com/libp2p/go-libp2p/p2p/net/swarm/testing"
"github.com/stretchr/testify/require"
)
func EchoStreamHandler(stream network.Stream) {
_, err := io.CopyBuffer(stream, stream, make([]byte, 64)) // use a small buffer here to avoid problems with flow control
if err == nil {
stream.Close()
} else {
stream.Reset()
}
}
func TestReconnect5(t *testing.T) {
runTest := func(t *testing.T, swarmOpt swarmt.Option) {
t.Helper()
const num = 5
hosts := make([]host.Host, 0, num)
for i := 0; i < num; i++ {
h, err := bhost.NewHost(swarmt.GenSwarm(t, swarmOpt), nil)
require.NoError(t, err)
defer h.Close()
h.Start()
hosts = append(hosts, h)
h.SetStreamHandler(protocol.TestingID, EchoStreamHandler)
}
for i := 0; i < 4; i++ {
runRound(t, hosts)
}
}
t.Run("using TCP", func(t *testing.T) {
if runtime.GOOS == "darwin" {
t.Skip("TCP RST handling is flaky in OSX, see https://github.com/golang/go/issues/50254")
}
runTest(t, swarmt.OptDisableQUIC)
})
t.Run("using QUIC", func(t *testing.T) {
runTest(t, swarmt.OptDisableTCP)
})
}
func runRound(t *testing.T, hosts []host.Host) {
for _, h1 := range hosts {
h1.SetStreamHandler(protocol.TestingID, EchoStreamHandler)
for _, h2 := range hosts {
h1.Peerstore().AddAddrs(h2.ID(), h2.Addrs(), time.Hour)
}
}
const (
numStreams = 5
maxDataLen = 64 << 10
)
rnd := rand.New(rand.NewSource(12345))
// exchange some data
for _, h1 := range hosts {
for _, h2 := range hosts {
if h1 == h2 {
continue
}
var wg sync.WaitGroup
wg.Add(numStreams)
for i := 0; i < numStreams; i++ {
data := make([]byte, rand.Intn(maxDataLen)+1)
rnd.Read(data)
go func() {
defer wg.Done()
str, err := h1.NewStream(context.Background(), h2.ID(), protocol.TestingID)
require.NoError(t, err)
defer str.Close()
_, err = str.Write(data)
require.NoError(t, err)
}()
}
wg.Wait()
}
}
// disconnect all hosts
for _, h1 := range hosts {
// close connection
cs := h1.Network().Conns()
for _, c := range cs {
c.Close()
}
}
require.Eventually(t, func() bool {
for _, h1 := range hosts {
for _, h2 := range hosts {
if len(h1.Network().ConnsToPeer(h2.ID())) > 0 {
return false
}
}
}
return true
}, 5000*time.Millisecond, 10*time.Millisecond)
}