chore: Use pool of bytes.Buffer instead of io.Pipe (#13543)

Author: grobinson-grafana

pkg/ingester-rf1/flush.go

@@ -1,6 +1,7 @@
 package ingesterrf1
 
 import (
+	"bytes"
 	"crypto/rand"
 	"errors"
 	"fmt"
@@ -12,12 +13,10 @@ import (
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/backoff"
 	"github.com/grafana/dskit/ring"
-	"github.com/grafana/dskit/runutil"
 	"github.com/oklog/ulid"
 	"golang.org/x/net/context"
 
 	"github.com/grafana/loki/v3/pkg/storage/wal"
-	util_log "github.com/grafana/loki/v3/pkg/util/log"
 )
 
 const (
@@ -40,6 +39,7 @@ const (
 func (i *Ingester) InitFlushWorkers() {
 	i.flushWorkersDone.Add(i.cfg.ConcurrentFlushes)
 	for j := 0; j < i.cfg.ConcurrentFlushes; j++ {
+		i.flushBuffers[j] = new(bytes.Buffer)
 		go i.flushWorker(j)
 	}
 }
@@ -90,7 +90,7 @@ func (i *Ingester) flushWorker(j int) {
 		n := it.Writer.InputSize()
 		humanized := humanize.Bytes(uint64(n))
 
-		err = i.flushItem(l, it)
+		err = i.flushItem(l, j, it)
 		d := time.Since(start)
 		if err != nil {
 			level.Error(l).Log("msg", "failed to flush", "size", humanized, "duration", d, "err", err)
@@ -103,13 +103,13 @@ func (i *Ingester) flushWorker(j int) {
 	}
 }
 
-func (i *Ingester) flushItem(l log.Logger, it *wal.PendingItem) error {
+func (i *Ingester) flushItem(l log.Logger, j int, it *wal.PendingItem) error {
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	defer cancelFunc()
 
 	b := backoff.New(ctx, i.cfg.FlushOpBackoff)
 	for b.Ongoing() {
-		err := i.flushSegment(ctx, it.Writer)
+		err := i.flushSegment(ctx, j, it.Writer)
 		if err == nil {
 			break
 		}
@@ -124,19 +124,23 @@ func (i *Ingester) flushItem(l log.Logger, it *wal.PendingItem) error {
 // If the flush is successful, metrics for this flush are to be reported.
 // If the flush isn't successful, the operation for this userID is requeued allowing this and all other unflushed
 // segments to have another opportunity to be flushed.
-func (i *Ingester) flushSegment(ctx context.Context, ch *wal.SegmentWriter) error {
+func (i *Ingester) flushSegment(ctx context.Context, j int, w *wal.SegmentWriter) error {
 	id := ulid.MustNew(ulid.Timestamp(time.Now()), rand.Reader)
-	r := ch.Reader()
 
 	start := time.Now()
 	defer func() {
-		runutil.CloseWithLogOnErr(util_log.Logger, r, "flushSegment")
 		i.metrics.flushDuration.Observe(time.Since(start).Seconds())
-		ch.Observe()
+		w.Observe()
 	}()
 
+	buf := i.flushBuffers[j]
+	defer buf.Reset()
+	if _, err := w.WriteTo(buf); err != nil {
+		return err
+	}
+
 	i.metrics.flushesTotal.Add(1)
-	if err := i.store.PutObject(ctx, fmt.Sprintf("loki-v2/wal/anon/"+id.String()), r); err != nil {
+	if err := i.store.PutObject(ctx, fmt.Sprintf("loki-v2/wal/anon/"+id.String()), buf); err != nil {
 		i.metrics.flushFailuresTotal.Inc()
 		return fmt.Errorf("store put chunk: %w", err)
 	}

pkg/ingester-rf1/ingester.go

@@ -1,6 +1,7 @@
 package ingesterrf1
 
 import (
+	"bytes"
 	"context"
 	"flag"
 	"fmt"
@@ -205,6 +206,7 @@ type Ingester struct {
 
 	// One queue per flush thread.  Fingerprint is used to
 	// pick a queue.
+	flushBuffers     []*bytes.Buffer
 	flushWorkersDone sync.WaitGroup
 
 	wal *wal.Manager

pkg/storage/wal/segment.go

@@ -299,32 +299,6 @@ func (b *SegmentWriter) Reset() {
 	b.inputSize.Store(0)
 }
 
-type EncodedSegmentReader struct {
-	*io.PipeReader
-	*io.PipeWriter
-}
-
-func (e *EncodedSegmentReader) Close() error {
-	err := e.PipeWriter.Close()
-	if err != nil {
-		return err
-	}
-	err = e.PipeReader.Close()
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
-func (b *SegmentWriter) Reader() io.ReadCloser {
-	pr, pw := io.Pipe()
-	go func() {
-		_, err := b.WriteTo(pw)
-		pw.CloseWithError(err)
-	}()
-	return &EncodedSegmentReader{PipeReader: pr, PipeWriter: pw}
-}
-
 // InputSize returns the total size of the input data written to the writer.
 // It doesn't account for timestamps and labels.
 func (b *SegmentWriter) InputSize() int64 {

pkg/storage/wal/segment_test.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"io"
 	"sort"
 	"sync"
 	"testing"
@@ -510,19 +509,4 @@ func BenchmarkWrites(b *testing.B) {
 			require.EqualValues(b, encodedLength, n)
 		}
 	})
-
-	bytesBuf := make([]byte, encodedLength)
-	b.Run("Reader", func(b *testing.B) {
-		b.ResetTimer()
-		b.ReportAllocs()
-		for i := 0; i < b.N; i++ {
-			var err error
-			reader := writer.Reader()
-
-			n, err := io.ReadFull(reader, bytesBuf)
-			require.NoError(b, err)
-			require.EqualValues(b, encodedLength, n)
-			require.NoError(b, reader.Close())
-		}
-	})
 }