diff --git a/session.go b/session.go index 988beba..c41bed5 100644 --- a/session.go +++ b/session.go @@ -48,9 +48,21 @@ type Session struct { streams map[uint32]*Stream // all streams in this session streamLock sync.Mutex // locks streams - die chan struct{} // flag session has died - dieOnce sync.Once - socketError atomic.Value // errors from underlying conn + die chan struct{} // flag session has died + dieOnce sync.Once + + // socket error handling + socketReadError atomic.Value + socketWriteError atomic.Value + chSocketReadError chan struct{} + chSocketWriteError chan struct{} + socketReadErrorOnce sync.Once + socketWriteErrorOnce sync.Once + + // smux protocol errors + protoError atomic.Value + chProtoError chan struct{} + protoErrorOnce sync.Once chAccepts chan *Stream @@ -73,6 +85,9 @@ func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session { s.bucket = int32(config.MaxReceiveBuffer) s.bucketNotify = make(chan struct{}, 1) s.writes = make(chan writeRequest) + s.chSocketReadError = make(chan struct{}) + s.chSocketWriteError = make(chan struct{}) + s.chProtoError = make(chan struct{}) if client { s.nextStreamID = 1 @@ -116,12 +131,10 @@ func (s *Session) OpenStream() (*Stream, error) { s.streamLock.Lock() defer s.streamLock.Unlock() select { + case <-s.chSocketWriteError: + return nil, s.socketWriteError.Load().(error) case <-s.die: - if err := s.socketError.Load(); err != nil { - return nil, errors.WithStack(err.(error)) - } else { - return nil, errors.WithStack(io.ErrClosedPipe) - } + return nil, errors.WithStack(io.ErrClosedPipe) default: s.streams[sid] = stream return stream, nil @@ -137,27 +150,27 @@ func (s *Session) AcceptStream() (*Stream, error) { defer timer.Stop() deadline = timer.C } + select { case stream := <-s.chAccepts: return stream, nil case <-deadline: return nil, errors.WithStack(errTimeout) + case <-s.chSocketReadError: + return nil, s.socketReadError.Load().(error) + case <-s.chProtoError: + return nil, s.protoError.Load().(error) case <-s.die: - if err := s.socketError.Load(); err != nil { - return nil, errors.WithStack(err.(error)) - } else { - return nil, errors.WithStack(io.ErrClosedPipe) - } + return nil, errors.WithStack(io.ErrClosedPipe) } } // Close is used to close the session and all streams. func (s *Session) Close() error { var once bool + var err error s.dieOnce.Do(func() { - if err := s.conn.Close(); err != nil { - s.socketError.Store(errors.WithStack(err)) - } + err = s.conn.Close() s.streamLock.Lock() for k := range s.streams { s.streams[k].sessionClose() @@ -166,15 +179,11 @@ func (s *Session) Close() error { close(s.die) }) - if err := s.socketError.Load(); err != nil { - return errors.WithStack(err.(error)) - } - - if !once { + if once { + return err + } else { return errors.WithStack(io.ErrClosedPipe) } - - return nil } // notifyBucket notifies recvLoop that bucket is available @@ -185,6 +194,27 @@ func (s *Session) notifyBucket() { } } +func (s *Session) notifyReadError(err error) { + s.socketReadErrorOnce.Do(func() { + s.socketReadError.Store(err) + close(s.chSocketReadError) + }) +} + +func (s *Session) notifyWriteError(err error) { + s.socketWriteErrorOnce.Do(func() { + s.socketWriteError.Store(err) + close(s.chSocketWriteError) + }) +} + +func (s *Session) notifyProtoError(err error) { + s.protoErrorOnce.Do(func() { + s.protoError.Store(err) + close(s.chProtoError) + }) +} + // IsClosed does a safe check to see if we have shutdown func (s *Session) IsClosed() bool { select { @@ -248,7 +278,7 @@ func (s *Session) recvLoop() { if _, err := io.ReadFull(s.conn, hdr[:]); err == nil { atomic.StoreInt32(&s.dataReady, 1) if hdr.Version() != version { - s.Close() + s.notifyProtoError(errors.WithStack(errInvalidProtocol)) return } sid := hdr.StreamID() @@ -268,7 +298,7 @@ func (s *Session) recvLoop() { case cmdFIN: s.streamLock.Lock() if stream, ok := s.streams[sid]; ok { - stream.markRST() + stream.fin() stream.notifyReadEvent() } s.streamLock.Unlock() @@ -284,19 +314,16 @@ func (s *Session) recvLoop() { } s.streamLock.Unlock() } else { - s.socketError.Store(errors.WithStack(err)) - s.Close() + s.notifyReadError(errors.WithStack(err)) return } } default: - s.socketError.Store(errors.WithStack(errInvalidProtocol)) - s.Close() + s.notifyProtoError(errors.WithStack(errInvalidProtocol)) return } } else { - s.socketError.Store(errors.WithStack(err)) - s.Close() + s.notifyReadError(errors.WithStack(err)) return } } @@ -371,8 +398,7 @@ func (s *Session) sendLoop() { // store conn error if err != nil { - s.socketError.Store(errors.WithStack(err)) - s.Close() + s.notifyWriteError(errors.WithStack(err)) return } } @@ -392,9 +418,11 @@ func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time) (int, e result: make(chan writeResult, 1), } select { + case s.writes <- req: case <-s.die: return 0, errors.WithStack(io.ErrClosedPipe) - case s.writes <- req: + case <-s.chSocketWriteError: + return 0, s.socketWriteError.Load().(error) case <-deadline: return 0, errors.WithStack(errTimeout) } @@ -402,9 +430,11 @@ func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time) (int, e select { case result := <-req.result: return result.n, errors.WithStack(result.err) - case <-deadline: - return 0, errors.WithStack(errTimeout) case <-s.die: return 0, errors.WithStack(io.ErrClosedPipe) + case <-s.chSocketWriteError: + return 0, s.socketWriteError.Load().(error) + case <-deadline: + return 0, errors.WithStack(errTimeout) } } diff --git a/session_test.go b/session_test.go index ebcb2fc..5849c59 100644 --- a/session_test.go +++ b/session_test.go @@ -5,14 +5,23 @@ import ( "encoding/binary" "fmt" "io" + "log" "math/rand" "net" + "net/http" + _ "net/http/pprof" "strings" "sync" "testing" "time" ) +func init() { + go func() { + log.Println(http.ListenAndServe("0.0.0.0:6060", nil)) + }() +} + // setupServer starts new server listening on a random localhost port and // returns address of the server, function to stop the server, new client // connection to this server or an error. @@ -820,6 +829,7 @@ func bench(b *testing.B, rd io.Reader, wr io.Writer) { buf2 := make([]byte, 128*1024) b.SetBytes(128 * 1024) b.ResetTimer() + b.ReportAllocs() var wg sync.WaitGroup wg.Add(1) diff --git a/stream.go b/stream.go index 2b16760..5da942c 100644 --- a/stream.go +++ b/stream.go @@ -12,15 +12,24 @@ import ( // Stream implements net.Conn type Stream struct { - id uint32 - rstflag int32 - sess *Session - buffers [][]byte - bufferLock sync.Mutex - frameSize int - chReadEvent chan struct{} // notify a read event - die chan struct{} // flag the stream has closed - dieLock sync.Mutex + id uint32 + sess *Session + buffers [][]byte + bufferLock sync.Mutex + frameSize int + + // notify a read event + chReadEvent chan struct{} + + // flag the stream has closed + die chan struct{} + dieOnce sync.Once + + // FIN + chFinEvent chan struct{} + finEventOnce sync.Once + + // deadlines readDeadline atomic.Value writeDeadline atomic.Value } @@ -33,6 +42,7 @@ func newStream(id uint32, frameSize int, sess *Session) *Stream { s.frameSize = frameSize s.sess = sess s.die = make(chan struct{}) + s.chFinEvent = make(chan struct{}) return s } @@ -44,49 +54,50 @@ func (s *Stream) ID() uint32 { // Read implements net.Conn func (s *Stream) Read(b []byte) (n int, err error) { if len(b) == 0 { + return 0, nil + } + + for { + s.bufferLock.Lock() + if len(s.buffers) > 0 { + n = copy(b, s.buffers[0]) + s.buffers[0] = s.buffers[0][n:] + if len(s.buffers[0]) == 0 { + s.buffers[0] = nil + s.buffers = s.buffers[1:] + } + } + s.bufferLock.Unlock() + + if n > 0 { + s.sess.returnTokens(n) + return n, nil + } + + var timer *time.Timer + var deadline <-chan time.Time + if d, ok := s.readDeadline.Load().(time.Time); ok && !d.IsZero() { + timer = time.NewTimer(time.Until(d)) + deadline = timer.C + } + select { + case <-s.chReadEvent: + if timer != nil { + timer.Stop() + } + case <-s.chFinEvent: + return 0, errors.WithStack(io.EOF) + case <-s.sess.chSocketReadError: + return 0, s.sess.socketReadError.Load().(error) + case <-s.sess.chProtoError: + return 0, s.sess.protoError.Load().(error) + case <-deadline: + return n, errors.WithStack(errTimeout) case <-s.die: return 0, errors.WithStack(io.ErrClosedPipe) - default: - return 0, nil } } - - var deadline <-chan time.Time - if d, ok := s.readDeadline.Load().(time.Time); ok && !d.IsZero() { - timer := time.NewTimer(time.Until(d)) - defer timer.Stop() - deadline = timer.C - } - -READ: - s.bufferLock.Lock() - if len(s.buffers) > 0 { - n = copy(b, s.buffers[0]) - s.buffers[0] = s.buffers[0][n:] - if len(s.buffers[0]) == 0 { - s.buffers[0] = nil - s.buffers = s.buffers[1:] - } - } - 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.WithStack(io.EOF) - } - - select { - case <-s.chReadEvent: - goto READ - case <-deadline: - return n, errors.WithStack(errTimeout) - case <-s.die: - return 0, errors.WithStack(io.ErrClosedPipe) - } } // Write implements net.Conn @@ -127,18 +138,19 @@ func (s *Stream) Write(b []byte) (n int, err error) { // Close implements net.Conn func (s *Stream) Close() error { - s.dieLock.Lock() - - select { - case <-s.die: - s.dieLock.Unlock() - return errors.WithStack(io.ErrClosedPipe) - default: + var once bool + var err error + s.dieOnce.Do(func() { close(s.die) - s.dieLock.Unlock() + once = true + }) + + if once { + _, err = s.sess.writeFrame(newFrame(cmdFIN, s.id)) s.sess.streamClosed(s.id) - _, err := s.sess.writeFrame(newFrame(cmdFIN, s.id)) - return errors.WithStack(err) + return err + } else { + return errors.WithStack(io.ErrClosedPipe) } } @@ -177,17 +189,8 @@ func (s *Stream) SetDeadline(t time.Time) error { return nil } -// session closes the stream -func (s *Stream) sessionClose() { - s.dieLock.Lock() - defer s.dieLock.Unlock() - - select { - case <-s.die: - default: - close(s.die) - } -} +// session closes +func (s *Stream) sessionClose() { s.dieOnce.Do(func() { close(s.die) }) } // LocalAddr satisfies net.Conn interface func (s *Stream) LocalAddr() net.Addr { @@ -236,7 +239,9 @@ func (s *Stream) notifyReadEvent() { } } -// mark this stream has been reset -func (s *Stream) markRST() { - atomic.StoreInt32(&s.rstflag, 1) +// mark this stream has been closed in protocol +func (s *Stream) fin() { + s.finEventOnce.Do(func() { + close(s.chFinEvent) + }) }