chore(dataobj): return Object on Builder.Flush

This commit updates Builder.Flush to return an Object rather than
writing bytes to an argument. To support this change, dataobj.Object has
two additional methods:

- dataobj.Object.Size, to retrieve the full encoded size of the object
- dataobj.Object.Reader, to read from the encoded object

This API change allows downstream users of dataobj.Builder to stream the
encoded object for analysis or uploading, without requiring the caller
to retain the entire object in memory at once.

In the future, this will be used to reduce the memory overhead of
dataobj-consumers by buffering flushed sections on an ephemeral disk
instead of in memory.

Author: rfratto

pkg/dataobj/builder.go

@@ -3,6 +3,9 @@ package dataobj
 import (
 	"bytes"
 	"fmt"
+	"io"
+
+	"github.com/grafana/loki/v3/pkg/dataobj/internal/util/bufpool"
 )
 
 // A Builder builds data objects from a set of incoming log data. Log data is
@@ -56,22 +59,48 @@ func (b *Builder) Bytes() int {
 	return b.encoder.Bytes()
 }
 
-// Flush flushes all buffered data to the buffer provided. Calling Flush can
-// result in a no-op if there is no buffered data to flush.
+// Flush constructs a new Object from the accumulated sections. Allocated
+// resources for the Object must be released by calling Close on the returned
+// io.Closer. After closing, the returned Object must no longer be read.
 //
-// [Builder.Reset] is called after a successful Flush to discard any pending
-// data and allow new data to be appended.
-func (b *Builder) Flush(output *bytes.Buffer) (int64, error) {
-	sz, err := b.encoder.Flush(output)
+// Flush returns an error if the object could not be constructed.
+// [Builder.Reset] is called after a successful flush to discard any pending
+// data, allowing new data to be appended.
+func (b *Builder) Flush() (*Object, io.Closer, error) {
+	flushSize, err := b.encoder.FlushSize()
+	if err != nil {
+		return nil, nil, fmt.Errorf("determining object size: %w", err)
+	}
+
+	buf := bufpool.Get(int(flushSize))
+
+	closer := func() error {
+		bufpool.Put(buf)
+		return nil
+	}
+
+	sz, err := b.encoder.Flush(buf)
 	if err != nil {
-		return sz, fmt.Errorf("building object: %w", err)
+		return nil, nil, fmt.Errorf("flushing object: %w", err)
+	}
+
+	obj, err := FromReaderAt(bytes.NewReader(buf.Bytes()), sz)
+	if err != nil {
+		bufpool.Put(buf)
+		return nil, nil, fmt.Errorf("error building object: %w", err)
 	}
 
 	b.Reset()
-	return sz, nil
+	return obj, funcIOCloser(closer), nil
 }
 
 // Reset discards pending data and resets the builder to an empty state.
 func (b *Builder) Reset() {
 	b.encoder.Reset()
 }
+
+type funcIOCloser func() error
+
+func (fc funcIOCloser) Close() error {
+	return fc()
+}

pkg/dataobj/consumer/logsobj/builder.go

@@ -2,11 +2,11 @@
 package logsobj
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"flag"
 	"fmt"
+	"io"
 	"time"
 
 	"github.com/grafana/dskit/flagext"
@@ -293,27 +293,24 @@ func (b *Builder) estimatedSize() int {
 	return size
 }
 
-type FlushStats struct {
-	MinTimestamp time.Time
-	MaxTimestamp time.Time
+// TimeRange returns the current time range of the builder.
+func (b *Builder) TimeRange() (time.Time, time.Time) {
+	return b.streams.TimeRange()
 }
 
 // Flush flushes all buffered data to the buffer provided. Calling Flush can result
 // in a no-op if there is no buffered data to flush.
 //
-// [Builder.Reset] is called after a successful Flush to discard any pending data and allow new data to be appended.
-func (b *Builder) Flush(output *bytes.Buffer) (FlushStats, error) {
+// [Builder.Reset] is called after a successful Flush to discard any pending
+// data and allow new data to be appended.
+func (b *Builder) Flush() (*dataobj.Object, io.Closer, error) {
 	if b.state == builderStateEmpty {
-		return FlushStats{}, ErrBuilderEmpty
+		return nil, nil, ErrBuilderEmpty
 	}
 
 	timer := prometheus.NewTimer(b.metrics.buildTime)
 	defer timer.ObserveDuration()
 
-	// Appending sections resets them, so we need to load the time range before
-	// appending.
-	minTime, maxTime := b.streams.TimeRange()
-
 	// Flush sections one more time in case they have data.
 	var flushErrors []error
 
@@ -322,34 +319,21 @@ func (b *Builder) Flush(output *bytes.Buffer) (FlushStats, error) {
 
 	if err := errors.Join(flushErrors...); err != nil {
 		b.metrics.flushFailures.Inc()
-		return FlushStats{}, fmt.Errorf("building object: %w", err)
+		return nil, nil, fmt.Errorf("building object: %w", err)
 	}
 
-	sz, err := b.builder.Flush(output)
+	obj, closer, err := b.builder.Flush()
 	if err != nil {
 		b.metrics.flushFailures.Inc()
-		return FlushStats{}, fmt.Errorf("building object: %w", err)
+		return nil, nil, fmt.Errorf("building object: %w", err)
 	}
 
-	b.metrics.builtSize.Observe(float64(sz))
-
-	var (
-		// We don't know if output was empty before calling Flush, so we only start
-		// reading from where we know writing began.
-
-		objReader = bytes.NewReader(output.Bytes()[output.Len()-int(sz):])
-		objLength = sz
-	)
-	obj, err := dataobj.FromReaderAt(objReader, objLength)
-	if err != nil {
-		b.metrics.flushFailures.Inc()
-		return FlushStats{}, fmt.Errorf("failed to create readable object: %w", err)
-	}
+	b.metrics.builtSize.Observe(float64(obj.Size()))
 
 	err = b.observeObject(context.Background(), obj)
 
 	b.Reset()
-	return FlushStats{MinTimestamp: minTime, MaxTimestamp: maxTime}, err
+	return obj, closer, err
 }
 
 func (b *Builder) observeObject(ctx context.Context, obj *dataobj.Object) error {

pkg/dataobj/consumer/logsobj/builder_test.go

@@ -1,7 +1,6 @@
 package logsobj
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"strings"
@@ -12,7 +11,6 @@ import (
 
 	"github.com/grafana/loki/pkg/push"
 
-	"github.com/grafana/loki/v3/pkg/dataobj"
 	"github.com/grafana/loki/v3/pkg/dataobj/sections/logs"
 	"github.com/grafana/loki/v3/pkg/dataobj/sections/streams"
 	"github.com/grafana/loki/v3/pkg/logproto"
@@ -29,9 +27,6 @@ var testBuilderConfig = BuilderConfig{
 }
 
 func TestBuilder(t *testing.T) {
-	buf := bytes.NewBuffer(nil)
-	dirtyBuf := bytes.NewBuffer([]byte("dirty"))
-
 	testStreams := []logproto.Stream{
 		{
 			Labels: `{cluster="test",app="foo"}`,
@@ -83,34 +78,10 @@ func TestBuilder(t *testing.T) {
 		for _, entry := range testStreams {
 			require.NoError(t, builder.Append(entry))
 		}
-		_, err = builder.Flush(buf)
-		require.NoError(t, err)
-	})
-
-	t.Run("Read", func(t *testing.T) {
-		obj, err := dataobj.FromReaderAt(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
-		require.NoError(t, err)
-		require.Equal(t, 1, obj.Sections().Count(streams.CheckSection))
-		require.Equal(t, 1, obj.Sections().Count(logs.CheckSection))
-	})
-
-	t.Run("BuildWithDirtyBuffer", func(t *testing.T) {
-		builder, err := NewBuilder(testBuilderConfig)
+		obj, closer, err := builder.Flush()
 		require.NoError(t, err)
+		defer closer.Close()
 
-		for _, entry := range testStreams {
-			require.NoError(t, builder.Append(entry))
-		}
-
-		_, err = builder.Flush(dirtyBuf)
-		require.NoError(t, err)
-
-		require.Equal(t, buf.Len(), dirtyBuf.Len()-5)
-	})
-
-	t.Run("ReadFromDirtyBuffer", func(t *testing.T) {
-		obj, err := dataobj.FromReaderAt(bytes.NewReader(dirtyBuf.Bytes()[5:]), int64(dirtyBuf.Len()-5))
-		require.NoError(t, err)
 		require.Equal(t, 1, obj.Sections().Count(streams.CheckSection))
 		require.Equal(t, 1, obj.Sections().Count(logs.CheckSection))
 	})

pkg/dataobj/consumer/partition_processor.go

@@ -194,20 +194,23 @@ func (p *partitionProcessor) initBuilder() error {
 	return initErr
 }
 
-func (p *partitionProcessor) flushStream(flushBuffer *bytes.Buffer) error {
-	stats, err := p.builder.Flush(flushBuffer)
+func (p *partitionProcessor) flushStream() error {
+	minTime, maxTime := p.builder.TimeRange()
+
+	obj, closer, err := p.builder.Flush()
 	if err != nil {
 		level.Error(p.logger).Log("msg", "failed to flush builder", "err", err)
 		return err
 	}
+	defer closer.Close()
 
-	objectPath, err := p.uploader.Upload(p.ctx, flushBuffer)
+	objectPath, err := p.uploader.Upload(p.ctx, obj)
 	if err != nil {
 		level.Error(p.logger).Log("msg", "failed to upload object", "err", err)
 		return err
 	}
 
-	if err := p.metastoreUpdater.Update(p.ctx, objectPath, stats.MinTimestamp, stats.MaxTimestamp); err != nil {
+	if err := p.metastoreUpdater.Update(p.ctx, objectPath, minTime, maxTime); err != nil {
 		level.Error(p.logger).Log("msg", "failed to update metastore", "err", err)
 		return err
 	}
@@ -284,17 +287,10 @@ func (p *partitionProcessor) processRecord(record *kgo.Record) {
 			return
 		}
 
-		func() {
-			flushBuffer := p.bufPool.Get().(*bytes.Buffer)
-			defer p.bufPool.Put(flushBuffer)
-
-			flushBuffer.Reset()
-
-			if err := p.flushStream(flushBuffer); err != nil {
-				level.Error(p.logger).Log("msg", "failed to flush stream", "err", err)
-				return
-			}
-		}()
+		if err := p.flushStream(); err != nil {
+			level.Error(p.logger).Log("msg", "failed to flush stream", "err", err)
+			return
+		}
 
 		if err := p.commitRecords(record); err != nil {
 			level.Error(p.logger).Log("msg", "failed to commit records", "err", err)
@@ -346,17 +342,10 @@ func (p *partitionProcessor) idleFlush() {
 		return // Avoid checking too frequently
 	}
 
-	func() {
-		flushBuffer := p.bufPool.Get().(*bytes.Buffer)
-		defer p.bufPool.Put(flushBuffer)
-
-		flushBuffer.Reset()
-
-		if err := p.flushStream(flushBuffer); err != nil {
-			level.Error(p.logger).Log("msg", "failed to flush stream", "err", err)
-			return
-		}
+	if err := p.flushStream(); err != nil {
+		level.Error(p.logger).Log("msg", "failed to flush stream", "err", err)
+		return
+	}
 
-		p.lastFlush = time.Now()
-	}()
+	p.lastFlush = time.Now()
 }

pkg/dataobj/dataobj.go

@@ -85,7 +85,9 @@ import (
 
 // An Object is a representation of a data object.
 type Object struct {
-	dec *decoder
+	rr   rangeReader
+	dec  *decoder
+	size int64
 
 	metadata *filemd.Metadata
 	sections []*Section
@@ -95,8 +97,14 @@ type Object struct {
 // FromBucket returns an error if the metadata of the Object cannot be read or
 // if the provided ctx times out.
 func FromBucket(ctx context.Context, bucket objstore.BucketReader, path string) (*Object, error) {
-	dec := &decoder{rr: &bucketRangeReader{bucket: bucket, path: path}}
-	obj := &Object{dec: dec}
+	rr := &bucketRangeReader{bucket: bucket, path: path}
+	size, err := rr.Size(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("getting size: %w", err)
+	}
+
+	dec := &decoder{rr: rr}
+	obj := &Object{rr: rr, dec: dec, size: size}
 	if err := obj.init(ctx); err != nil {
 		return nil, err
 	}
@@ -107,8 +115,9 @@ func FromBucket(ctx context.Context, bucket objstore.BucketReader, path string)
 // specifies the size of the data object in bytes. FromReaderAt returns an
 // error if the metadata of the Object cannot be read.
 func FromReaderAt(r io.ReaderAt, size int64) (*Object, error) {
-	dec := &decoder{rr: &readerAtRangeReader{size: size, r: r}}
-	obj := &Object{dec: dec}
+	rr := &readerAtRangeReader{size: size, r: r}
+	dec := &decoder{rr: rr}
+	obj := &Object{rr: rr, dec: dec, size: size}
 	if err := obj.init(context.Background()); err != nil {
 		return nil, err
 	}
@@ -139,6 +148,14 @@ func (o *Object) init(ctx context.Context) error {
 	return nil
 }
 
+// Size returns the size of the data object in bytes.
+func (o *Object) Size() int64 { return o.size }
+
 // Sections returns the list of sections available in the Object. The slice of
 // returned sections must not be mutated.
 func (o *Object) Sections() Sections { return o.sections }
+
+// Reader returns a reader for the entire raw data object.
+func (o *Object) Reader(ctx context.Context) (io.ReadCloser, error) {
+	return o.rr.Read(ctx)
+}

pkg/dataobj/encoder.go

@@ -196,9 +196,25 @@ func (enc *encoder) Bytes() int {
 	return enc.data.Len()
 }
 
+// FlushSize returns the exact number of bytes that would be written upon
+// calling [encoder.Flush].
+func (enc *encoder) FlushSize() (int64, error) {
+	return enc.flush(streamio.Discard)
+}
+
 // Flush flushes any buffered data to the underlying writer. After flushing,
 // enc is reset.
 func (enc *encoder) Flush(w streamio.Writer) (int64, error) {
+	size, err := enc.flush(w)
+	if err != nil {
+		return size, err
+	}
+
+	enc.Reset()
+	return size, nil
+}
+
+func (enc *encoder) flush(w streamio.Writer) (int64, error) {
 	cw := countingWriter{w: w}
 
 	if enc.data == nil {
@@ -241,7 +257,6 @@ func (enc *encoder) Flush(w streamio.Writer) (int64, error) {
 		return cw.count, fmt.Errorf("writing magic tailer: %w", err)
 	}
 
-	enc.Reset()
 	return cw.count, nil
 }