From 30af872b011067e8cc5978a4f30afdf0d45dc93a Mon Sep 17 00:00:00 2001 From: Steven Scott Date: Thu, 13 Sep 2018 08:13:24 -0700 Subject: [PATCH] Return write buffer to pool on write error Fix bug where connection did not return the write buffer to the pool after a write error. Add test for the same. Rename messsageWriter.fatal method to endMessage and consolidate all message cleanup code there. This ensures that the buffer is returned to pool on all code paths. Rename Conn.prepMessage to beginMessage for symmetry with endMessage. Move some duplicated code at calls to prepMessage to beginMessage. Bonus improvement: Adjust message and buffer size in TestWriteBufferPool to test that pool works with fragmented messages. --- conn.go | 47 ++++++++++++++++++++-------------------- conn_test.go | 60 +++++++++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 81 insertions(+), 26 deletions(-) diff --git a/conn.go b/conn.go index d2a21c14..7eb05cf7 100644 --- a/conn.go +++ b/conn.go @@ -451,7 +451,8 @@ func (c *Conn) WriteControl(messageType int, data []byte, deadline time.Time) er return err } -func (c *Conn) prepWrite(messageType int) error { +// beginMessage prepares a connection and message writer for a new message. +func (c *Conn) beginMessage(mw *messageWriter, messageType int) error { // Close previous writer if not already closed by the application. It's // probably better to return an error in this situation, but we cannot // change this without breaking existing applications. @@ -471,6 +472,10 @@ func (c *Conn) prepWrite(messageType int) error { return err } + mw.c = c + mw.frameType = messageType + mw.pos = maxFrameHeaderSize + if c.writeBuf == nil { wpd, ok := c.writePool.Get().(writePoolData) if ok { @@ -491,16 +496,11 @@ func (c *Conn) prepWrite(messageType int) error { // All message types (TextMessage, BinaryMessage, CloseMessage, PingMessage and // PongMessage) are supported. func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) { - if err := c.prepWrite(messageType); err != nil { + var mw messageWriter + if err := c.beginMessage(&mw, messageType); err != nil { return nil, err } - - mw := &messageWriter{ - c: c, - frameType: messageType, - pos: maxFrameHeaderSize, - } - c.writer = mw + c.writer = &mw if c.newCompressionWriter != nil && c.enableWriteCompression && isData(messageType) { w := c.newCompressionWriter(c.writer, c.compressionLevel) mw.compress = true @@ -517,10 +517,16 @@ type messageWriter struct { err error } -func (w *messageWriter) fatal(err error) error { +func (w *messageWriter) endMessage(err error) error { if w.err != nil { - w.err = err - w.c.writer = nil + return err + } + c := w.c + w.err = err + c.writer = nil + if c.writePool != nil { + c.writePool.Put(writePoolData{buf: c.writeBuf}) + c.writeBuf = nil } return err } @@ -534,7 +540,7 @@ func (w *messageWriter) flushFrame(final bool, extra []byte) error { // Check for invalid control frames. if isControl(w.frameType) && (!final || length > maxControlFramePayloadSize) { - return w.fatal(errInvalidControlFrame) + return w.endMessage(errInvalidControlFrame) } b0 := byte(w.frameType) @@ -579,7 +585,7 @@ func (w *messageWriter) flushFrame(final bool, extra []byte) error { copy(c.writeBuf[maxFrameHeaderSize-4:], key[:]) maskBytes(key, 0, c.writeBuf[maxFrameHeaderSize:w.pos]) if len(extra) > 0 { - return c.writeFatal(errors.New("websocket: internal error, extra used in client mode")) + return w.endMessage(c.writeFatal(errors.New("websocket: internal error, extra used in client mode"))) } } @@ -600,15 +606,11 @@ func (w *messageWriter) flushFrame(final bool, extra []byte) error { c.isWriting = false if err != nil { - return w.fatal(err) + return w.endMessage(err) } if final { - c.writer = nil - if c.writePool != nil { - c.writePool.Put(writePoolData{buf: c.writeBuf}) - c.writeBuf = nil - } + w.endMessage(errWriteClosed) return nil } @@ -709,7 +711,6 @@ func (w *messageWriter) Close() error { if err := w.flushFrame(true, nil); err != nil { return err } - w.err = errWriteClosed return nil } @@ -742,10 +743,10 @@ func (c *Conn) WriteMessage(messageType int, data []byte) error { if c.isServer && (c.newCompressionWriter == nil || !c.enableWriteCompression) { // Fast path with no allocations and single frame. - if err := c.prepWrite(messageType); err != nil { + var mw messageWriter + if err := c.beginMessage(&mw, messageType); err != nil { return err } - mw := messageWriter{c: c, frameType: messageType, pos: maxFrameHeaderSize} n := copy(c.writeBuf[mw.pos:], data) mw.pos += n data = data[n:] diff --git a/conn_test.go b/conn_test.go index 5ca06738..ad906b14 100644 --- a/conn_test.go +++ b/conn_test.go @@ -196,11 +196,16 @@ func (p *simpleBufferPool) Put(v interface{}) { } func TestWriteBufferPool(t *testing.T) { + const message = "Now is the time for all good people to come to the aid of the party." + var buf bytes.Buffer var pool simpleBufferPool - wc := newConn(fakeNetConn{Writer: &buf}, true, 1024, 1024, &pool, nil, nil) rc := newTestConn(&buf, nil, false) + // Specify writeBufferSize smaller than message size to ensure that pooling + // works with fragmented messages. + wc := newConn(fakeNetConn{Writer: &buf}, true, 1024, len(message)-1, &pool, nil, nil) + if wc.writeBuf != nil { t.Fatal("writeBuf not nil after create") } @@ -218,8 +223,6 @@ func TestWriteBufferPool(t *testing.T) { writeBufAddr := &wc.writeBuf[0] - const message = "Hello World!" - if _, err := io.WriteString(w, message); err != nil { t.Fatalf("io.WriteString(w, message) returned %v", err) } @@ -269,6 +272,7 @@ func TestWriteBufferPool(t *testing.T) { } } +// TestWriteBufferPoolSync ensures that *sync.Pool works as a buffer pool. func TestWriteBufferPoolSync(t *testing.T) { var buf bytes.Buffer var pool sync.Pool @@ -290,6 +294,56 @@ func TestWriteBufferPoolSync(t *testing.T) { } } +// errorWriter is an io.Writer than returns an error on all writes. +type errorWriter struct{} + +func (ew errorWriter) Write(p []byte) (int, error) { return 0, errors.New("Error!") } + +// TestWriteBufferPoolError ensures that buffer is returned to pool after error +// on write. +func TestWriteBufferPoolError(t *testing.T) { + + // Part 1: Test NextWriter/Write/Close + + var pool simpleBufferPool + wc := newConn(fakeNetConn{Writer: errorWriter{}}, true, 1024, 1024, &pool, nil, nil) + + w, err := wc.NextWriter(TextMessage) + if err != nil { + t.Fatalf("wc.NextWriter() returned %v", err) + } + + if wc.writeBuf == nil { + t.Fatal("writeBuf is nil after NextWriter") + } + + writeBufAddr := &wc.writeBuf[0] + + if _, err := io.WriteString(w, "Hello"); err != nil { + t.Fatalf("io.WriteString(w, message) returned %v", err) + } + + if err := w.Close(); err == nil { + t.Fatalf("w.Close() did not return error") + } + + if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr { + t.Fatal("writeBuf not returned to pool") + } + + // Part 2: Test WriteMessage + + wc = newConn(fakeNetConn{Writer: errorWriter{}}, true, 1024, 1024, &pool, nil, nil) + + if err := wc.WriteMessage(TextMessage, []byte("Hello")); err == nil { + t.Fatalf("wc.WriteMessage did not return error") + } + + if wpd, ok := pool.v.(writePoolData); !ok || len(wpd.buf) == 0 || &wpd.buf[0] != writeBufAddr { + t.Fatal("writeBuf not returned to pool") + } +} + func TestCloseFrameBeforeFinalMessageFrame(t *testing.T) { const bufSize = 512