storage/ingest: refactor recordConsumer.Consume to take iter.Seq (#12482)

#### What this PR does

The PR updates internals of `pkg/storage/ingest`. Now, the
`recordConsumer.Consume` accepts an iterator over Kafka records,
`iter.Seq[*kgo.Record]`.

Following the discussion in
#12443 (comment), we
want the block-builder to reuse the `pusherConsumer`, but we want to
avoid exposing the existing `record` data structs. Instead of expanding
the existing interface, I propose we refactor the `recordConsumer` and
get rid of the intermediate `record` data type completely.

---

I checked the benchmarks of the `PusherConsumer`, and, although this
wasn't planned, I got some positive results from the refactoring (_I
didn't investigate if the results represent real use case, or this is
just a side effect of the benchmarks themselves_):

```
goos: darwin
goarch: arm64
pkg: github.com/grafana/mimir/pkg/storage/ingest
cpu: Apple M2
                                   │   1.bench   │              2.bench               │
                                   │   sec/op    │   sec/op     vs base               │
PusherConsumer/sequential_pusher-8   332.6µ ± 0%   333.7µ ± 0%       ~ (p=0.089 n=10)
PusherConsumer/parallel_pusher-8     332.1µ ± 0%   334.0µ ± 0%  +0.59% (p=0.003 n=10)
geomean                              332.4µ        333.9µ       +0.46%

                                   │   1.bench    │               2.bench               │
                                   │     B/op     │     B/op      vs base               │
PusherConsumer/sequential_pusher-8   41.05Ki ± 0%   40.41Ki ± 0%  -1.57% (p=0.000 n=10)
PusherConsumer/parallel_pusher-8     41.64Ki ± 0%   40.97Ki ± 0%  -1.61% (p=0.000 n=10)
geomean                              41.35Ki        40.69Ki       -1.59%

                                   │  1.bench   │              2.bench              │
                                   │ allocs/op  │ allocs/op   vs base               │
PusherConsumer/sequential_pusher-8   564.0 ± 0%   519.0 ± 0%  -7.98% (p=0.000 n=10)
PusherConsumer/parallel_pusher-8     570.0 ± 0%   525.0 ± 0%  -7.89% (p=0.000 n=10)
geomean                              567.0        522.0       -7.94%
```

#### Checklist

- [x] Tests updated.
- [ ] Documentation added.
- [ ] `CHANGELOG.md` updated - the order of entries should be
`[CHANGE]`, `[FEATURE]`, `[ENHANCEMENT]`, `[BUGFIX]`. If changelog entry
is not needed, please add the `changelog-not-needed` label to the PR.
- [ ]
[`about-versioning.md`](https://github.com/grafana/mimir/blob/main/docs/sources/mimir/configure/about-versioning.md)
updated with experimental features.

Signed-off-by: Vladimir Varankin <[email protected]>

Author: narqo

pkg/storage/ingest/pusher.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"iter"
 	"math/rand/v2"
 	"sync"
 	"time"
@@ -17,6 +18,7 @@ import (
 	"github.com/grafana/dskit/middleware"
 	"github.com/grafana/dskit/multierror"
 	"github.com/grafana/dskit/user"
+	"github.com/twmb/franz-go/pkg/kgo"
 
 	"github.com/grafana/mimir/pkg/mimirpb"
 	util_log "github.com/grafana/mimir/pkg/util/log"
@@ -60,7 +62,7 @@ func newPusherConsumer(pusher Pusher, kafkaCfg KafkaConfig, metrics *pusherConsu
 
 // Consume implements the recordConsumer interface.
 // It'll use a separate goroutine to unmarshal the next record while we push the current record to storage.
-func (c pusherConsumer) Consume(ctx context.Context, records []record) (returnErr error) {
+func (c pusherConsumer) Consume(ctx context.Context, records iter.Seq[*kgo.Record]) (returnErr error) {
 	defer func(processingStart time.Time) {
 		c.metrics.processingTimeSeconds.Observe(time.Since(processingStart).Seconds())
 	}(time.Now())
@@ -71,7 +73,7 @@ func (c pusherConsumer) Consume(ctx context.Context, records []record) (returnEr
 		ctx      context.Context
 		tenantID string
 		err      error
-		index    int
+		offset   int64
 	}
 
 	recordsChannel := make(chan parsedRecord)
@@ -80,10 +82,10 @@ func (c pusherConsumer) Consume(ctx context.Context, records []record) (returnEr
 	ctx, cancel := context.WithCancelCause(ctx)
 
 	// Now, unmarshal the records into the channel.
-	go func(unmarshalCtx context.Context, records []record, ch chan<- parsedRecord) {
+	go func(unmarshalCtx context.Context, records iter.Seq[*kgo.Record], ch chan<- parsedRecord) {
 		defer close(ch)
 
-		for index, r := range records {
+		for rec := range records {
 			// Before we being unmarshalling the write request check if the context was cancelled.
 			select {
 			case <-unmarshalCtx.Done():
@@ -93,18 +95,18 @@ func (c pusherConsumer) Consume(ctx context.Context, records []record) (returnEr
 			}
 
 			parsed := parsedRecord{
-				ctx:                  r.ctx,
-				tenantID:             r.tenantID,
 				PreallocWriteRequest: &mimirpb.PreallocWriteRequest{},
-				index:                index,
+				// This context carries the tracing data for this individual record;
+				// kotel populates this data when it fetches the messages.
+				ctx:      rec.Context,
+				tenantID: string(rec.Key),
+				offset:   rec.Offset,
 			}
 
-			if r.version > LatestRecordVersion {
-				parsed.err = fmt.Errorf("received a record with an unsupported version: %d, max supported version: %d", r.version, LatestRecordVersion)
-			}
+			recVersion := ParseRecordVersion(rec)
 
 			// We don't free the WriteRequest slices because they are being freed by a level below.
-			err := DeserializeRecordContent(r.content, parsed.PreallocWriteRequest, r.version)
+			err := DeserializeRecordContent(rec.Value, parsed.PreallocWriteRequest, recVersion)
 			if err != nil {
 				parsed.err = fmt.Errorf("parsing ingest consumer write request: %w", err)
 			}
@@ -121,8 +123,8 @@ func (c pusherConsumer) Consume(ctx context.Context, records []record) (returnEr
 	// We accumulate the total bytes across all records per tenant to determine the number of timeseries we expected to receive.
 	// Then, we'll use that to determine the number of shards we need to parallelize the writes.
 	var bytesPerTenant = make(map[string]int)
-	for _, r := range records {
-		bytesPerTenant[r.tenantID] += len(r.content)
+	for rec := range records {
+		bytesPerTenant[string(rec.Key)] += len(rec.Value)
 	}
 
 	// Create and start the storage writer.
@@ -149,7 +151,7 @@ func (c pusherConsumer) Consume(ctx context.Context, records []record) (returnEr
 		err := c.pushToStorage(r.ctx, r.tenantID, &r.WriteRequest, writer)
 		if err != nil {
 			cancel(cancellation.NewErrorf("error while pushing to storage")) // Stop the unmarshalling goroutine.
-			return fmt.Errorf("consuming record at index %d for tenant %s: %w", r.index, r.tenantID, err)
+			return fmt.Errorf("consuming record at offset %d for tenant %s: %w", r.offset, r.tenantID, err)
 		}
 	}