Browse Source

Do not embed the recvBuffer

Armon Dadgar 9 years ago
parent
commit
b1464d5400
1 changed files with 9 additions and 4 deletions
  1. 9 4
      stream.go

+ 9 - 4
stream.go

@@ -33,7 +33,7 @@ type Stream struct {
 	state     streamState
 	stateLock sync.Mutex
 
-	recvBuf  bytes.Buffer
+	recvBuf  *bytes.Buffer
 	recvLock sync.Mutex
 
 	controlHdr     header
@@ -91,7 +91,7 @@ START:
 	case streamRemoteClose:
 		fallthrough
 	case streamClosed:
-		if s.recvBuf.Len() == 0 {
+		if s.recvBuf == nil || s.recvBuf.Len() == 0 {
 			s.stateLock.Unlock()
 			return 0, io.EOF
 		}
@@ -103,7 +103,7 @@ START:
 
 	// If there is no data available, block
 	s.recvLock.Lock()
-	if s.recvBuf.Len() == 0 {
+	if s.recvBuf == nil || s.recvBuf.Len() == 0 {
 		s.recvLock.Unlock()
 		goto WAIT
 	}
@@ -397,7 +397,12 @@ func (s *Stream) readData(hdr header, flags uint16, conn io.Reader) error {
 
 	// Copy into buffer
 	s.recvLock.Lock()
-	if _, err := io.Copy(&s.recvBuf, conn); err != nil {
+	if s.recvBuf == nil {
+		// Allocate the receive buffer just-in-time to fit the full data frame.
+		// This way we can read in the whole packet without further allocations.
+		s.recvBuf = bytes.NewBuffer(make([]byte, 0, length))
+	}
+	if _, err := io.Copy(s.recvBuf, conn); err != nil {
 		s.session.logger.Printf("[ERR] yamux: Failed to read stream data: %v", err)
 		s.recvLock.Unlock()
 		return err