From 8ff34859ff3c85ffe8a85f359e1f70714f4b485f Mon Sep 17 00:00:00 2001 From: xtaci Date: Fri, 2 Sep 2016 23:17:34 +0800 Subject: [PATCH] use sync.Cond to control receiving --- mux.go | 8 ++-- mux_test.go | 2 +- session.go | 126 +++++++++++++++++++++++++++------------------------- 3 files changed, 71 insertions(+), 65 deletions(-) diff --git a/mux.go b/mux.go index 5fbd011..5bf24f5 100644 --- a/mux.go +++ b/mux.go @@ -22,8 +22,8 @@ type Config struct { MaxFrameSize int // MaxFrameTokens is used to control the maximum - // number of frame in the buffer pool - MaxFrameTokens int + // number of data in the buffer pool + MaxReceiveBuffer int } // DefaultConfig is used to return a default configuration @@ -32,7 +32,7 @@ func DefaultConfig() *Config { KeepAliveInterval: 10 * time.Second, KeepAliveTimeout: 30 * time.Second, MaxFrameSize: 4096, - MaxFrameTokens: 4096, + MaxReceiveBuffer: 4194304, } } @@ -50,7 +50,7 @@ func VerifyConfig(config *Config) error { if config.MaxFrameSize > 65535 { return errors.New("max frame size must not be larger than 65535") } - if config.MaxFrameTokens <= 0 { + if config.MaxReceiveBuffer <= 0 { return errors.New("max frame tokens must be positive") } return nil diff --git a/mux_test.go b/mux_test.go index 30fc6e6..638e67c 100644 --- a/mux_test.go +++ b/mux_test.go @@ -51,7 +51,7 @@ func TestConfig(t *testing.T) { } config = DefaultConfig() - config.MaxFrameTokens = 0 + config.MaxReceiveBuffer = 0 err = VerifyConfig(config) t.Log(err) if err == nil { diff --git a/session.go b/session.go index 676e261..09712b7 100644 --- a/session.go +++ b/session.go @@ -27,7 +27,8 @@ type Session struct { config *Config nextStreamID uint32 // next stream identifier - tbf chan struct{} // token bucket + bucket int32 + tbfCond *sync.Cond frameQueues map[uint32][]Frame // stream input frame queue streams map[uint32]*Stream // all streams in this session streamLock sync.Mutex // locks streams && frameQueues @@ -49,10 +50,8 @@ func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session { s.frameQueues = make(map[uint32][]Frame) s.chAccepts = make(chan *Stream, defaultAcceptBacklog) s.chClosedStream = make(chan uint32, defaultCloseWait) - s.tbf = make(chan struct{}, config.MaxFrameTokens) - for i := 0; i < config.MaxFrameTokens; i++ { - s.tbf <- struct{}{} - } + s.bucket = int32(config.MaxReceiveBuffer) + s.tbfCond = sync.NewCond(&sync.Mutex{}) if client { s.nextStreamID = 1 } else { @@ -109,6 +108,7 @@ func (s *Session) Close() error { s.writeFrame(newFrame(cmdTerminate, 0)) s.conn.Close() close(s.die) + s.tbfCond.Signal() } return nil } @@ -150,8 +150,9 @@ func (s *Session) nioread(sid uint32) *Frame { if len(frames) > 0 { f := frames[0] s.frameQueues[sid] = frames[1:] - s.tbf <- struct{}{} + atomic.AddInt32(&s.bucket, int32(len(f.data))) s.streamLock.Unlock() + s.tbfCond.Signal() return &f } s.streamLock.Unlock() @@ -187,12 +188,13 @@ func (s *Session) monitor() { case sid := <-s.chClosedStream: s.streamLock.Lock() delete(s.streams, sid) - ntokens := len(s.frameQueues[sid]) + fq := s.frameQueues[sid] + for k := range fq { // return remaining tokens to the bucket + atomic.AddInt32(&s.bucket, int32(len(fq[k].data))) + } + s.tbfCond.Signal() delete(s.frameQueues, sid) s.streamLock.Unlock() - for i := 0; i < ntokens; i++ { // return remaining tokens to the bucket - s.tbf <- struct{}{} - } case <-s.die: return } @@ -203,61 +205,65 @@ func (s *Session) monitor() { func (s *Session) recvLoop() { buffer := make([]byte, (1<<16)+headerSize) for { - select { - case <-s.tbf: - s.tbf <- struct{}{} - if f, err := s.readFrame(buffer); err == nil { - atomic.StoreInt32(&s.dataReady, 1) + s.tbfCond.L.Lock() + for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() { + s.tbfCond.Wait() + } + s.tbfCond.L.Unlock() - switch f.cmd { - case cmdNOP: - case cmdTerminate: - s.Close() - return - case cmdSYN: - rstflag := false - s.streamLock.Lock() - if _, ok := s.streams[f.sid]; !ok { - s.streams[f.sid] = newStream(f.sid, s.config.MaxFrameSize, s) - s.chAccepts <- s.streams[f.sid] - } else { // stream exists, RST the peer - rstflag = true - } - s.streamLock.Unlock() + if s.IsClosed() { + return + } - if rstflag { - s.writeFrame(newFrame(cmdRST, f.sid)) - } - case cmdRST: - s.streamLock.Lock() - if _, ok := s.streams[f.sid]; ok { - s.streams[f.sid].Close() - } else { // must do nothing if stream is absent - } - s.streamLock.Unlock() - case cmdPSH: - rstflag := false - s.streamLock.Lock() - if stream, ok := s.streams[f.sid]; ok { - <-s.tbf // remove a token - s.frameQueues[f.sid] = append(s.frameQueues[f.sid], f) - stream.notifyReadEvent() - } else { // stream is absent - rstflag = true - } - s.streamLock.Unlock() - if rstflag { - s.writeFrame(newFrame(cmdRST, f.sid)) - } - default: - s.Close() - return + if f, err := s.readFrame(buffer); err == nil { + atomic.StoreInt32(&s.dataReady, 1) + + switch f.cmd { + case cmdNOP: + case cmdTerminate: + s.Close() + return + case cmdSYN: + rstflag := false + s.streamLock.Lock() + if _, ok := s.streams[f.sid]; !ok { + s.streams[f.sid] = newStream(f.sid, s.config.MaxFrameSize, s) + s.chAccepts <- s.streams[f.sid] + } else { // stream exists, RST the peer + rstflag = true } - } else { + s.streamLock.Unlock() + + if rstflag { + s.writeFrame(newFrame(cmdRST, f.sid)) + } + case cmdRST: + s.streamLock.Lock() + if _, ok := s.streams[f.sid]; ok { + s.streams[f.sid].Close() + } else { // must do nothing if stream is absent + } + s.streamLock.Unlock() + case cmdPSH: + rstflag := false + s.streamLock.Lock() + if stream, ok := s.streams[f.sid]; ok { + atomic.AddInt32(&s.bucket, -int32(len(f.data)-1)) + s.frameQueues[f.sid] = append(s.frameQueues[f.sid], f) + stream.notifyReadEvent() + } else { // stream is absent + rstflag = true + } + s.streamLock.Unlock() + if rstflag { + s.writeFrame(newFrame(cmdRST, f.sid)) + } + default: s.Close() return } - case <-s.die: + } else { + s.Close() return } } @@ -273,7 +279,7 @@ func (s *Session) keepalive() { case <-tickerPing.C: s.writeFrame(newFrame(cmdNOP, 0)) case <-tickerTimeout.C: - if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) && len(s.tbf) == s.config.MaxFrameTokens { + if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) { s.Close() return }