update test

xtaci · Sep 8, 2016 · c2aeb5b · c2aeb5b
1 parent f75ca5b
commit c2aeb5b
Showing 1 changed file with 16 additions and 3 deletions.
diff --git a/session_test.go b/session_test.go
@@ -488,7 +488,6 @@ func BenchmarkConnSmux(b *testing.B) {
 	bench(b, cs, ss)
 }
 
-/*
 func BenchmarkConnTCP(b *testing.B) {
 	cs, ss, err := getTCPConnectionPair()
 	if err != nil {
@@ -498,7 +497,6 @@ func BenchmarkConnTCP(b *testing.B) {
 	defer ss.Close()
 	bench(b, cs, ss)
 }
-*/
 
 func getSmuxStreamPair() (*Stream, *Stream, error) {
 	c1, c2, err := getTCPConnectionPair()
@@ -562,10 +560,25 @@ func getTCPConnectionPair() (net.Conn, net.Conn, error) {
 
 func bench(b *testing.B, rd io.Reader, wr io.Writer) {
 	buf := make([]byte, 128*1024)
+	buf2 := make([]byte, 128*1024)
 	b.SetBytes(128 * 1024)
 	b.ResetTimer()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		count := 0
+		for {
+			n, _ := rd.Read(buf2)
+			count += n
+			if count == 128*1024*b.N {
+				return
+			}
+		}
+	}()
 	for i := 0; i < b.N; i++ {
 		wr.Write(buf)
-		rd.Read(buf)
 	}
+	wg.Wait()
 }