add support for sending error codes on session close

const.go

@@ -3,6 +3,7 @@
 import (
 	"encoding/binary"
 	"fmt"
+	"time"
 )
 
 type Error struct {
@@ -22,6 +23,61 @@
 	return ye.temporary
 }
 
+type GoAwayError struct {
+	ErrorCode uint32
+	Remote    bool
+}
+
+func (e *GoAwayError) Error() string {
+	if e.Remote {
+		return fmt.Sprintf("remote sent go away, code: %d", e.ErrorCode)
+	}
+	return fmt.Sprintf("sent go away, code: %d", e.ErrorCode)
+}
+
+func (e *GoAwayError) Timeout() bool {
+	return false
+}
+
+func (e *GoAwayError) Temporary() bool {
+	return false
+}
+
+func (e *GoAwayError) Is(target error) bool {
+	// to maintain compatibility with errors returned by previous versions
+	if e.Remote && target == ErrRemoteGoAway {
+		return true
+	} else if !e.Remote && target == ErrSessionShutdown {
+		return true
+	}
+
+	if err, ok := target.(*GoAwayError); ok {
+		return *e == *err
+	}
+	return false
+}
+
+// A StreamError is used for errors returned from Read and Write calls after the stream is Reset
+type StreamError struct {
+	ErrorCode uint32
+	Remote    bool
+}
+
+func (s *StreamError) Error() string {
+	if s.Remote {
+		return fmt.Sprintf("stream reset by remote, error code: %d", s.ErrorCode)
+	}
+	return fmt.Sprintf("stream reset, error code: %d", s.ErrorCode)
+}
+
+func (s *StreamError) Is(target error) bool {
+	if target == ErrStreamReset {
+		return true
+	}
+	e, ok := target.(*StreamError)
+	return ok && *e == *s
+}
+
 var (
 	// ErrInvalidVersion means we received a frame with an
 	// invalid version
@@ -33,7 +89,7 @@
 
 	// ErrSessionShutdown is used if there is a shutdown during
 	// an operation
-	ErrSessionShutdown = &Error{msg: "session shutdown"}
+	ErrSessionShutdown = &GoAwayError{ErrorCode: goAwayNormal, Remote: false}
 
 	// ErrStreamsExhausted is returned if we have no more
 	// stream ids to issue
@@ -56,7 +112,7 @@
 	ErrUnexpectedFlag = &Error{msg: "unexpected flag"}
 
 	// ErrRemoteGoAway is used when we get a go away from the other side
-	ErrRemoteGoAway = &Error{msg: "remote end is not accepting connections"}
+	ErrRemoteGoAway = &GoAwayError{Remote: true, ErrorCode: goAwayNormal}
 
 	// ErrStreamReset is sent if a stream is reset. This can happen
 	// if the backlog is exceeded, or if there was a remote GoAway.
@@ -117,6 +173,7 @@
 	// It's not an implementation choice, the value defined in the specification.
 	initialStreamWindow = 256 * 1024
 	maxStreamWindow     = 16 * 1024 * 1024
+	goAwayWaitTime      = 100 * time.Millisecond
 )
 
 const (

session.go

@@ -102,6 +102,8 @@ type Session struct {
 	// recvDoneCh is closed when recv() exits to avoid a race
 	// between stream registration and stream shutdown
 	recvDoneCh chan struct{}
+	// recvErr is the error the receive loop ended with
+	recvErr error
 
 	// sendDoneCh is closed when send() exits to avoid a race
 	// between returning from a Stream.Write and exiting from the send loop
@@ -284,8 +286,22 @@ func (s *Session) AcceptStream() (*Stream, error) {
 }
 
 // Close is used to close the session and all streams.
-// Attempts to send a GoAway before closing the connection.
+// Attempts to send a GoAway before closing the connection. The GoAway may not actually be sent depending on the
+// semantics of the underlying net.Conn. For TCP connections, it may be dropped depending on LINGER value or
+// if there's unread data in the kernel receive buffer.
 func (s *Session) Close() error {
+	return s.close(ErrSessionShutdown, false, goAwayNormal)
+}
+
+// CloseWithError is used to close the session and all streams after sending a GoAway message with errCode.
+// The GoAway may not actually be sent depending on the semantics of the underlying net.Conn.
+// For TCP connections, it may be dropped depending on LINGER value or if there's unread data in the kernel
+// receive buffer.
+func (s *Session) CloseWithError(errCode uint32) error {
+	return s.close(&GoAwayError{Remote: false, ErrorCode: errCode}, true, errCode)
+}
+
+func (s *Session) close(shutdownErr error, sendGoAway bool, errCode uint32) error {
 	s.shutdownLock.Lock()
 	defer s.shutdownLock.Unlock()
 
@@ -294,35 +310,37 @@ func (s *Session) Close() error {
 	}
 	s.shutdown = true
 	if s.shutdownErr == nil {
-		s.shutdownErr = ErrSessionShutdown
+		s.shutdownErr = shutdownErr
 	}
 	close(s.shutdownCh)
-	s.conn.Close()
 	s.stopKeepalive()
-	<-s.recvDoneCh
+
+	if sendGoAway {
+		// wait for write loop to exit
+		// We need to write the current frame completely before sending a goaway.
+		// This will wait for at most s.config.ConnectionWriteTimeout
+		<-s.sendDoneCh
+		ga := s.goAway(errCode)
+		if err := s.conn.SetWriteDeadline(time.Now().Add(goAwayWaitTime)); err == nil {
+			_, _ = s.conn.Write(ga[:]) // there's nothing we can do on error here
+		}
+		s.conn.SetWriteDeadline(time.Time{})
+	}
+
+	s.conn.Close()
 	<-s.sendDoneCh
+	<-s.recvDoneCh
 
 	s.streamLock.Lock()
 	defer s.streamLock.Unlock()
 	for id, stream := range s.streams {
-		stream.forceClose()
+		stream.forceClose(fmt.Errorf("%w: connection closed: %w", ErrStreamReset, s.shutdownErr))
 		delete(s.streams, id)
 		stream.memorySpan.Done()
 	}
 	return nil
 }
 
-// exitErr is used to handle an error that is causing the
-// session to terminate.
-func (s *Session) exitErr(err error) {
-	s.shutdownLock.Lock()
-	if s.shutdownErr == nil {
-		s.shutdownErr = err
-	}
-	s.shutdownLock.Unlock()
-	s.Close()
-}
-
 // GoAway can be used to prevent accepting further
 // connections. It does not close the underlying conn.
 func (s *Session) GoAway() error {
@@ -451,7 +469,7 @@ func (s *Session) startKeepalive() {
 
 		if err != nil {
 			s.logger.Printf("[ERR] yamux: keepalive failed: %v", err)
-			s.exitErr(ErrKeepAliveTimeout)
+			s.close(ErrKeepAliveTimeout, false, 0)
 		}
 	})
 }
@@ -516,7 +534,19 @@ func (s *Session) sendMsg(hdr header, body []byte, deadline <-chan struct{}) err
 // send is a long running goroutine that sends data
 func (s *Session) send() {
 	if err := s.sendLoop(); err != nil {
-		s.exitErr(err)
+		// Prefer the recvLoop error over the sendLoop error. The receive loop might have the error code
+		// received in a GoAway frame received just before the TCP RST that closed the sendLoop
+		s.shutdownLock.Lock()
+		if s.shutdownErr == nil {
+			s.conn.Close()
+			<-s.recvDoneCh
+			if _, ok := s.recvErr.(*GoAwayError); ok {
+				err = s.recvErr
+			}
+			s.shutdownErr = err
+		}
+		s.shutdownLock.Unlock()
+		s.close(err, false, 0)
 	}
 }
 
@@ -644,7 +674,7 @@ func (s *Session) sendLoop() (err error) {
 // recv is a long running goroutine that accepts new data
 func (s *Session) recv() {
 	if err := s.recvLoop(); err != nil {
-		s.exitErr(err)
+		s.close(err, false, 0)
 	}
 }
 
@@ -666,7 +696,10 @@ func (s *Session) recvLoop() (err error) {
 			err = fmt.Errorf("panic in yamux receive loop: %s", rerr)
 		}
 	}()
-	defer close(s.recvDoneCh)
+	defer func() {
+		s.recvErr = err
+		close(s.recvDoneCh)
+	}()
 	var hdr header
 	for {
 		// fmt.Printf("ReadFull from %#v\n", s.reader)
@@ -782,17 +815,17 @@ func (s *Session) handleGoAway(hdr header) error {
 	switch code {
 	case goAwayNormal:
 		atomic.SwapInt32(&s.remoteGoAway, 1)
+		// Don't close connection on normal go away. Let the existing streams
+		// complete gracefully.
+		return nil
 	case goAwayProtoErr:
 		s.logger.Printf("[ERR] yamux: received protocol error go away")
-		return fmt.Errorf("yamux protocol error")
 	case goAwayInternalErr:
 		s.logger.Printf("[ERR] yamux: received internal error go away")
-		return fmt.Errorf("remote yamux internal error")
 	default:
-		s.logger.Printf("[ERR] yamux: received unexpected go away")
-		return fmt.Errorf("unexpected go away received")
+		s.logger.Printf("[ERR] yamux: received go away with error code: %d", code)
 	}
-	return nil
+	return &GoAwayError{Remote: true, ErrorCode: code}
 }
 
 // incomingStream is used to create a new incoming stream