Make Close() idempotent and close connection if session fails to close

This commit is contained in:
Elara 2023-01-04 14:08:18 -08:00
parent dfc248e5bb
commit 9944612ff5
1 changed files with 14 additions and 8 deletions

View File

@ -16,6 +16,7 @@ var _ drpc.Conn = &Conn{}
type Conn struct { type Conn struct {
conn io.ReadWriteCloser conn io.ReadWriteCloser
sess *yamux.Session sess *yamux.Session
isClosed bool
closed chan struct{} closed chan struct{}
} }
@ -36,14 +37,19 @@ func New(conn io.ReadWriteCloser) (*Conn, error) {
// Close closes the multiplexer session // Close closes the multiplexer session
// and the underlying connection. // and the underlying connection.
func (m *Conn) Close() error { func (m *Conn) Close() error {
if !m.isClosed {
m.isClosed = true
defer close(m.closed) defer close(m.closed)
err := m.sess.Close() err := m.sess.Close()
if err != nil { if err != nil {
m.conn.Close()
return err return err
} }
return m.conn.Close() return m.conn.Close()
} }
return nil
}
// Closed returns a channel that will be closed // Closed returns a channel that will be closed
// when the connection is closed // when the connection is closed