feat: reuse buffer memory on Writer.Flush (#252)

bench_test.go

@@ -1,6 +1,7 @@
 package avro_test
 
 import (
+	"io"
 	"os"
 	"testing"
 
@@ -142,3 +143,32 @@ func BenchmarkSuperheroGenericEncode(b *testing.B) {
 		_, _ = avro.Marshal(schema, super)
 	}
 }
+
+func BenchmarkSuperheroWriteFlush(b *testing.B) {
+	schema, err := avro.ParseFiles("testdata/superhero.avsc")
+	if err != nil {
+		panic(err)
+	}
+
+	super := &Superhero{
+		ID:            234765,
+		AffiliationID: 9867,
+		Name:          "Wolverine",
+		Life:          85.25,
+		Energy:        32.75,
+		Powers: []*Superpower{
+			{ID: 2345, Name: "Bone Claws", Damage: 5, Energy: 1.15, Passive: false},
+			{ID: 2346, Name: "Regeneration", Damage: -2, Energy: 0.55, Passive: true},
+			{ID: 2347, Name: "Adamant skeleton", Damage: -10, Energy: 0, Passive: true},
+		},
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	w := avro.NewWriter(io.Discard, 128)
+	for i := 0; i < b.N; i++ {
+		w.WriteVal(schema, super)
+		_ = w.Flush()
+	}
+}

writer.go

@@ -66,13 +66,18 @@ func (w *Writer) Flush() error {
 	}
 
 	n, err := w.out.Write(w.buf)
+	if n < len(w.buf) && err == nil {
+		err = io.ErrShortWrite
+	}
 	if err != nil {
 		if w.Error == nil {
 			w.Error = err
 		}
 		return err
 	}
-	w.buf = w.buf[n:]
+
+	w.buf = w.buf[:0]
+
 	return nil
 }
 

writer_test.go

@@ -3,6 +3,7 @@ package avro_test
 import (
 	"bytes"
 	"errors"
+	"io"
 	"testing"
 
 	"github.com/hamba/avro/v2"
@@ -41,6 +42,19 @@ func TestWriter_Flush(t *testing.T) {
 	assert.Equal(t, []byte("test"), buf.Bytes())
 }
 
+type shortWrite struct{}
+
+func (shortWrite) Write(p []byte) (n int, err error) { return len(p) - 1, nil } // short write (breaks io.Writer contract; see io.ErrShortWrite)
+
+func TestWriter_FlushShortWrite(t *testing.T) {
+	w := avro.NewWriter(shortWrite{}, 10)
+	_, _ = w.Write([]byte("test"))
+
+	err := w.Flush()
+
+	require.ErrorIs(t, err, io.ErrShortWrite)
+}
+
 func TestWriter_FlushNoWriter(t *testing.T) {
 	w := avro.NewWriter(nil, 10)
 	_, _ = w.Write([]byte("test"))
@@ -70,6 +84,20 @@ func TestWriter_FlushReturnsUnderlyingWriterError(t *testing.T) {
 	assert.Error(t, w.Error)
 }
 
+func TestWriter_FlushReuseMemory(t *testing.T) {
+	var buf bytes.Buffer
+	w := avro.NewWriter(&buf, 10)
+	_, _ = w.Write(bytes.Repeat([]byte("test"), 10))
+
+	bufAddr := &w.Buffer()[:1][0] // buffer alloc address
+	err := w.Flush()
+	bufAddr2 := &w.Buffer()[:1][0]
+
+	require.NoError(t, err)
+
+	assert.Equal(t, bufAddr, bufAddr2)
+}
+
 func TestWriter_Write(t *testing.T) {
 	w := avro.NewWriter(nil, 50)