Fix goroutine leak from deadlock when closing

Author: nhooyr

conn_test.go

@@ -348,7 +348,7 @@ func TestConn(t *testing.T) {
 		}()
 		defer cancel()
 
-		c1.SetReadLimit(131072)
+		c1.SetReadLimit(1 << 30)
 
 		exp := xrand.String(xrand.Int(131072))
 
@@ -401,8 +401,6 @@ func TestConn(t *testing.T) {
 		}()
 		defer cancel()
 
-		c1.SetReadLimit(131072)
-
 		exp := ptypes.DurationProto(100)
 		err = wspb.Write(ctx, c1, exp)
 		if err != nil {

read.go

@@ -69,7 +69,9 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
 //
 // When the limit is hit, the connection will be closed with StatusMessageTooBig.
 func (c *Conn) SetReadLimit(n int64) {
-	c.msgReader.limitReader.limit.Store(n)
+	// We add read one more byte than the limit in case
+	// there is a fin frame that needs to be read.
+	c.msgReader.limitReader.limit.Store(n + 1)
 }
 
 const defaultReadLimit = 32768
@@ -80,7 +82,7 @@ func newMsgReader(c *Conn) *msgReader {
 		fin: true,
 	}
 
-	mr.limitReader = newLimitReader(c, readerFunc(mr.read), defaultReadLimit)
+	mr.limitReader = newLimitReader(c, readerFunc(mr.read), defaultReadLimit+1)
 	return mr
 }
 

write.go

@@ -50,7 +50,8 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error {
 type msgWriter struct {
 	c *Conn
 
-	mu *mu
+	mu       *mu
+	activeMu *mu
 
 	ctx    context.Context
 	opcode opcode
@@ -63,8 +64,9 @@ type msgWriter struct {
 
 func newMsgWriter(c *Conn) *msgWriter {
 	mw := &msgWriter{
-		c:  c,
-		mu: newMu(c),
+		c:        c,
+		mu:       newMu(c),
+		activeMu: newMu(c),
 	}
 	return mw
 }
@@ -147,6 +149,12 @@ func (mw *msgWriter) returnFlateWriter() {
 func (mw *msgWriter) Write(p []byte) (_ int, err error) {
 	defer errd.Wrap(&err, "failed to write")
 
+	err = mw.activeMu.Lock(mw.ctx)
+	if err != nil {
+		return 0, err
+	}
+	defer mw.activeMu.Unlock()
+
 	if mw.closed {
 		return 0, xerrors.New("cannot use closed writer")
 	}
@@ -173,6 +181,12 @@ func (mw *msgWriter) write(p []byte) (int, error) {
 func (mw *msgWriter) Close() (err error) {
 	defer errd.Wrap(&err, "failed to close writer")
 
+	err = mw.activeMu.Lock(mw.ctx)
+	if err != nil {
+		return err
+	}
+	defer mw.activeMu.Unlock()
+
 	if mw.closed {
 		return xerrors.New("cannot use closed writer")
 	}
@@ -201,7 +215,7 @@ func (mw *msgWriter) Close() (err error) {
 }
 
 func (mw *msgWriter) close() {
-	mw.mu.Lock(context.Background())
+	mw.activeMu.Lock(context.Background())
 	mw.returnFlateWriter()
 }