diff --git a/session.go b/session.go index e3dfea5..19b008c 100644 --- a/session.go +++ b/session.go @@ -1,7 +1,6 @@ package smux import ( - "bytes" "io" "sync" "sync/atomic" @@ -31,9 +30,8 @@ type Session struct { bucket int32 // token bucket bucketCond *sync.Cond // used for waiting for tokens - streamBuffers map[uint32]*bytes.Buffer // stream input buffer - streams map[uint32]*Stream // all streams in this session - streamLock sync.Mutex // locks streams && frameQueues + streams map[uint32]*Stream // all streams in this session + streamLock sync.Mutex // locks streams && frameQueues die chan struct{} // flag session has died dieLock sync.Mutex @@ -49,7 +47,6 @@ func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session { s.conn = conn s.config = config s.streams = make(map[uint32]*Stream) - s.streamBuffers = make(map[uint32]*bytes.Buffer) s.chAccepts = make(chan *Stream, defaultAcceptBacklog) s.chClosedStream = make(chan uint32, defaultCloseWait) s.bucket = int32(config.MaxReceiveBuffer) @@ -73,11 +70,9 @@ func (s *Session) OpenStream() (*Stream, error) { sid := atomic.AddUint32(&s.nextStreamID, 2) stream := newStream(sid, s.config.MaxFrameSize, s) - streamBuffer := new(bytes.Buffer) s.streamLock.Lock() s.streams[sid] = stream - s.streamBuffers[sid] = streamBuffer s.streamLock.Unlock() s.writeFrame(newFrame(cmdSYN, sid)) @@ -146,18 +141,11 @@ func (s *Session) streamClosed(sid uint32) { }() } -// nonblocking read from session pool, for streams -func (s *Session) nioread(sid uint32, p []byte) (n int) { - s.streamLock.Lock() - if streamBuffer, ok := s.streamBuffers[sid]; ok { - n, _ = streamBuffer.Read(p) - if n > 0 { - atomic.AddInt32(&s.bucket, int32(n)) - s.bucketCond.Signal() - } +// returnTokens +func (s *Session) returnTokens(n int) { + if atomic.AddInt32(&s.bucket, int32(n)) > 0 { + s.bucketCond.Signal() } - s.streamLock.Unlock() - return } // session read a frame from underlying connection @@ -189,12 +177,12 @@ func (s *Session) monitor() { select { case sid := <-s.chClosedStream: s.streamLock.Lock() - delete(s.streams, sid) - if n := s.streamBuffers[sid].Len(); n > 0 { // return remaining tokens to the bucket - atomic.AddInt32(&s.bucket, int32(n)) - s.bucketCond.Signal() + if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket + if atomic.AddInt32(&s.bucket, int32(n)) > 0 { + s.bucketCond.Signal() + } } - delete(s.streamBuffers, sid) + delete(s.streams, sid) s.streamLock.Unlock() case <-s.die: return @@ -225,9 +213,7 @@ func (s *Session) recvLoop() { s.streamLock.Lock() if _, ok := s.streams[f.sid]; !ok { stream := newStream(f.sid, s.config.MaxFrameSize, s) - streamBuffer := new(bytes.Buffer) s.streams[f.sid] = stream - s.streamBuffers[f.sid] = streamBuffer go func() { s.chAccepts <- stream }() } else { // stream exists, RST the peer go s.writeFrame(newFrame(cmdRST, f.sid)) @@ -244,7 +230,7 @@ func (s *Session) recvLoop() { s.streamLock.Lock() if stream, ok := s.streams[f.sid]; ok { atomic.AddInt32(&s.bucket, -int32(len(f.data))) - s.streamBuffers[f.sid].Write(f.data) + stream.pushBytes(f.data) stream.notifyReadEvent() } else { // stream is absent go s.writeFrame(newFrame(cmdRST, f.sid)) diff --git a/stream.go b/stream.go index 17d6717..9a21155 100644 --- a/stream.go +++ b/stream.go @@ -13,6 +13,8 @@ type Stream struct { id uint32 rstflag int32 sess *Session + buffer bytes.Buffer + bufferLock sync.Mutex frameSize int chReadEvent chan struct{} // notify a read event die chan struct{} // flag the stream has closed @@ -39,8 +41,13 @@ READ: default: } - if n = s.sess.nioread(s.id, b); n > 0 { - return n, err + s.bufferLock.Lock() + n, err = s.buffer.Read(b) + s.bufferLock.Unlock() + + if n > 0 { + s.sess.returnTokens(n) + return n, nil } else if atomic.LoadInt32(&s.rstflag) == 1 { s.Close() return 0, errors.New(errBrokenPipe) @@ -106,6 +113,22 @@ func (s *Stream) sessionClose() error { return nil } +// pushBytes a slice into buffer +func (s *Stream) pushBytes(p []byte) { + s.bufferLock.Lock() + s.buffer.Write(p) + s.bufferLock.Unlock() +} + +// recycleTokens transform remaining bytes to tokens(will truncate buffer) +func (s *Stream) recycleTokens() (n int) { + s.bufferLock.Lock() + n = s.buffer.Len() + s.buffer.Reset() + s.bufferLock.Unlock() + return +} + // split large byte buffer into smaller frames func (s *Stream) split(bts []byte, cmd byte, sid uint32) []Frame { var frames []Frame