Squashed commit of the following:

xtaci · xtaci · commit fe8b9a22334d · 2019-12-28T22:16:15.000+08:00
commit 4d53a56 Author: xtaci <daniel820313@gmail.com> Date: Sat Dec 28 22:14:17 2019 +0800 add function PeekSize commit 9ca9c83 Author: xtaci <daniel820313@gmail.com> Date: Sat Dec 28 21:53:41 2019 +0800 add Peak
diff --git a/session_test.go b/session_test.go
@@ -160,7 +160,7 @@ func TestPoll(t *testing.T) {
 		}
 	}()
 
-	buf := make([]byte, 128)
+	buf := make([]byte, 65536)
 	events := make([]*Stream, 128)
 	for {
 		n, err := session.PollWait(events)
@@ -171,15 +171,20 @@ func TestPoll(t *testing.T) {
 		for i := 0; i < n; i++ {
 			stream := events[i]
 			for {
-				n, err := stream.TryRead(buf)
+				size := stream.PeekSize()
+				nr, err := stream.TryRead(buf)
 				if err == ErrWouldBlock {
 					break
 				}
 
+				if size != nr {
+					t.Fatal("incorrect peak")
+				}
+
 				if err != nil {
 					t.Fatal(err)
 				}
-				received += n
+				received += nr
 				if received == len(tx)*N {
 					session.Close()
 					return
diff --git a/stream.go b/stream.go
@@ -79,6 +79,16 @@ func (s *Stream) Read(b []byte) (n int, err error) {
 	}
 }
 
+// PeekSize returns the next size for Read
+func (s *Stream) PeekSize() int {
+	s.bufferLock.Lock()
+	defer s.bufferLock.Unlock()
+	if len(s.buffers) > 0 {
+		return len(s.buffers[0])
+	}
+	return 0
+}
+
 // TryRead is the nonblocking version of Read
 func (s *Stream) TryRead(b []byte) (n int, err error) {
 	if s.sess.config.Version == 2 {
