diff --git a/session.go b/session.go index 1a2cf25..8bf7eb1 100644 --- a/session.go +++ b/session.go @@ -36,8 +36,8 @@ type Session struct { config *Config nextStreamID uint32 // next stream identifier - bucket int32 // token bucket - bucketCond *sync.Cond // used for waiting for tokens + bucket int32 // token bucket + bucketNotify chan struct{} // used for waiting for tokens streams map[uint32]*Stream // all streams in this session streamLock sync.Mutex // locks streams @@ -61,7 +61,7 @@ func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session { s.streams = make(map[uint32]*Stream) s.chAccepts = make(chan *Stream, defaultAcceptBacklog) s.bucket = int32(config.MaxReceiveBuffer) - s.bucketCond = sync.NewCond(&sync.Mutex{}) + s.bucketNotify = make(chan struct{}, 1) s.writes = make(chan writeRequest) if client { @@ -129,11 +129,19 @@ func (s *Session) Close() (err error) { s.streams[k].sessionClose() } s.streamLock.Unlock() - s.bucketCond.Signal() + s.notifyBucket() return s.conn.Close() } } +// notifyBucket notifies recvLoop that bucket is available +func (s *Session) notifyBucket() { + select { + case s.bucketNotify <- struct{}{}: + default: + } +} + // IsClosed does a safe check to see if we have shutdown func (s *Session) IsClosed() bool { select { @@ -166,7 +174,7 @@ func (s *Session) streamClosed(sid uint32) { s.streamLock.Lock() 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() + s.notifyBucket() } } delete(s.streams, sid) @@ -178,7 +186,7 @@ func (s *Session) returnTokens(n int) { oldvalue := atomic.LoadInt32(&s.bucket) newvalue := atomic.AddInt32(&s.bucket, int32(n)) if oldvalue <= 0 && newvalue > 0 { - s.bucketCond.Signal() + s.notifyBucket() } } @@ -211,14 +219,8 @@ func (s *Session) readFrame(buffer []byte) (f Frame, err error) { func (s *Session) recvLoop() { buffer := make([]byte, (1<<16)+headerSize) for { - s.bucketCond.L.Lock() for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() { - s.bucketCond.Wait() - } - s.bucketCond.L.Unlock() - - if s.IsClosed() { - return + <-s.bucketNotify } if f, err := s.readFrame(buffer); err == nil { @@ -272,7 +274,7 @@ func (s *Session) keepalive() { select { case <-tickerPing.C: s.writeFrame(newFrame(cmdNOP, 0)) - s.bucketCond.Signal() // force a signal to the recvLoop + s.notifyBucket() // force a signal to the recvLoop case <-tickerTimeout.C: if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) { s.Close()