test: make the tests reliable

* fix races
* fix timeouts
* disable some tests when the race detector is enabled (too slow, too many goroutines)

bench_test.go

@@ -48,24 +48,25 @@ func BenchmarkAccept(b *testing.B) {
 func BenchmarkSendRecv(b *testing.B) {
 	client, server := testClientServer()
 	defer client.Close()
-	defer server.Close()
 
 	sendBuf := make([]byte, 512)
 	recvBuf := make([]byte, 512)
 
 	doneCh := make(chan struct{})
 	go func() {
+		defer close(doneCh)
+		defer server.Close()
 		stream, err := server.AcceptStream()
 		if err != nil {
 			return
 		}
 		defer stream.Close()
 		for i := 0; i < b.N; i++ {
 			if _, err := io.ReadFull(stream, recvBuf); err != nil {
-				b.Fatalf("err: %v", err)
+				b.Errorf("err: %v", err)
+				return
 			}
 		}
-		close(doneCh)
 	}()
 
 	stream, err := client.Open()
@@ -95,6 +96,8 @@ func BenchmarkSendRecvLarge(b *testing.B) {
 	recvDone := make(chan struct{})
 
 	go func() {
+		defer close(recvDone)
+		defer server.Close()
 		stream, err := server.AcceptStream()
 		if err != nil {
 			return
@@ -103,11 +106,11 @@ func BenchmarkSendRecvLarge(b *testing.B) {
 		for i := 0; i < b.N; i++ {
 			for j := 0; j < sendSize/recvSize; j++ {
 				if _, err := io.ReadFull(stream, recvBuf); err != nil {
-					b.Fatalf("err: %v", err)
+					b.Errorf("err: %v", err)
+					return
 				}
 			}
 		}
-		close(recvDone)
 	}()
 
 	stream, err := client.Open()

session_norace_test.go

@@ -0,0 +1,163 @@
+//+build !race
+
+package yamux
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestSession_PingOfDeath(t *testing.T) {
+	client, server := testClientServerConfig(testConfNoKeepAlive())
+	defer client.Close()
+	defer server.Close()
+
+	count := 10000
+
+	var wg sync.WaitGroup
+	begin := make(chan struct{})
+	for i := 0; i < count; i++ {
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			<-begin
+			if _, err := server.Ping(); err != nil {
+				t.Error(err)
+			}
+		}()
+		go func() {
+			defer wg.Done()
+			<-begin
+			if _, err := client.Ping(); err != nil {
+				t.Error(err)
+			}
+		}()
+	}
+	close(begin)
+	wg.Wait()
+}
+
+func TestSendData_VeryLarge(t *testing.T) {
+	client, server := testClientServer()
+	defer client.Close()
+	defer server.Close()
+
+	var n int64 = 1 * 1024 * 1024 * 1024
+	var workers int = 16
+
+	wg := &sync.WaitGroup{}
+	wg.Add(workers * 2)
+
+	for i := 0; i < workers; i++ {
+		go func() {
+			defer wg.Done()
+			stream, err := server.AcceptStream()
+			if err != nil {
+				t.Errorf("err: %v", err)
+				return
+			}
+			defer stream.Close()
+
+			buf := make([]byte, 4)
+			_, err = io.ReadFull(stream, buf)
+			if err != nil {
+				t.Errorf("err: %v", err)
+				return
+			}
+			if !bytes.Equal(buf, []byte{0, 1, 2, 3}) {
+				t.Errorf("bad header")
+				return
+			}
+
+			recv, err := io.Copy(ioutil.Discard, stream)
+			if err != nil {
+				t.Errorf("err: %v", err)
+				return
+			}
+			if recv != n {
+				t.Errorf("bad: %v", recv)
+				return
+			}
+		}()
+	}
+	for i := 0; i < workers; i++ {
+		go func() {
+			defer wg.Done()
+			stream, err := client.Open()
+			if err != nil {
+				t.Errorf("err: %v", err)
+				return
+			}
+			defer stream.Close()
+
+			_, err = stream.Write([]byte{0, 1, 2, 3})
+			if err != nil {
+				t.Errorf("err: %v", err)
+				return
+			}
+
+			unlimited := &UnlimitedReader{}
+			sent, err := io.Copy(stream, io.LimitReader(unlimited, n))
+			if err != nil {
+				t.Errorf("err: %v", err)
+				return
+			}
+			if sent != n {
+				t.Errorf("bad: %v", sent)
+				return
+			}
+		}()
+	}
+
+	doneCh := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(doneCh)
+	}()
+	select {
+	case <-doneCh:
+	case <-time.After(20 * time.Second):
+		server.Close()
+		client.Close()
+		wg.Wait()
+		t.Fatal("timeout")
+	}
+}
+
+func TestLargeWindow(t *testing.T) {
+	conf := DefaultConfig()
+	conf.MaxStreamWindowSize *= 2
+
+	client, server := testClientServerConfig(conf)
+	defer client.Close()
+	defer server.Close()
+
+	stream, err := client.Open()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer stream.Close()
+
+	stream2, err := server.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer stream2.Close()
+
+	err = stream.SetWriteDeadline(time.Now().Add(10 * time.Millisecond))
+	if err != nil {
+		t.Fatal(err)
+	}
+	buf := make([]byte, conf.MaxStreamWindowSize)
+	n, err := stream.Write(buf)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if n != len(buf) {
+		t.Fatalf("short write: %d", n)
+	}
+}