package ttransport import ( "bytes" "context" "fmt" "io" "os" "runtime" "strconv" "sync" "testing" "time" crand "crypto/rand" mrand "math/rand" "github.com/libp2p/go-libp2p-testing/race" "github.com/libp2p/go-libp2p/core/network" "github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/core/transport" ma "github.com/multiformats/go-multiaddr" ) var randomness []byte var StressTestTimeout = 1 * time.Minute func init() { // read 1MB of randomness randomness = make([]byte, 1<<20) if _, err := crand.Read(randomness); err != nil { panic(err) } if timeout := os.Getenv("TEST_STRESS_TIMEOUT_MS"); timeout != "" { if v, err := strconv.ParseInt(timeout, 10, 32); err == nil { StressTestTimeout = time.Duration(v) * time.Millisecond } } } type Options struct { ConnNum int StreamNum int MsgNum int MsgMin int MsgMax int } func fullClose(t *testing.T, s network.MuxedStream) { if err := s.CloseWrite(); err != nil { t.Error(err) s.Reset() return } b, err := io.ReadAll(s) if err != nil { t.Error(err) } if len(b) != 0 { t.Error("expected to be done reading") } if err := s.Close(); err != nil { t.Error(err) } } func randBuf(size int) []byte { n := len(randomness) - size if size < 1 { panic(fmt.Errorf("requested too large buffer (%d). max is %d", size, len(randomness))) } start := mrand.Intn(n) return randomness[start : start+size] } func echoStream(t *testing.T, s network.MuxedStream) { // echo everything if _, err := io.Copy(s, s); err != nil { t.Error(err) } } func echo(t *testing.T, c transport.CapableConn) { var wg sync.WaitGroup defer wg.Wait() for { str, err := c.AcceptStream() if err != nil { break } wg.Add(1) go func() { defer wg.Done() defer str.Close() echoStream(t, str) }() } } func serve(t *testing.T, l transport.Listener) { var wg sync.WaitGroup defer wg.Wait() for { c, err := l.Accept() if err != nil { return } defer c.Close() wg.Add(1) go func() { defer wg.Done() echo(t, c) }() } } func SubtestStress(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID, opt Options) { msgsize := 1 << 11 rateLimitN := 5000 // max of 5k funcs, because -race has 8k max. rateLimitChan := make(chan struct{}, rateLimitN) for i := 0; i < rateLimitN; i++ { rateLimitChan <- struct{}{} } rateLimit := func(f func()) { <-rateLimitChan f() rateLimitChan <- struct{}{} } writeStream := func(s network.MuxedStream, bufs chan<- []byte) { for i := 0; i < opt.MsgNum; i++ { buf := randBuf(msgsize) bufs <- buf if _, err := s.Write(buf); err != nil { t.Errorf("s.Write(buf): %s", err) return } } } readStream := func(s network.MuxedStream, bufs <-chan []byte) { buf2 := make([]byte, msgsize) i := 0 for buf1 := range bufs { i++ if _, err := io.ReadFull(s, buf2); err != nil { t.Errorf("io.ReadFull(s, buf2): %s", err) return } if !bytes.Equal(buf1, buf2) { t.Errorf("buffers not equal (%x != %x)", buf1[:3], buf2[:3]) return } } } openStreamAndRW := func(c network.MuxedConn) { s, err := c.OpenStream(context.Background()) if err != nil { t.Errorf("failed to create NewStream: %s", err) return } bufs := make(chan []byte, opt.MsgNum) go func() { writeStream(s, bufs) close(bufs) }() readStream(s, bufs) fullClose(t, s) } openConnAndRW := func() { var wg sync.WaitGroup defer wg.Wait() l, err := ta.Listen(maddr) if err != nil { t.Error(err) return } defer l.Close() wg.Add(1) go func() { defer wg.Done() serve(t, l) }() c, err := tb.Dial(context.Background(), l.Multiaddr(), peerA) if err != nil { t.Error(err) return } defer c.Close() // serve the outgoing conn, because some muxers assume // that we _always_ call serve. (this is an error?) wg.Add(1) go func() { defer wg.Done() echo(t, c) }() var openWg sync.WaitGroup for i := 0; i < opt.StreamNum; i++ { openWg.Add(1) go rateLimit(func() { defer openWg.Done() openStreamAndRW(c) }) } openWg.Wait() } var wg sync.WaitGroup defer wg.Wait() for i := 0; i < opt.ConnNum; i++ { wg.Add(1) go rateLimit(func() { defer wg.Done() openConnAndRW() }) } } func SubtestStreamOpenStress(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) { l, err := ta.Listen(maddr) if err != nil { t.Fatal(err) } defer l.Close() count := 10000 workers := 5 if race.WithRace() { // the race detector can only deal with 8128 simultaneous goroutines, so let's make sure we don't go overboard. count = 1000 } var ( connA, connB transport.CapableConn ) accepted := make(chan error, 1) go func() { var err error connA, err = l.Accept() accepted <- err }() connB, err = tb.Dial(context.Background(), l.Multiaddr(), peerA) if err != nil { t.Fatal(err) } err = <-accepted if err != nil { t.Fatal(err) } defer func() { if connA != nil { connA.Close() } if connB != nil { connB.Close() } }() var wg sync.WaitGroup wg.Add(1) go func() { defer wg.Done() for j := 0; j < workers; j++ { wg.Add(1) go func() { defer wg.Done() for i := 0; i < count; i++ { s, err := connA.OpenStream(context.Background()) if err != nil { t.Error(err) return } wg.Add(1) go func() { defer wg.Done() fullClose(t, s) }() } }() } }() wg.Add(1) go func() { defer wg.Done() for i := 0; i < count*workers; i++ { str, err := connB.AcceptStream() if err != nil { break } wg.Add(1) go func() { defer wg.Done() fullClose(t, str) }() } }() timeout := time.After(StressTestTimeout) done := make(chan struct{}) go func() { wg.Wait() close(done) }() select { case <-timeout: t.Fatal("timed out receiving streams") case <-done: } } func SubtestStreamReset(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) { var wg sync.WaitGroup defer wg.Wait() l, err := ta.Listen(maddr) if err != nil { t.Fatal(err) } defer l.Close() wg.Add(1) go func() { defer wg.Done() muxa, err := l.Accept() if err != nil { t.Error(err) return } defer muxa.Close() s, err := muxa.OpenStream(context.Background()) if err != nil { t.Error(err) return } defer s.Close() // Some transports won't open the stream until we write. That's // fine. _, _ = s.Write([]byte("foo")) time.Sleep(time.Millisecond * 50) _, err = s.Write([]byte("bar")) if err == nil { t.Error("should have failed to write") } }() muxb, err := tb.Dial(context.Background(), l.Multiaddr(), peerA) if err != nil { t.Fatal(err) } defer muxb.Close() str, err := muxb.AcceptStream() if err != nil { t.Error(err) return } str.Reset() } func SubtestStress1Conn1Stream1Msg(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) { SubtestStress(t, ta, tb, maddr, peerA, Options{ ConnNum: 1, StreamNum: 1, MsgNum: 1, MsgMax: 100, MsgMin: 100, }) } func SubtestStress1Conn1Stream100Msg(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) { SubtestStress(t, ta, tb, maddr, peerA, Options{ ConnNum: 1, StreamNum: 1, MsgNum: 100, MsgMax: 100, MsgMin: 100, }) } func SubtestStress1Conn100Stream100Msg(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) { SubtestStress(t, ta, tb, maddr, peerA, Options{ ConnNum: 1, StreamNum: 100, MsgNum: 100, MsgMax: 100, MsgMin: 100, }) } func SubtestStressManyConn10Stream50Msg(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) { connNum := 5 if runtime.GOOS == "linux" { // Linux can handle a higher number of conns here than other platforms in CI. // See https://github.com/libp2p/go-libp2p/issues/1498. connNum = 50 } SubtestStress(t, ta, tb, maddr, peerA, Options{ ConnNum: connNum, StreamNum: 10, MsgNum: 50, MsgMax: 100, MsgMin: 100, }) } func SubtestStress1Conn1000Stream10Msg(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) { SubtestStress(t, ta, tb, maddr, peerA, Options{ ConnNum: 1, StreamNum: 1000, MsgNum: 10, MsgMax: 100, MsgMin: 100, }) } func SubtestStress1Conn100Stream100Msg10MB(t *testing.T, ta, tb transport.Transport, maddr ma.Multiaddr, peerA peer.ID) { SubtestStress(t, ta, tb, maddr, peerA, Options{ ConnNum: 1, StreamNum: 100, MsgNum: 100, MsgMax: 10000, MsgMin: 1000, }) }