Merge pull request #168 from kevpar/deadlock

kevpar · web-flow · commit ef5734239ebc · 2024-05-13T17:22:09.000-05:00
client: Fix deadlock when writing to pipe blocks
diff --git a/client.go b/client.go
@@ -386,25 +386,44 @@ func (c *Client) receiveLoop() error {
 // createStream creates a new stream and registers it with the client
 // Introduce stream types for multiple or single response
 func (c *Client) createStream(flags uint8, b []byte) (*stream, error) {
-	c.streamLock.Lock()
+	// sendLock must be held across both allocation of the stream ID and sending it across the wire.
+	// This ensures that new stream IDs sent on the wire are always increasing, which is a
+	// requirement of the TTRPC protocol.
+	// This use of sendLock could be split into another mutex that covers stream creation + first send,
+	// and just use sendLock to guard writing to the wire, but for now it seems simpler to have fewer mutexes.
+	c.sendLock.Lock()
+	defer c.sendLock.Unlock()
 
 	// Check if closed since lock acquired to prevent adding
 	// anything after cleanup completes
 	select {
 	case <-c.ctx.Done():
-		c.streamLock.Unlock()
 		return nil, ErrClosed
 	default:
 	}
 
-	// Stream ID should be allocated at same time
-	s := newStream(c.nextStreamID, c)
-	c.streams[s.id] = s
-	c.nextStreamID = c.nextStreamID + 2
+	var s *stream
+	if err := func() error {
+		// In the future this could be replaced with a sync.Map instead of streamLock+map.
+		c.streamLock.Lock()
+		defer c.streamLock.Unlock()
 
-	c.sendLock.Lock()
-	defer c.sendLock.Unlock()
-	c.streamLock.Unlock()
+		// Check if closed since lock acquired to prevent adding
+		// anything after cleanup completes
+		select {
+		case <-c.ctx.Done():
+			return ErrClosed
+		default:
+		}
+
+		s = newStream(c.nextStreamID, c)
+		c.streams[s.id] = s
+		c.nextStreamID = c.nextStreamID + 2
+
+		return nil
+	}(); err != nil {
+		return nil, err
+	}
 
 	if err := c.channel.send(uint32(s.id), messageTypeRequest, flags, b); err != nil {
 		return s, filterCloseErr(err)
