release buffer properly

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 {

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
 }