feat(wal): Benchmark and improve WAL writes using Reset.

pkg/pattern/iter/batch_test.go

@@ -1,6 +1,7 @@
 package iter
 
 import (
+	"sort"
 	"testing"
 
 	"github.com/go-kit/log"
@@ -211,6 +212,12 @@ func TestReadMetricsBatch(t *testing.T) {
 			it := NewSumMergeSampleIterator(tt.seriesIter)
 			got, err := ReadMetricsBatch(it, tt.batchSize, log.NewNopLogger())
 			require.NoError(t, err)
+			sort.Slice(tt.expected.Series, func(i, j int) bool {
+				return tt.expected.Series[i].Labels < tt.expected.Series[j].Labels
+			})
+			sort.Slice(got.Series, func(i, j int) bool {
+				return got.Series[i].Labels < got.Series[j].Labels
+			})
 			require.Equal(t, tt.expected.Series, got.Series)
 		})
 	}

pkg/storage/wal/chunks/chunks_test.go

@@ -1,18 +1,16 @@
 package chunks
 
 import (
-	"bufio"
 	"bytes"
 	"fmt"
-	"os"
-	"path/filepath"
 	"strconv"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/loki/v3/pkg/logproto"
+	"github.com/grafana/loki/v3/pkg/storage/wal/testdata"
 )
 
 func TestChunkReaderWriter(t *testing.T) {
@@ -121,11 +119,11 @@ func TestChunkReaderWriter(t *testing.T) {
 }
 
 func TestChunkReaderWriterWithLogGenerator(t *testing.T) {
-	filenames := testDataFile()
+	filenames := testdata.Files()
 
 	for _, filename := range filenames {
 		t.Run(filename, func(t *testing.T) {
-			gen := newLogGenerator(t, filename)
+			gen := testdata.NewLogGenerator(t, filename)
 			defer gen.Close()
 
 			var entries []*logproto.Entry
@@ -196,10 +194,10 @@ var (
 
 // Benchmark reads with log generator
 func BenchmarkReadChunkWithLogGenerator(b *testing.B) {
-	filenames := testDataFile()
+	filenames := testdata.Files()
 	for _, filename := range filenames {
 		b.Run(filename, func(b *testing.B) {
-			gen := newLogGenerator(b, filename)
+			gen := testdata.NewLogGenerator(b, filename)
 			defer gen.Close()
 
 			var entries []*logproto.Entry
@@ -239,12 +237,12 @@ func BenchmarkReadChunkWithLogGenerator(b *testing.B) {
 
 // Benchmark with log generator
 func BenchmarkWriteChunkWithLogGenerator(b *testing.B) {
-	filenames := testDataFile()
+	filenames := testdata.Files()
 
 	for _, filename := range filenames {
 		for _, count := range []int{1000, 10000, 100000} {
 			b.Run(fmt.Sprintf("%s-%d", filename, count), func(b *testing.B) {
-				gen := newLogGenerator(b, filename)
+				gen := testdata.NewLogGenerator(b, filename)
 				defer gen.Close()
 
 				var entries []*logproto.Entry
@@ -278,24 +276,6 @@ func BenchmarkWriteChunkWithLogGenerator(b *testing.B) {
 	}
 }
 
-func testDataFile() []string {
-	testdataDir := "../testdata"
-	files, err := os.ReadDir(testdataDir)
-	if err != nil {
-		panic(err)
-	}
-
-	var fileNames []string
-	for _, file := range files {
-		if !file.IsDir() {
-			filePath := filepath.Join(testdataDir, file.Name())
-			fileNames = append(fileNames, filePath)
-		}
-	}
-
-	return fileNames
-}
-
 // generateLogEntries generates a slice of logproto.Entry with the given count.
 func generateLogEntries(count int) []*logproto.Entry {
 	entries := make([]*logproto.Entry, count)
@@ -307,39 +287,3 @@ func generateLogEntries(count int) []*logproto.Entry {
 	}
 	return entries
 }
-
-type logGenerator struct {
-	f *os.File
-	s *bufio.Scanner
-}
-
-func (g *logGenerator) Next() (bool, []byte) {
-	if g.s.Scan() {
-		return true, g.s.Bytes()
-	}
-	g.reset()
-	return g.s.Scan(), g.s.Bytes()
-}
-
-func (g *logGenerator) Close() {
-	if g.f != nil {
-		g.f.Close()
-	}
-	g.f = nil
-}
-
-func (g *logGenerator) reset() {
-	_, _ = g.f.Seek(0, 0)
-	g.s = bufio.NewScanner(g.f)
-}
-
-func newLogGenerator(t testing.TB, filename string) *logGenerator {
-	t.Helper()
-	file, err := os.Open(filename)
-	require.NoError(t, err)
-
-	return &logGenerator{
-		f: file,
-		s: bufio.NewScanner(file),
-	}
-}

pkg/storage/wal/index/buffer.go

@@ -86,6 +86,11 @@ func (fw *BufferWriter) Close() error {
 	return nil
 }
 
+func (fw *BufferWriter) Reset() {
+	fw.pos = 0
+	fw.buf.Reset()
+}
+
 func (fw *BufferWriter) Remove() error {
 	return nil
 }

pkg/storage/wal/index/index.go

@@ -118,8 +118,6 @@ type PostingsEncoder func(*encoding.Encbuf, []uint32) error
 // Writer implements the IndexWriter interface for the standard
 // serialization format.
 type Writer struct {
-	ctx context.Context
-
 	// For the main index file.
 	f *BufferWriter
 
@@ -197,9 +195,8 @@ func NewTOCFromByteSlice(bs ByteSlice) (*TOC, error) {
 
 // NewWriter returns a new Writer to the given filename. It serializes data in format version 2.
 // It uses the given encoder to encode each postings list.
-func NewWriterWithEncoder(ctx context.Context, encoder PostingsEncoder) (*Writer, error) {
+func NewWriterWithEncoder(encoder PostingsEncoder) (*Writer, error) {
 	iw := &Writer{
-		ctx:   ctx,
 		f:     NewBufferWriter(),
 		fP:    NewBufferWriter(),
 		fPO:   NewBufferWriter(),
@@ -222,8 +219,8 @@ func NewWriterWithEncoder(ctx context.Context, encoder PostingsEncoder) (*Writer
 
 // NewWriter creates a new index writer using the default encoder. See
 // NewWriterWithEncoder.
-func NewWriter(ctx context.Context) (*Writer, error) {
-	return NewWriterWithEncoder(ctx, EncodePostingsRaw)
+func NewWriter() (*Writer, error) {
+	return NewWriterWithEncoder(EncodePostingsRaw)
 }
 
 func (w *Writer) write(bufs ...[]byte) error {
@@ -242,15 +239,36 @@ func (w *Writer) Buffer() ([]byte, io.Closer, error) {
 	return w.f.Buffer()
 }
 
+func (w *Writer) Reset() error {
+	w.f.Reset()
+	w.fP.Reset()
+	w.fPO.Reset()
+	w.buf1.Reset()
+	w.buf2.Reset()
+	w.stage = idxStageNone
+	w.toc = TOC{}
+	w.postingsStart = 0
+	w.numSymbols = 0
+	w.symbols = nil
+	w.symbolFile = nil
+	w.lastSymbol = ""
+	w.symbolCache = make(map[string]symbolCacheEntry, 1<<8)
+	w.labelIndexes = w.labelIndexes[:0]
+	w.labelNames = make(map[string]uint64, 1<<8)
+	w.lastSeries = nil
+	w.lastSeriesRef = 0
+	w.lastChunkRef = 0
+	w.cntPO = 0
+	w.crc32.Reset()
+	if err := w.writeMeta(); err != nil {
+		return err
+	}
+	return nil
+}
+
 // ensureStage handles transitions between write stages and ensures that IndexWriter
 // methods are called in an order valid for the implementation.
 func (w *Writer) ensureStage(s indexWriterStage) error {
-	select {
-	case <-w.ctx.Done():
-		return w.ctx.Err()
-	default:
-	}
-
 	if w.stage == s {
 		return nil
 	}
@@ -691,7 +709,6 @@ func (w *Writer) writePostingsOffsetTable() error {
 	if err := w.fPO.Remove(); err != nil {
 		return err
 	}
-	w.fPO = nil
 
 	err = w.writeLengthAndHash(startPos)
 	if err != nil {
@@ -854,11 +871,7 @@ func (w *Writer) writePostingsToTmpFiles() error {
 				}
 			}
 		}
-		select {
-		case <-w.ctx.Done():
-			return w.ctx.Err()
-		default:
-		}
+
 	}
 	return nil
 }
@@ -936,7 +949,6 @@ func (w *Writer) writePostings() error {
 	if err := w.fP.Remove(); err != nil {
 		return err
 	}
-	w.fP = nil
 	return nil
 }
 

pkg/storage/wal/index/index_test.go

@@ -135,7 +135,7 @@ func (m mockIndex) Series(ref storage.SeriesRef, builder *labels.ScratchBuilder,
 
 func TestIndexRW_Create_Open(t *testing.T) {
 	// An empty index must still result in a readable file.
-	iw, err := NewWriter(context.Background())
+	iw, err := NewWriter()
 	require.NoError(t, err)
 	require.NoError(t, iw.Close())
 
@@ -160,7 +160,7 @@ func TestIndexRW_Postings(t *testing.T) {
 			labels: labels.FromStrings("a", "1", "b", strconv.Itoa(i)),
 		})
 	}
-	ir, buf, _ := createReader(ctx, t, input)
+	ir, buf, _ := createReader(t, input)
 
 	p, err := ir.Postings(ctx, "a", "1")
 	require.NoError(t, err)
@@ -271,7 +271,7 @@ func TestPostingsMany(t *testing.T) {
 			labels: labels.FromStrings("i", v, "foo", "bar"),
 		})
 	}
-	ir, _, symbols := createReader(ctx, t, input)
+	ir, _, symbols := createReader(t, input)
 
 	cases := []struct {
 		in []string
@@ -353,7 +353,7 @@ func TestPersistence_index_e2e(t *testing.T) {
 		})
 	}
 
-	ir, _, _ := createReader(ctx, t, input)
+	ir, _, _ := createReader(t, input)
 
 	// Population procedure as done by compaction.
 	var (
@@ -435,7 +435,7 @@ func TestPersistence_index_e2e(t *testing.T) {
 }
 
 func TestWriter_ShouldReturnErrorOnSeriesWithDuplicatedLabelNames(t *testing.T) {
-	w, err := NewWriter(context.Background())
+	w, err := NewWriter()
 	require.NoError(t, err)
 
 	require.NoError(t, w.AddSymbol("__name__"))
@@ -523,7 +523,7 @@ func BenchmarkReader_ShardedPostings(b *testing.B) {
 			labels: labels.FromStrings("const", fmt.Sprintf("%10d", 1), "unique", fmt.Sprintf("%10d", i)),
 		})
 	}
-	ir, _, _ := createReader(ctx, b, input)
+	ir, _, _ := createReader(b, input)
 	b.ResetTimer()
 
 	for n := 0; n < b.N; n++ {
@@ -540,7 +540,7 @@ func TestDecoder_Postings_WrongInput(t *testing.T) {
 }
 
 func TestChunksRefOrdering(t *testing.T) {
-	idx, err := NewWriter(context.Background())
+	idx, err := NewWriter()
 	require.NoError(t, err)
 
 	require.NoError(t, idx.AddSymbol("1"))
@@ -558,7 +558,7 @@ func TestChunksRefOrdering(t *testing.T) {
 }
 
 func TestChunksTimeOrdering(t *testing.T) {
-	idx, err := NewWriter(context.Background())
+	idx, err := NewWriter()
 	require.NoError(t, err)
 
 	require.NoError(t, idx.AddSymbol("1"))
@@ -585,10 +585,10 @@ func TestChunksTimeOrdering(t *testing.T) {
 
 // createFileReader creates a temporary index file. It writes the provided input to this file.
 // It returns a Reader for this file, the file's name, and the symbol map.
-func createReader(ctx context.Context, tb testing.TB, input indexWriterSeriesSlice) (*Reader, []byte, map[string]struct{}) {
+func createReader(tb testing.TB, input indexWriterSeriesSlice) (*Reader, []byte, map[string]struct{}) {
 	tb.Helper()
 
-	iw, err := NewWriter(ctx)
+	iw, err := NewWriter()
 	require.NoError(tb, err)
 
 	symbols := map[string]struct{}{}