fix deadlock between onStreamCompleted and Session.Close

receive_stream.go

@@ -73,44 +73,52 @@ func (s *receiveStream) StreamID() protocol.StreamID {
 
 // Read implements io.Reader. It is not thread safe!
 func (s *receiveStream) Read(p []byte) (int, error) {
+	completed, n, err := s.readImpl(p)
+	if completed {
+		s.sender.onStreamCompleted(s.streamID)
+	}
+	return n, err
+}
+
+func (s *receiveStream) readImpl(p []byte) (bool /*stream completed */, int, error) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
 	if s.finRead {
-		return 0, io.EOF
+		return false, 0, io.EOF
 	}
 	if s.canceledRead {
-		return 0, s.cancelReadErr
+		return false, 0, s.cancelReadErr
 	}
 	if s.resetRemotely {
-		return 0, s.resetRemotelyErr
+		return false, 0, s.resetRemotelyErr
 	}
 	if s.closedForShutdown {
-		return 0, s.closeForShutdownErr
+		return false, 0, s.closeForShutdownErr
 	}
 
 	bytesRead := 0
 	for bytesRead < len(p) {
 		frame := s.frameQueue.Head()
 		if frame == nil && bytesRead > 0 {
-			return bytesRead, s.closeForShutdownErr
+			return false, bytesRead, s.closeForShutdownErr
 		}
 
 		for {
 			// Stop waiting on errors
 			if s.closedForShutdown {
-				return bytesRead, s.closeForShutdownErr
+				return false, bytesRead, s.closeForShutdownErr
 			}
 			if s.canceledRead {
-				return bytesRead, s.cancelReadErr
+				return false, bytesRead, s.cancelReadErr
 			}
 			if s.resetRemotely {
-				return bytesRead, s.resetRemotelyErr
+				return false, bytesRead, s.resetRemotelyErr
 			}
 
 			deadline := s.readDeadline
 			if !deadline.IsZero() && !time.Now().Before(deadline) {
-				return bytesRead, errDeadline
+				return false, bytesRead, errDeadline
 			}
 
 			if frame != nil {
@@ -132,10 +140,10 @@ func (s *receiveStream) Read(p []byte) (int, error) {
 		}
 
 		if bytesRead > len(p) {
-			return bytesRead, fmt.Errorf("BUG: bytesRead (%d) > len(p) (%d) in stream.Read", bytesRead, len(p))
+			return false, bytesRead, fmt.Errorf("BUG: bytesRead (%d) > len(p) (%d) in stream.Read", bytesRead, len(p))
 		}
 		if s.readPosInFrame > int(frame.DataLen()) {
-			return bytesRead, fmt.Errorf("BUG: readPosInFrame (%d) > frame.DataLen (%d) in stream.Read", s.readPosInFrame, frame.DataLen())
+			return false, bytesRead, fmt.Errorf("BUG: readPosInFrame (%d) > frame.DataLen (%d) in stream.Read", s.readPosInFrame, frame.DataLen())
 		}
 
 		s.mutex.Unlock()
@@ -158,12 +166,11 @@ func (s *receiveStream) Read(p []byte) (int, error) {
 			s.frameQueue.Pop()
 			s.finRead = frame.FinBit
 			if frame.FinBit {
-				s.sender.onStreamCompleted(s.streamID)
+				return true, bytesRead, io.EOF
-				return bytesRead, io.EOF
 			}
 		}
 	}
-	return bytesRead, nil
+	return false, bytesRead, nil
 }
 
 func (s *receiveStream) CancelRead(errorCode protocol.ApplicationErrorCode) error {
@@ -204,25 +211,33 @@ func (s *receiveStream) handleStreamFrame(frame *wire.StreamFrame) error {
 }
 
 func (s *receiveStream) handleRstStreamFrame(frame *wire.RstStreamFrame) error {
+	completed, err := s.handleRstStreamFrameImpl(frame)
+	if completed {
+		s.sender.onStreamCompleted(s.streamID)
+	}
+	return err
+}
+
+func (s *receiveStream) handleRstStreamFrameImpl(frame *wire.RstStreamFrame) (bool /*completed */, error) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
 	if s.closedForShutdown {
-		return nil
+		return false, nil
 	}
 	if err := s.flowController.UpdateHighestReceived(frame.ByteOffset, true); err != nil {
-		return err
+		return false, err
 	}
 	// In gQUIC, error code 0 has a special meaning.
 	// The peer will reliably continue transmitting, but is not interested in reading from the stream.
 	// We should therefore just continue reading from the stream, until we encounter the FIN bit.
 	if !s.version.UsesIETFFrameFormat() && frame.ErrorCode == 0 {
-		return nil
+		return false, nil
 	}
 
 	// ignore duplicate RST_STREAM frames for this stream (after checking their final offset)
 	if s.resetRemotely {
-		return nil
+		return false, nil
 	}
 	s.resetRemotely = true
 	s.resetRemotelyErr = streamCanceledError{
@@ -230,8 +245,7 @@ func (s *receiveStream) handleRstStreamFrame(frame *wire.RstStreamFrame) error {
 		error:     fmt.Errorf("Stream %d was reset with error code %d", s.streamID, frame.ErrorCode),
 	}
 	s.signalRead()
-	s.sender.onStreamCompleted(s.streamID)
+	return true, nil
-	return nil
 }
 
 func (s *receiveStream) CloseRemote(offset protocol.ByteCount) {

send_stream.go

@@ -133,11 +133,19 @@ func (s *sendStream) Write(p []byte) (int, error) {
 // popStreamFrame returns the next STREAM frame that is supposed to be sent on this stream
 // maxBytes is the maximum length this frame (including frame header) will have.
 func (s *sendStream) popStreamFrame(maxBytes protocol.ByteCount) (*wire.StreamFrame, bool /* has more data to send */) {
+	completed, frame, hasMoreData := s.popStreamFrameImpl(maxBytes)
+	if completed {
+		s.sender.onStreamCompleted(s.streamID)
+	}
+	return frame, hasMoreData
+}
+
+func (s *sendStream) popStreamFrameImpl(maxBytes protocol.ByteCount) (bool /* completed */, *wire.StreamFrame, bool /* has more data to send */) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
 	if s.closeForShutdownErr != nil {
-		return nil, false
+		return false, nil, false
 	}
 
 	frame := &wire.StreamFrame{
@@ -147,7 +155,7 @@ func (s *sendStream) popStreamFrame(maxBytes protocol.ByteCount) (*wire.StreamFr
 	}
 	maxDataLen := frame.MaxDataLen(maxBytes, s.version)
 	if maxDataLen == 0 { // a STREAM frame must have at least one byte of data
-		return nil, s.dataForWriting != nil
+		return false, nil, s.dataForWriting != nil
 	}
 	frame.Data, frame.FinBit = s.getDataForWriting(maxDataLen)
 	if len(frame.Data) == 0 && !frame.FinBit {
@@ -156,24 +164,24 @@ func (s *sendStream) popStreamFrame(maxBytes protocol.ByteCount) (*wire.StreamFr
 		// - there's data for writing, but the stream is stream-level flow control blocked
 		// - there's data for writing, but the stream is connection-level flow control blocked
 		if s.dataForWriting == nil {
-			return nil, false
+			return false, nil, false
 		}
 		isBlocked, _ := s.flowController.IsBlocked()
-		return nil, !isBlocked
+		return false, nil, !isBlocked
 	}
 	if frame.FinBit {
 		s.finSent = true
-		s.sender.onStreamCompleted(s.streamID)
+		return true, frame, s.dataForWriting != nil
 	} else if s.streamID != s.version.CryptoStreamID() { // TODO(#657): Flow control for the crypto stream
 		if isBlocked, offset := s.flowController.IsBlocked(); isBlocked {
 			s.sender.queueControlFrame(&wire.StreamBlockedFrame{
 				StreamID: s.streamID,
 				Offset:   offset,
 			})
-			return frame, false
+			return false, frame, false
 		}
 	}
-	return frame, s.dataForWriting != nil
+	return false, frame, s.dataForWriting != nil
 }
 
 func (s *sendStream) getDataForWriting(maxBytes protocol.ByteCount) ([]byte, bool /* should send FIN */) {
@@ -218,18 +226,22 @@ func (s *sendStream) Close() error {
 
 func (s *sendStream) CancelWrite(errorCode protocol.ApplicationErrorCode) error {
 	s.mutex.Lock()
-	defer s.mutex.Unlock()
+	completed, err := s.cancelWriteImpl(errorCode, fmt.Errorf("Write on stream %d canceled with error code %d", s.streamID, errorCode))
+	s.mutex.Unlock()
 
-	return s.cancelWriteImpl(errorCode, fmt.Errorf("Write on stream %d canceled with error code %d", s.streamID, errorCode))
+	if completed {
+		s.sender.onStreamCompleted(s.streamID)
+	}
+	return err
 }
 
 // must be called after locking the mutex
-func (s *sendStream) cancelWriteImpl(errorCode protocol.ApplicationErrorCode, writeErr error) error {
+func (s *sendStream) cancelWriteImpl(errorCode protocol.ApplicationErrorCode, writeErr error) (bool /*completed */, error) {
 	if s.canceledWrite {
-		return nil
+		return false, nil
 	}
 	if s.finishedWriting {
-		return fmt.Errorf("CancelWrite for closed stream %d", s.streamID)
+		return false, fmt.Errorf("CancelWrite for closed stream %d", s.streamID)
 	}
 	s.canceledWrite = true
 	s.cancelWriteErr = writeErr
@@ -241,14 +253,13 @@ func (s *sendStream) cancelWriteImpl(errorCode protocol.ApplicationErrorCode, wr
 	})
 	// TODO(#991): cancel retransmissions for this stream
 	s.ctxCancel()
-	s.sender.onStreamCompleted(s.streamID)
+	return true, nil
-	return nil
 }
 
 func (s *sendStream) handleStopSendingFrame(frame *wire.StopSendingFrame) {
-	s.mutex.Lock()
+	if completed := s.handleStopSendingFrameImpl(frame); completed {
-	defer s.mutex.Unlock()
+		s.sender.onStreamCompleted(s.streamID)
-	s.handleStopSendingFrameImpl(frame)
+	}
 }
 
 func (s *sendStream) handleMaxStreamDataFrame(frame *wire.MaxStreamDataFrame) {
@@ -261,7 +272,10 @@ func (s *sendStream) handleMaxStreamDataFrame(frame *wire.MaxStreamDataFrame) {
 }
 
 // must be called after locking the mutex
-func (s *sendStream) handleStopSendingFrameImpl(frame *wire.StopSendingFrame) {
+func (s *sendStream) handleStopSendingFrameImpl(frame *wire.StopSendingFrame) bool /*completed*/ {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
 	writeErr := streamCanceledError{
 		errorCode: frame.ErrorCode,
 		error:     fmt.Errorf("Stream %d was reset with error code %d", s.streamID, frame.ErrorCode),
@@ -270,7 +284,8 @@ func (s *sendStream) handleStopSendingFrameImpl(frame *wire.StopSendingFrame) {
 	if !s.version.UsesIETFFrameFormat() {
 		errorCode = errorCodeStoppingGQUIC
 	}
-	s.cancelWriteImpl(errorCode, writeErr)
+	completed, _ := s.cancelWriteImpl(errorCode, writeErr)
+	return completed
 }
 
 func (s *sendStream) Context() context.Context {

stream.go

@@ -19,6 +19,7 @@ const (
 type streamSender interface {
 	queueControlFrame(wire.Frame)
 	onHasStreamData(protocol.StreamID)
+	// must be called without holding the mutex that is acquired by closeForShutdown
 	onStreamCompleted(protocol.StreamID)
 }