Skip to content

Commit

Permalink
Ensure flush happens and correctly lock connection for a series of un…
Browse files Browse the repository at this point in the history
…flushed writes (#161)

* Ensure flush happens even with errors

Fixes #160

* Use correct lock around unflushed frame writes

* formatting

* Move flushing to endSendUnflushed
  • Loading branch information
lukebakken authored Jan 31, 2023
1 parent beca949 commit 2d6096a
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 7 deletions.
16 changes: 12 additions & 4 deletions channel.go
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,18 @@ func (ch *Channel) sendOpen(msg message) (err error) {
return ch.sendClosed(msg)
}

ch.connection.startSendUnflushed()

// Flush the buffer only after all the Frames that comprise the Message
// have been written to maximise benefits of using a buffered writer.
defer func() {
if endError := ch.connection.endSendUnflushed(); endError != nil {
if err == nil {
err = endError
}
}
}()

// We use sendUnflushed() in this method as sending the message requires
// sending multiple Frames (methodFrame, headerFrame, N x bodyFrame).
// Flushing after each Frame is inefficient, as it negates much of the
Expand Down Expand Up @@ -275,10 +287,6 @@ func (ch *Channel) sendOpen(msg message) (err error) {
return
}
}

// Flush the buffer only after all the Frames that comprise the Message
// have been written to maximise benefits of using a buffered writer.
err = ch.connection.flush()
} else {
// If the channel is closed, use Channel.sendClosed()
if ch.IsClosed() {
Expand Down
17 changes: 14 additions & 3 deletions connection.go
Original file line number Diff line number Diff line change
Expand Up @@ -448,6 +448,19 @@ func (c *Connection) send(f frame) error {
return err
}

// This method is intended to be used with sendUnflushed() to start a sequence
// of sendUnflushed() calls
func (c *Connection) startSendUnflushed() {
c.sendM.Lock()
}

// This method is intended to be used with sendUnflushed() to end a sequence
// of sendUnflushed() calls and flush the connection
func (c *Connection) endSendUnflushed() error {
defer c.sendM.Unlock()
return c.flush()
}

// sendUnflushed performs an *Unflushed* write. It is otherwise equivalent to
// send(), and we provide a separate flush() function to explicitly flush the
// buffer after all Frames are written.
Expand All @@ -468,9 +481,7 @@ func (c *Connection) sendUnflushed(f frame) error {
return ErrClosed
}

c.sendM.Lock()
err := f.write(c.writer.w)
c.sendM.Unlock()
err := c.writer.WriteFrameNoFlush(f)

if err != nil {
// shutdown could be re-entrant from signaling notify chans
Expand Down
5 changes: 5 additions & 0 deletions write.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,11 @@ import (
"time"
)

func (w *writer) WriteFrameNoFlush(frame frame) (err error) {
err = frame.write(w.w)
return
}

func (w *writer) WriteFrame(frame frame) (err error) {
if err = frame.write(w.w); err != nil {
return
Expand Down

0 comments on commit 2d6096a

Please sign in to comment.