optimize stream.Write

session.go

@@ -195,12 +195,10 @@ func (s *Session) recvLoop() {
 				if _, ok := s.streams[f.sid]; !ok {
 					stream := newStream(f.sid, s.config.MaxFrameSize, s)
 					s.streams[f.sid] = stream
-					go func() {
-						select {
-						case s.chAccepts <- stream:
-						case <-s.die:
-						}
-					}()
+					select {
+					case s.chAccepts <- stream:
+					case <-s.die:
+					}
 				}
 				s.streamLock.Unlock()
 			case cmdRST:

stream.go

@@ -2,6 +2,7 @@ package smux
 
 import (
 	"bytes"
+	"encoding/binary"
 	"sync"
 	"sync/atomic"
 
@@ -70,14 +71,21 @@ func (s *Stream) Write(b []byte) (n int, err error) {
 	}
 
 	frames := s.split(b, cmdPSH, s.id)
-	// combine the frames
-	var buffer bytes.Buffer
+	// preallocate buffer
+	buffer := make([]byte, len(frames)*headerSize+len(b))
+	bts := buffer
+
+	// combine frames into a large blob
 	for k := range frames {
-		bts, _ := frames[k].MarshalBinary()
-		buffer.Write(bts)
+		bts[0] = version
+		bts[1] = frames[k].cmd
+		binary.LittleEndian.PutUint16(bts[2:], uint16(len(frames[k].data)))
+		binary.LittleEndian.PutUint32(bts[4:], frames[k].sid)
+		copy(bts[headerSize:], frames[k].data)
+		bts = bts[len(frames[k].data)+headerSize:]
 	}
 
-	if _, err = s.sess.writeBinary(buffer.Bytes()); err != nil {
+	if _, err = s.sess.writeBinary(buffer); err != nil {
 		return 0, err
 	}
 	return len(b), nil