Update blockbuilder to periodically flush wals and sort traces

modules/blockbuilder/blockbuilder.go

@@ -28,9 +28,29 @@ const (
 	blockBuilderServiceName = "block-builder"
 	ConsumerGroup           = "block-builder"
 	pollTimeout             = 2 * time.Second
+	cutTime                 = 10 * time.Second
 )
 
 var (
+	metricFetchDuration = promauto.NewHistogramVec(prometheus.HistogramOpts{
+		Namespace:                   "tempo",
+		Subsystem:                   "block_builder",
+		Name:                        "fetch_duration_seconds",
+		Help:                        "Time spent fetching from Kafka.",
+		NativeHistogramBucketFactor: 1.1,
+	}, []string{"partition"})
+	metricFetchBytesTotal = promauto.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: "tempo",
+		Subsystem: "block_builder",
+		Name:      "fetch_bytes_total",
+		Help:      "Total number of bytes fetched from Kafka",
+	}, []string{"partition"})
+	metricFetchRecordsTotal = promauto.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: "tempo",
+		Subsystem: "block_builder",
+		Name:      "fetch_records_total",
+		Help:      "Total number of records fetched from Kafka",
+	}, []string{"partition"})
 	metricPartitionLag = promauto.NewGaugeVec(prometheus.GaugeOpts{
 		Namespace: "tempo",
 		Subsystem: "block_builder",
@@ -180,6 +200,12 @@ func (b *BlockBuilder) consume(ctx context.Context) error {
 	level.Info(b.logger).Log("msg", "starting consume cycle", "cycle_end", end, "active_partitions", partitions)
 	defer func(t time.Time) { metricConsumeCycleDuration.Observe(time.Since(t).Seconds()) }(time.Now())
 
+	// Clear all previous remnants
+	err := b.wal.Clear()
+	if err != nil {
+		return err
+	}
+
 	for _, partition := range partitions {
 		// Consume partition while data remains.
 		// TODO - round-robin one consumption per partition instead to equalize catch-up time.
@@ -207,10 +233,12 @@ func (b *BlockBuilder) consumePartition(ctx context.Context, partition int32, ov
 		dur         = b.cfg.ConsumeCycleDuration
 		topic       = b.cfg.IngestStorageConfig.Kafka.Topic
 		group       = b.cfg.IngestStorageConfig.Kafka.ConsumerGroup
+		partLabel   = strconv.Itoa(int(partition))
 		startOffset kgo.Offset
 		init        bool
 		writer      *writer
 		lastRec     *kgo.Record
+		nextCut     time.Time
 		end         time.Time
 	)
 
@@ -247,6 +275,7 @@ func (b *BlockBuilder) consumePartition(ctx context.Context, partition int32, ov
 outer:
 	for {
 		fetches := func() kgo.Fetches {
+			defer func(t time.Time) { metricFetchDuration.WithLabelValues(partLabel).Observe(time.Since(t).Seconds()) }(time.Now())
 			ctx2, cancel := context.WithTimeout(ctx, pollTimeout)
 			defer cancel()
 			return b.kafkaClient.PollFetches(ctx2)
@@ -257,7 +286,7 @@ outer:
 				// No more data
 				break
 			}
-			metricFetchErrors.WithLabelValues(strconv.Itoa(int(partition))).Inc()
+			metricFetchErrors.WithLabelValues(partLabel).Inc()
 			return false, err
 		}
 
@@ -267,19 +296,23 @@ outer:
 
 		for iter := fetches.RecordIter(); !iter.Done(); {
 			rec := iter.Next()
+			metricFetchBytesTotal.WithLabelValues(partLabel).Add(float64(len(rec.Value)))
+			metricFetchRecordsTotal.WithLabelValues(partLabel).Inc()
 
 			level.Debug(b.logger).Log(
 				"msg", "processing record",
 				"partition", rec.Partition,
 				"offset", rec.Offset,
 				"timestamp", rec.Timestamp,
+				"len", len(rec.Value),
 			)
 
 			// Initialize on first record
 			if !init {
 				end = rec.Timestamp.Add(dur) // When block will be cut
-				metricPartitionLagSeconds.WithLabelValues(strconv.Itoa(int(partition))).Set(time.Since(rec.Timestamp).Seconds())
+				metricPartitionLagSeconds.WithLabelValues(partLabel).Set(time.Since(rec.Timestamp).Seconds())
 				writer = newPartitionSectionWriter(b.logger, uint64(partition), uint64(rec.Offset), b.cfg.BlockConfig, b.overrides, b.wal, b.enc)
+				nextCut = rec.Timestamp.Add(cutTime)
 				init = true
 			}
 
@@ -295,7 +328,16 @@ outer:
 				break outer
 			}
 
-			err := b.pushTraces(rec.Key, rec.Value, writer)
+			if rec.Timestamp.After(nextCut) {
+				// Cut before appending this trace
+				err = writer.cutidle(rec.Timestamp.Add(-cutTime), false)
+				if err != nil {
+					return false, err
+				}
+				nextCut = rec.Timestamp.Add(cutTime)
+			}
+
+			err := b.pushTraces(rec.Timestamp, rec.Key, rec.Value, writer)
 			if err != nil {
 				return false, err
 			}
@@ -313,6 +355,12 @@ outer:
 		return false, nil
 	}
 
+	// Cut any remaining
+	err = writer.cutidle(time.Time{}, true)
+	if err != nil {
+		return false, err
+	}
+
 	err = writer.flush(ctx, b.writer)
 	if err != nil {
 		return false, err
@@ -370,14 +418,14 @@ func (b *BlockBuilder) stopping(err error) error {
 	return err
 }
 
-func (b *BlockBuilder) pushTraces(tenantBytes, reqBytes []byte, p partitionSectionWriter) error {
+func (b *BlockBuilder) pushTraces(ts time.Time, tenantBytes, reqBytes []byte, p partitionSectionWriter) error {
 	req, err := b.decoder.Decode(reqBytes)
 	if err != nil {
 		return fmt.Errorf("failed to decode trace: %w", err)
 	}
 	defer b.decoder.Reset()
 
-	return p.pushBytes(string(tenantBytes), req)
+	return p.pushBytes(ts, string(tenantBytes), req)
 }
 
 func (b *BlockBuilder) getAssignedActivePartitions() []int32 {

modules/blockbuilder/blockbuilder_test.go

@@ -302,7 +302,7 @@ func TestBlockbuilder_committingFails(t *testing.T) {
 	requireLastCommitEquals(t, ctx, client, producedRecords[len(producedRecords)-1].Offset+1)
 }
 
-func blockbuilderConfig(t *testing.T, address string) Config {
+func blockbuilderConfig(t testing.TB, address string) Config {
 	cfg := Config{}
 	flagext.DefaultValues(&cfg)
 
@@ -327,7 +327,7 @@ type ownEverythingSharder struct{}
 
 func (o *ownEverythingSharder) Owns(string) bool { return true }
 
-func newStore(ctx context.Context, t *testing.T) storage.Store {
+func newStore(ctx context.Context, t testing.TB) storage.Store {
 	tmpDir := t.TempDir()
 	s, err := storage.NewStore(storage.Config{
 		Trace: tempodb.Config{
@@ -404,7 +404,7 @@ type mockOverrides struct {
 
 func (m *mockOverrides) DedicatedColumns(_ string) backend.DedicatedColumns { return m.dc }
 
-func newKafkaClient(t *testing.T, config ingest.KafkaConfig) *kgo.Client {
+func newKafkaClient(t testing.TB, config ingest.KafkaConfig) *kgo.Client {
 	writeClient, err := kgo.NewClient(
 		kgo.SeedBrokers(config.Address),
 		kgo.AllowAutoTopicCreation(),
@@ -427,7 +427,7 @@ func countFlushedTraces(store storage.Store) int {
 }
 
 // nolint: revive
-func sendReq(t *testing.T, ctx context.Context, client *kgo.Client) []*kgo.Record {
+func sendReq(t testing.TB, ctx context.Context, client *kgo.Client) []*kgo.Record {
 	traceID := generateTraceID(t)
 
 	req := test.MakePushBytesRequest(t, 10, traceID)
@@ -463,7 +463,7 @@ func sendTracesFor(t *testing.T, ctx context.Context, client *kgo.Client, dur, i
 	}
 }
 
-func generateTraceID(t *testing.T) []byte {
+func generateTraceID(t testing.TB) []byte {
 	traceID := make([]byte, 16)
 	_, err := rand.Read(traceID)
 	require.NoError(t, err)
@@ -478,3 +478,44 @@ func requireLastCommitEquals(t testing.TB, ctx context.Context, client *kgo.Clie
 	require.True(t, ok)
 	require.Equal(t, expectedOffset, offset.At)
 }
+
+func BenchmarkBlockBuilder(b *testing.B) {
+	var (
+		ctx        = context.Background()
+		_, address = testkafka.CreateCluster(b, 1, testTopic)
+		store      = newStore(ctx, b)
+		cfg        = blockbuilderConfig(b, address)
+		client     = newKafkaClient(b, cfg.IngestStorageConfig.Kafka)
+	)
+
+	cfg.ConsumeCycleDuration = 1 * time.Hour
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+
+		var records []*kgo.Record
+
+		for i := 0; i < 1000; i++ {
+			records = append(records, sendReq(b, ctx, client)...)
+		}
+
+		var size int
+		for _, r := range records {
+			size += len(r.Value)
+		}
+
+		b.ResetTimer()
+
+		bb := New(cfg, test.NewTestingLogger(b), newPartitionRingReader(), &mockOverrides{}, store)
+
+		// Startup (without starting the background consume cycle)
+		err := bb.starting(ctx)
+		require.NoError(b, err)
+
+		err = bb.consume(ctx)
+		require.NoError(b, err)
+
+		b.SetBytes(int64(size))
+	}
+}

modules/blockbuilder/partition_writer.go

@@ -18,7 +18,7 @@ import (
 )
 
 type partitionSectionWriter interface {
-	pushBytes(tenant string, req *tempopb.PushBytesRequest) error
+	pushBytes(ts time.Time, tenant string, req *tempopb.PushBytesRequest) error
 	flush(ctx context.Context, store tempodb.Writer) error
 }
 
@@ -50,7 +50,7 @@ func newPartitionSectionWriter(logger log.Logger, partition, cycleEndTs uint64,
 	}
 }
 
-func (p *writer) pushBytes(tenant string, req *tempopb.PushBytesRequest) error {
+func (p *writer) pushBytes(ts time.Time, tenant string, req *tempopb.PushBytesRequest) error {
 	level.Debug(p.logger).Log(
 		"msg", "pushing bytes",
 		"tenant", tenant,
@@ -69,28 +69,20 @@ func (p *writer) pushBytes(tenant string, req *tempopb.PushBytesRequest) error {
 			return fmt.Errorf("failed to unmarshal trace: %w", err)
 		}
 
-		var start, end uint64
-		for _, b := range tr.ResourceSpans {
-			for _, ss := range b.ScopeSpans {
-				for _, s := range ss.Spans {
-					if start == 0 || s.StartTimeUnixNano < start {
-						start = s.StartTimeUnixNano
-					}
-					if s.EndTimeUnixNano > end {
-						end = s.EndTimeUnixNano
-					}
-				}
-			}
+		if err := i.AppendTrace(req.Ids[j], tr, ts); err != nil {
+			return err
 		}
+	}
 
-		startSeconds := uint32(start / uint64(time.Second))
-		endSeconds := uint32(end / uint64(time.Second))
+	return nil
+}
 
-		if err := i.AppendTrace(req.Ids[j], tr, startSeconds, endSeconds); err != nil {
+func (p *writer) cutidle(since time.Time, immediate bool) error {
+	for _, i := range p.m {
+		if err := i.CutIdle(since, immediate); err != nil {
 			return err
 		}
 	}
-
 	return nil
 }