try to fix stream close

MatsuriDayo · Aug 22, 2022 · 0b38a8e · 0b38a8e
1 parent 1592141
commit 0b38a8e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/pkg/core/client.go b/pkg/core/client.go
@@ -166,7 +166,7 @@ func (c *Client) openStreamWithReconnect() (quic.Connection, quic.Stream, error)
 	stream, err := c.quicSession.OpenStream()
 	if err == nil {
 		// All good
-		return c.quicSession, &wrappedQUICStream{stream}, nil
+		return c.quicSession, &wrappedQUICStream{true, stream}, nil
 	}
 	// Something is wrong
 	if nErr, ok := err.(net.Error); ok && nErr.Temporary() {
@@ -180,7 +180,7 @@ func (c *Client) openStreamWithReconnect() (quic.Connection, quic.Stream, error)
 	}
 	// We are not going to try again even if it still fails the second time
 	stream, err = c.quicSession.OpenStream()
-	return c.quicSession, &wrappedQUICStream{stream}, err
+	return c.quicSession, &wrappedQUICStream{true, stream}, err
 }
 
 func (c *Client) DialTCP(addr string) (net.Conn, error) {

diff --git a/pkg/core/server_client.go b/pkg/core/server_client.go
@@ -91,7 +91,7 @@ func (c *serverClient) Run() error {
 			c.ConnGauge.Inc()
 		}
 		go func() {
-			stream := &wrappedQUICStream{stream}
+			stream := &wrappedQUICStream{false, stream}
 			c.handleStream(stream)
 			_ = stream.Close()
 			if c.ConnGauge != nil {

diff --git a/pkg/core/stream.go b/pkg/core/stream.go
@@ -10,7 +10,8 @@ import (
 // Handle stream close properly
 // Ref: https://github.com/libp2p/go-libp2p-quic-transport/blob/master/stream.go
 type wrappedQUICStream struct {
-	Stream quic.Stream
+	isClient bool
+	Stream   quic.Stream
 }
 
 func (s *wrappedQUICStream) StreamID() quic.StreamID {
@@ -34,6 +35,9 @@ func (s *wrappedQUICStream) Write(p []byte) (n int, err error) {
 }
 
 func (s *wrappedQUICStream) Close() error {
+	if s.isClient {
+		s.Stream.CancelWrite(0)
+	}
 	s.Stream.CancelRead(0)
 	return s.Stream.Close()
 }