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

View File

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