From 9944612ff58b743cab2cac640ba1390f7f0c9869 Mon Sep 17 00:00:00 2001
From: Elara Musayelyan <elara@elara.ws>
Date: Wed, 4 Jan 2023 14:08:18 -0800
Subject: [PATCH] Make Close() idempotent and close connection if session fails
 to close

---
 muxconn/muxconn.go | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/muxconn/muxconn.go b/muxconn/muxconn.go
index d2aa29d..d42c9f1 100644
--- a/muxconn/muxconn.go
+++ b/muxconn/muxconn.go
@@ -14,9 +14,10 @@ var _ drpc.Conn = &Conn{}
 // Conn implements drpc.Conn using the yamux
 // multiplexer to allow concurrent RPCs
 type Conn struct {
-	conn   io.ReadWriteCloser
-	sess   *yamux.Session
-	closed chan struct{}
+	conn     io.ReadWriteCloser
+	sess     *yamux.Session
+	isClosed bool
+	closed   chan struct{}
 }
 
 // New returns a new multiplexed DRPC connection
@@ -36,13 +37,18 @@ func New(conn io.ReadWriteCloser) (*Conn, error) {
 // Close closes the multiplexer session
 // and the underlying connection.
 func (m *Conn) Close() error {
-	defer close(m.closed)
+	if !m.isClosed {
+		m.isClosed = true
+		defer close(m.closed)
 
-	err := m.sess.Close()
-	if err != nil {
-		return err
+		err := m.sess.Close()
+		if err != nil {
+			m.conn.Close()
+			return err
+		}
+		return m.conn.Close()
 	}
-	return m.conn.Close()
+	return nil
 }
 
 // Closed returns a channel that will be closed