chore: cleanup loki imports (#4930)

* vendor wal reader that is only used in tests.

* vendor enc and dec wrappers

* move to own file and add comment

* add comment

* move comment

* fix

* fix

Author: kalleep

internal/component/common/loki/wal/buf.go

@@ -0,0 +1,71 @@
+package wal
+
+// KEEP IN SYNC WITH:
+// https://github.com/grafana/loki/blob/main/pkg/util/encoding/encoding.go
+// Local modifications should be minimized.
+
+import (
+	"encoding/binary"
+	"hash/crc32"
+
+	"github.com/prometheus/prometheus/tsdb/encoding"
+)
+
+func encWith(b []byte) (res encbuf) {
+	res.B = b
+	return res
+}
+
+// encbuf extends encoding.Encbuf with support for multi byte encoding
+type encbuf struct {
+	encoding.Encbuf
+}
+
+func (e *encbuf) PutString(s string) { e.B = append(e.B, s...) }
+
+func (e *encbuf) Skip(i int) {
+	e.B = e.B[:len(e.B)+i]
+}
+
+func decWith(b []byte) (res decbuf) {
+	res.B = b
+	return res
+}
+
+// decbuf extends encoding.Decbuf with support for multi byte decoding
+type decbuf struct {
+	encoding.Decbuf
+}
+
+func (d *decbuf) Bytes(n int) []byte {
+	if d.E != nil {
+		return nil
+	}
+	if len(d.B) < n {
+		d.E = encoding.ErrInvalidSize
+		return nil
+	}
+	x := d.B[:n]
+	d.B = d.B[n:]
+	return x
+}
+
+func (d *decbuf) CheckCrc(castagnoliTable *crc32.Table) error {
+	if d.E != nil {
+		return d.E
+	}
+	if len(d.B) < 4 {
+		d.E = encoding.ErrInvalidSize
+		return d.E
+	}
+
+	offset := len(d.B) - 4
+	expCRC := binary.BigEndian.Uint32(d.B[offset:])
+	d.B = d.B[:offset]
+
+	if d.Crc32(castagnoliTable) != expCRC {
+		d.E = encoding.ErrInvalidChecksum
+		return d.E
+	}
+	return nil
+}

internal/component/common/loki/wal/encoding.go

@@ -5,11 +5,9 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/grafana/loki/pkg/push"
 	"github.com/prometheus/prometheus/tsdb/chunks"
 	"github.com/prometheus/prometheus/tsdb/record"
-
-	"github.com/grafana/loki/pkg/push"
-	"github.com/grafana/loki/v3/pkg/util/encoding"
 )
 
 // RecordType represents the type of the WAL/Checkpoint record.
@@ -82,7 +80,7 @@ type RefEntries struct {
 }
 
 func (r *Record) EncodeSeries(b []byte) []byte {
-	buf := encoding.EncWith(b)
+	buf := encWith(b)
 	buf.PutByte(byte(WALRecordSeries))
 	buf.PutUvarintStr(r.UserID)
 
@@ -96,7 +94,7 @@ func (r *Record) EncodeSeries(b []byte) []byte {
 }
 
 func (r *Record) EncodeEntries(version RecordType, b []byte) []byte {
-	buf := encoding.EncWith(b)
+	buf := encWith(b)
 	buf.PutByte(byte(version))
 	buf.PutUvarintStr(r.UserID)
 
@@ -152,7 +150,7 @@ func DecodeEntries(b []byte, version RecordType, rec *Record) error {
 		return nil
 	}
 
-	dec := encoding.DecWith(b)
+	dec := decWith(b)
 	baseTime := dec.Be64int64()
 
 	for len(dec.B) > 0 && dec.Err() == nil {
@@ -220,7 +218,7 @@ func DecodeRecord(b []byte, walRec *Record) (err error) {
 		dec     record.Decoder
 		rSeries []record.RefSeries
 
-		decbuf = encoding.DecWith(b)
+		decbuf = decWith(b)
 		t      = RecordType(decbuf.Byte())
 	)
 

internal/component/common/loki/wal/reader.go

@@ -1,59 +1,45 @@
 package wal
 
-import (
-	"fmt"
-
-	"github.com/prometheus/common/model"
+// KEEP IN SYNC WITH:
+// https://github.com/grafana/loki/blob/main/pkg/util/wal/reader.go
+// Local modifications should be minimized.
 
-	"github.com/grafana/alloy/internal/loki/util"
-	walUtils "github.com/grafana/loki/v3/pkg/util/wal"
+import (
+	"errors"
+	"io"
 
-	"github.com/grafana/alloy/internal/component/common/loki"
+	"github.com/prometheus/prometheus/tsdb/wlog"
 )
 
-// ReadWAL will read all entries in the WAL located under dir. Mainly used for testing
-func ReadWAL(dir string) ([]loki.Entry, error) {
-	reader, closeFn, err := walUtils.NewWalReader(dir, -1)
-	if err != nil {
-		return nil, err
-	}
-	defer func() { closeFn.Close() }()
-
-	seenSeries := make(map[uint64]model.LabelSet)
-	seenEntries := []loki.Entry{}
-
-	for reader.Next() {
-		var walRec = Record{}
-		bytes := reader.Record()
-		err = DecodeRecord(bytes, &walRec)
+func newWalReader(dir string, startSegment int) (*wlog.Reader, io.Closer, error) {
+	var (
+		segmentReader io.ReadCloser
+		err           error
+	)
+	if startSegment < 0 {
+		segmentReader, err = wlog.NewSegmentsReader(dir)
 		if err != nil {
-			return nil, fmt.Errorf("error decoding wal record: %w", err)
+			return nil, nil, err
 		}
-
-		// first read series
-		for _, series := range walRec.Series {
-			if _, ok := seenSeries[uint64(series.Ref)]; !ok {
-				seenSeries[uint64(series.Ref)] = util.MapToModelLabelSet(series.Labels.Map())
-			}
+	} else {
+		first, last, err := wlog.Segments(dir)
+		if err != nil {
+			return nil, nil, err
 		}
-
-		for _, entries := range walRec.RefEntries {
-			for _, entry := range entries.Entries {
-				labels, ok := seenSeries[uint64(entries.Ref)]
-				if !ok {
-					return nil, fmt.Errorf("found entry without matching series")
-				}
-				seenEntries = append(seenEntries, loki.Entry{
-					Labels: labels,
-					Entry:  entry,
-				})
-			}
+		if startSegment > last {
+			return nil, nil, errors.New("start segment is beyond the last WAL segment")
+		}
+		if first > startSegment {
+			startSegment = first
+		}
+		segmentReader, err = wlog.NewSegmentsRangeReader(wlog.SegmentRange{
+			Dir:   dir,
+			First: startSegment,
+			Last:  -1, // Till the end.
+		})
+		if err != nil {
+			return nil, nil, err
 		}
-
-		// reset entry
-		walRec.Series = walRec.Series[:]
-		walRec.RefEntries = walRec.RefEntries[:]
 	}
-
-	return seenEntries, nil
+	return wlog.NewReader(segmentReader), segmentReader, nil
 }

internal/component/common/loki/wal/writer_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/alloy/internal/component/common/loki"
+	"github.com/grafana/alloy/internal/loki/util"
 	"github.com/grafana/alloy/internal/runtime/logging/level"
 )
 
@@ -251,7 +252,7 @@ func eventuallyReadWAL(t *testing.T, expectedEntries int, dir string) []loki.Ent
 	var readEntries []loki.Entry
 	require.Eventually(t, func() bool {
 		// read WAL and assert over read entries
-		readEntries, err = ReadWAL(dir)
+		readEntries, err = readWAL(dir)
 		if err != nil {
 			return false
 		}
@@ -260,6 +261,49 @@ func eventuallyReadWAL(t *testing.T, expectedEntries int, dir string) []loki.Ent
 	return readEntries
 }
 
+// readWAL will read all entries in the WAL located under dir.
+func readWAL(dir string) ([]loki.Entry, error) {
+	reader, closeFn, err := newWalReader(dir, -1)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { closeFn.Close() }()
+
+	seenSeries := make(map[uint64]model.LabelSet)
+	seenEntries := []loki.Entry{}
+
+	for reader.Next() {
+		var walRec = Record{}
+		bytes := reader.Record()
+		err = DecodeRecord(bytes, &walRec)
+		if err != nil {
+			return nil, fmt.Errorf("error decoding wal record: %w", err)
+		}
+
+		// first read series
+		for _, series := range walRec.Series {
+			if _, ok := seenSeries[uint64(series.Ref)]; !ok {
+				seenSeries[uint64(series.Ref)] = util.MapToModelLabelSet(series.Labels.Map())
+			}
+		}
+
+		for _, entries := range walRec.RefEntries {
+			for _, entry := range entries.Entries {
+				labels, ok := seenSeries[uint64(entries.Ref)]
+				if !ok {
+					return nil, fmt.Errorf("found entry without matching series")
+				}
+				seenEntries = append(seenEntries, loki.Entry{
+					Labels: labels,
+					Entry:  entry,
+				})
+			}
+		}
+	}
+
+	return seenEntries, nil
+}
+
 func BenchmarkWriter_WriteEntries(b *testing.B) {
 	type testCase struct {
 		lines          int