diff --git a/client.go b/client.go index 482a68e3d..532476684 100644 --- a/client.go +++ b/client.go @@ -386,25 +386,36 @@ 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() + 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 { + c.streamLock.Lock() + defer c.streamLock.Unlock() - c.sendLock.Lock() - defer c.sendLock.Unlock() - c.streamLock.Unlock() + 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)