make yamux close with error non blocking

libp2p · Nov 21, 2024 · 3c53bf6 · 3c53bf6
1 parent 0e1eb18
commit 3c53bf6
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 13 deletions.
diff --git a/go.mod b/go.mod
@@ -32,7 +32,7 @@ require (
 	github.com/libp2p/go-nat v0.2.0
 	github.com/libp2p/go-netroute v0.2.1
 	github.com/libp2p/go-reuseport v0.4.0
-	github.com/libp2p/go-yamux/v4 v4.0.2-0.20241120100319-39abe7ed206a
+	github.com/libp2p/go-yamux/v4 v4.0.2-0.20241121144427-4ad80596076b
 	github.com/libp2p/zeroconf/v2 v2.2.0
 	github.com/marten-seemann/tcp v0.0.0-20210406111302-dfbc87cc63fd
 	github.com/mikioh/tcpinfo v0.0.0-20190314235526-30a79bb1804b

diff --git a/go.sum b/go.sum
@@ -194,8 +194,8 @@ github.com/libp2p/go-netroute v0.2.1 h1:V8kVrpD8GK0Riv15/7VN6RbUQ3URNZVosw7H2v9t
 github.com/libp2p/go-netroute v0.2.1/go.mod h1:hraioZr0fhBjG0ZRXJJ6Zj2IVEVNx6tDTFQfSmcq7mQ=
 github.com/libp2p/go-reuseport v0.4.0 h1:nR5KU7hD0WxXCJbmw7r2rhRYruNRl2koHw8fQscQm2s=
 github.com/libp2p/go-reuseport v0.4.0/go.mod h1:ZtI03j/wO5hZVDFo2jKywN6bYKWLOy8Se6DrI2E1cLU=
-github.com/libp2p/go-yamux/v4 v4.0.2-0.20241120100319-39abe7ed206a h1:zc7jPWFFQibZbACDyQdEAWg7yG/fjx5Jmg6djtpjKog=
-github.com/libp2p/go-yamux/v4 v4.0.2-0.20241120100319-39abe7ed206a/go.mod h1:PGP+3py2ZWDKABvqstBZtMnixEHNC7U/odnGylzur5o=
+github.com/libp2p/go-yamux/v4 v4.0.2-0.20241121144427-4ad80596076b h1:HWXTC6mwQRnPIgm0rgMAOIi/jebt/5AOa81EfU/fPPE=
+github.com/libp2p/go-yamux/v4 v4.0.2-0.20241121144427-4ad80596076b/go.mod h1:PGP+3py2ZWDKABvqstBZtMnixEHNC7U/odnGylzur5o=
 github.com/libp2p/zeroconf/v2 v2.2.0 h1:Cup06Jv6u81HLhIj1KasuNM/RHHrJ8T7wOTS4+Tv53Q=
 github.com/libp2p/zeroconf/v2 v2.2.0/go.mod h1:fuJqLnUwZTshS3U/bMRJ3+ow/v9oid1n0DmyYyNO1Xs=
 github.com/lunixbochs/vtclean v1.0.0/go.mod h1:pHhQNgMf3btfWnGBVipUOjRYhoOsdGqdm/+2c2E2WMI=

diff --git a/p2p/muxer/yamux/conn.go b/p2p/muxer/yamux/conn.go
@@ -23,8 +23,8 @@ func (c *conn) Close() error {
 	return c.yamux().Close()
 }
 
-func (c *conn) CloseWithError(errCode network.ConnErrorCode) error {
-	return c.yamux().CloseWithError(uint32(errCode))
+func (c *conn) CloseWithErrorChan(errCode network.ConnErrorCode) chan error {
+	return c.yamux().CloseWithErrorChan(uint32(errCode))
 }
 
 // IsClosed checks if yamux.Session is in closed state.
@@ -36,7 +36,7 @@ func (c *conn) IsClosed() bool {
 func (c *conn) OpenStream(ctx context.Context) (network.MuxedStream, error) {
 	s, err := c.yamux().OpenStream(ctx)
 	if err != nil {
-		return nil, parseResetError(err)
+		return nil, err
 	}
 
 	return (*stream)(s), nil
@@ -45,7 +45,7 @@ func (c *conn) OpenStream(ctx context.Context) (network.MuxedStream, error) {
 // AcceptStream accepts a stream opened by the other side.
 func (c *conn) AcceptStream() (network.MuxedStream, error) {
 	s, err := c.yamux().AcceptStream()
-	return (*stream)(s), parseResetError(err)
+	return (*stream)(s), err
 }
 
 func (c *conn) yamux() *yamux.Session {

diff --git a/p2p/net/upgrader/conn.go b/p2p/net/upgrader/conn.go
@@ -65,9 +65,19 @@ func (t *transportConn) ConnState() network.ConnectionState {
 }
 
 func (t *transportConn) CloseWithError(errCode network.ConnErrorCode) error {
-	defer t.scope.Done()
-	if ce, ok := t.MuxedConn.(network.CloseWithErrorer); ok {
-		return ce.CloseWithError(errCode)
+	if ce, ok := t.MuxedConn.(interface {
+		CloseWithErrorChan(errCode network.ConnErrorCode) chan error
+	}); ok {
+		err := t.scope.ReserveMemory(10_000, network.ReservationPriorityMedium)
+		if err != nil {
+			return t.Close()
+		}
+		errCh := ce.CloseWithErrorChan(errCode)
+		go func() {
+			defer t.scope.Done()
+			<-errCh
+		}()
+		return nil
 	}
 	return t.Close()
 }
diff --git a/test-plans/go.mod b/test-plans/go.mod
@@ -46,7 +46,7 @@ require (
 	github.com/libp2p/go-nat v0.2.0 // indirect
 	github.com/libp2p/go-netroute v0.2.1 // indirect
 	github.com/libp2p/go-reuseport v0.4.0 // indirect
-	github.com/libp2p/go-yamux/v4 v4.0.2-0.20241120100319-39abe7ed206a // indirect
+	github.com/libp2p/go-yamux/v4 v4.0.2-0.20241121144427-4ad80596076b // indirect
 	github.com/marten-seemann/tcp v0.0.0-20210406111302-dfbc87cc63fd // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/miekg/dns v1.1.62 // indirect

diff --git a/test-plans/go.sum b/test-plans/go.sum
@@ -150,8 +150,8 @@ github.com/libp2p/go-netroute v0.2.1 h1:V8kVrpD8GK0Riv15/7VN6RbUQ3URNZVosw7H2v9t
 github.com/libp2p/go-netroute v0.2.1/go.mod h1:hraioZr0fhBjG0ZRXJJ6Zj2IVEVNx6tDTFQfSmcq7mQ=
 github.com/libp2p/go-reuseport v0.4.0 h1:nR5KU7hD0WxXCJbmw7r2rhRYruNRl2koHw8fQscQm2s=
 github.com/libp2p/go-reuseport v0.4.0/go.mod h1:ZtI03j/wO5hZVDFo2jKywN6bYKWLOy8Se6DrI2E1cLU=
-github.com/libp2p/go-yamux/v4 v4.0.2-0.20241120100319-39abe7ed206a h1:zc7jPWFFQibZbACDyQdEAWg7yG/fjx5Jmg6djtpjKog=
-github.com/libp2p/go-yamux/v4 v4.0.2-0.20241120100319-39abe7ed206a/go.mod h1:PGP+3py2ZWDKABvqstBZtMnixEHNC7U/odnGylzur5o=
+github.com/libp2p/go-yamux/v4 v4.0.2-0.20241121144427-4ad80596076b h1:HWXTC6mwQRnPIgm0rgMAOIi/jebt/5AOa81EfU/fPPE=
+github.com/libp2p/go-yamux/v4 v4.0.2-0.20241121144427-4ad80596076b/go.mod h1:PGP+3py2ZWDKABvqstBZtMnixEHNC7U/odnGylzur5o=
 github.com/lunixbochs/vtclean v1.0.0/go.mod h1:pHhQNgMf3btfWnGBVipUOjRYhoOsdGqdm/+2c2E2WMI=
 github.com/mailru/easyjson v0.0.0-20190312143242-1de009706dbe/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc=
 github.com/marten-seemann/tcp v0.0.0-20210406111302-dfbc87cc63fd h1:br0buuQ854V8u83wA0rVZ8ttrq5CpaPZdvrK0LP2lOk=