diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index 3c5352427..93eecfa39 100644 --- a/common/buf/multi_buffer.go +++ b/common/buf/multi_buffer.go @@ -1,9 +1,6 @@ package buf -import ( - "io" - "net" -) +import "io" type MultiBufferWriter interface { WriteMultiBuffer(MultiBuffer) (int, error) @@ -51,19 +48,6 @@ func (mb *MultiBuffer) Read(b []byte) (int, error) { return totalBytes, nil } -func (mb MultiBuffer) WriteTo(writer io.Writer) (int, error) { - if mw, ok := writer.(MultiBufferWriter); ok { - return mw.WriteMultiBuffer(mb) - } - bs := make([][]byte, len(mb)) - for i, b := range mb { - bs[i] = b.Bytes() - } - nbs := net.Buffers(bs) - nBytes, err := nbs.WriteTo(writer) - return int(nBytes), err -} - func (mb MultiBuffer) Len() int { size := 0 for _, b := range mb { diff --git a/common/buf/writer.go b/common/buf/writer.go index 84fa55ba6..141c8b354 100644 --- a/common/buf/writer.go +++ b/common/buf/writer.go @@ -1,6 +1,9 @@ package buf -import "io" +import ( + "io" + "net" +) // BufferToBytesWriter is a Writer that writes alloc.Buffer into underlying writer. type BufferToBytesWriter struct { @@ -8,9 +11,19 @@ type BufferToBytesWriter struct { } // Write implements Writer.Write(). Write() takes ownership of the given buffer. -func (v *BufferToBytesWriter) Write(buffer MultiBuffer) error { - _, err := buffer.WriteTo(v.writer) - //buffer.Release() +func (v *BufferToBytesWriter) Write(mb MultiBuffer) error { + if mw, ok := v.writer.(MultiBufferWriter); ok { + _, err := mw.WriteMultiBuffer(mb) + return err + } + + defer mb.Release() + bs := make([][]byte, len(mb)) + for i, b := range mb { + bs[i] = b.Bytes() + } + nbs := net.Buffers(bs) + _, err := nbs.WriteTo(v.writer) return err }