proper error handling

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)
 	}
 }

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)

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)
+	})
 }