feat: write to dataobj partitions based on segmentation key

grobinson-grafana · grobinson-grafana · commit 83ed9276e686 · 2025-11-21T13:02:53.000Z
diff --git a/pkg/distributor/dataobj_tee.go b/pkg/distributor/dataobj_tee.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/grafana/dskit/ring"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/twmb/franz-go/pkg/kgo"
@@ -16,31 +15,43 @@ import (
 )
 
 type DataObjTeeConfig struct {
-	Enabled          bool   `yaml:"enabled"`
-	Topic            string `yaml:"topic"`
-	MaxBufferedBytes int    `yaml:"max_buffered_bytes"`
+	Enabled               bool   `yaml:"enabled"`
+	Topic                 string `yaml:"topic"`
+	MaxBufferedBytes      int    `yaml:"max_buffered_bytes"`
+	PerPartitionRateBytes int    `yaml:"per_partition_rate_bytes"`
 }
 
 func (c *DataObjTeeConfig) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&c.Enabled, "distributor.dataobj-tee.enabled", false, "Enable data object tee.")
 	f.StringVar(&c.Topic, "distributor.dataobj-tee.topic", "", "Topic for data object tee.")
 	f.IntVar(&c.MaxBufferedBytes, "distributor.dataobj-tee.max-buffered-bytes", 100<<20, "Maximum number of bytes to buffer.")
+	f.IntVar(&c.PerPartitionRateBytes, "distributor.dataobj-tee.per-partition-rate-bytes", 1024*1024, "The per-tenant partition rate (bytes/sec).")
 }
 
 func (c *DataObjTeeConfig) Validate() error {
-	if c.Enabled && c.Topic == "" {
+	if !c.Enabled {
+		return nil
+	}
+	if c.Topic == "" {
 		return errors.New("the topic is required")
 	}
+	if c.MaxBufferedBytes < 0 {
+		return errors.New("max buffered bytes cannot be negative")
+	}
+	if c.PerPartitionRateBytes < 0 {
+		return errors.New("per partition rate bytes cannot be negative")
+	}
 	return nil
 }
 
 // DataObjTee is a tee that duplicates streams to the data object topic.
 // It is a temporary solution while we work on segmentation keys.
 type DataObjTee struct {
-	cfg        *DataObjTeeConfig
-	client     *kgo.Client
-	ringReader ring.PartitionRingReader
-	logger     log.Logger
+	cfg          *DataObjTeeConfig
+	limitsClient *ingestLimits
+	kafkaClient  *kgo.Client
+	resolver     *SegmentationPartitionResolver
+	logger       log.Logger
 
 	// Metrics.
 	failures prometheus.Counter
@@ -50,16 +61,18 @@ type DataObjTee struct {
 // NewDataObjTee returns a new DataObjTee.
 func NewDataObjTee(
 	cfg *DataObjTeeConfig,
-	client *kgo.Client,
-	ringReader ring.PartitionRingReader,
+	resolver *SegmentationPartitionResolver,
+	limitsClient *ingestLimits,
+	kafkaClient *kgo.Client,
 	logger log.Logger,
 	reg prometheus.Registerer,
 ) (*DataObjTee, error) {
 	return &DataObjTee{
-		cfg:        cfg,
-		client:     client,
-		ringReader: ringReader,
-		logger:     logger,
+		cfg:          cfg,
+		resolver:     resolver,
+		kafkaClient:  kafkaClient,
+		limitsClient: limitsClient,
+		logger:       logger,
 		failures: promauto.With(reg).NewCounter(prometheus.CounterOpts{
 			Name: "loki_distributor_dataobj_tee_duplicate_stream_failures_total",
 			Help: "Total number of streams that could not be duplicated.",
@@ -71,18 +84,47 @@ func NewDataObjTee(
 	}, nil
 }
 
+// A SegmentedStream is a KeyedStream with a segmentation key.
+type SegmentedStream struct {
+	KeyedStream
+	SegmentationKey SegmentationKey
+}
+
 // Duplicate implements the [Tee] interface.
-func (t *DataObjTee) Duplicate(_ context.Context, tenant string, streams []KeyedStream) {
-	for _, s := range streams {
-		go t.duplicate(tenant, s)
+func (t *DataObjTee) Duplicate(ctx context.Context, tenant string, streams []KeyedStream) {
+	segmentationKeyStreams := make([]SegmentedStream, 0, len(streams))
+	for _, stream := range streams {
+		segmentationKey, err := GetSegmentationKey(stream)
+		if err != nil {
+			level.Error(t.logger).Log("msg", "failed to get segmentation key", "err", err)
+			t.failures.Inc()
+			return
+		}
+		segmentationKeyStreams = append(segmentationKeyStreams, SegmentedStream{
+			KeyedStream:     stream,
+			SegmentationKey: segmentationKey,
+		})
+	}
+	rates, err := t.limitsClient.UpdateRates(ctx, tenant, segmentationKeyStreams)
+	if err != nil {
+		level.Error(t.logger).Log("msg", "failed to update rates", "err", err)
+	}
+	// fastRates is a temporary lookup table that lets us find the rate
+	// for a segmentation key in constant time.
+	fastRates := make(map[uint64]uint64, len(rates))
+	for _, rate := range rates {
+		fastRates[rate.StreamHash] = rate.Rate
+	}
+	for _, s := range segmentationKeyStreams {
+		go t.duplicate(ctx, tenant, s, fastRates[s.SegmentationKey.Sum64()])
 	}
 }
 
-func (t *DataObjTee) duplicate(tenant string, stream KeyedStream) {
+func (t *DataObjTee) duplicate(ctx context.Context, tenant string, stream SegmentedStream, rateBytes uint64) {
 	t.total.Inc()
-	partition, err := t.ringReader.PartitionRing().ActivePartitionForKey(stream.HashKey)
+	partition, err := t.resolver.Resolve(ctx, stream.SegmentationKey, rateBytes)
 	if err != nil {
-		level.Error(t.logger).Log("msg", "failed to get partition", "err", err)
+		level.Error(t.logger).Log("msg", "failed to resolve partition", "err", err)
 		t.failures.Inc()
 		return
 	}
@@ -92,7 +134,7 @@ func (t *DataObjTee) duplicate(tenant string, stream KeyedStream) {
 		t.failures.Inc()
 		return
 	}
-	results := t.client.ProduceSync(context.TODO(), records...)
+	results := t.kafkaClient.ProduceSync(ctx, records...)
 	if err := results.FirstErr(); err != nil {
 		level.Error(t.logger).Log("msg", "failed to produce records", "err", err)
 		t.failures.Inc()
diff --git a/pkg/distributor/distributor.go b/pkg/distributor/distributor.go
@@ -283,6 +283,8 @@ func New(
 		return nil, fmt.Errorf("partition ring is required for kafka writes")
 	}
 
+	ingestLimits := newIngestLimits(limitsFrontendClient, registerer)
+
 	var kafkaWriter KafkaProducer
 	if cfg.KafkaEnabled {
 		kafkaClient, err := kafka_client.NewWriterClient("distributor", cfg.KafkaConfig, 20, logger, registerer)
@@ -295,10 +297,17 @@ func New(
 		)
 
 		if cfg.DataObjTeeConfig.Enabled {
+			resolver := NewSegmentationPartitionResolver(
+				uint64(cfg.DataObjTeeConfig.PerPartitionRateBytes),
+				dataObjConsumerPartitionRing,
+				registerer,
+				logger,
+			)
 			dataObjTee, err := NewDataObjTee(
 				&cfg.DataObjTeeConfig,
+				resolver,
+				ingestLimits,
 				kafkaClient,
-				dataObjConsumerPartitionRing,
 				logger,
 				registerer,
 			)
@@ -379,7 +388,7 @@ func New(
 		writeFailuresManager:  writefailures.NewManager(logger, registerer, cfg.WriteFailuresLogging, configs, "distributor"),
 		kafkaWriter:           kafkaWriter,
 		partitionRing:         partitionRing,
-		ingestLimits:          newIngestLimits(limitsFrontendClient, registerer),
+		ingestLimits:          ingestLimits,
 		numMetadataPartitions: numMetadataPartitions,
 	}
 
diff --git a/pkg/distributor/ingest_limits.go b/pkg/distributor/ingest_limits.go
@@ -194,7 +194,7 @@ func newExceedsLimitsRequest(tenant string, streams []KeyedStream) (*proto.Excee
 // UpdateRates updates the rates for the streams and returns a slice of the
 // updated rates for all streams. Any streams that could not have rates updated
 // have a rate of zero.
-func (l *ingestLimits) UpdateRates(ctx context.Context, tenant string, streams []KeyedStream) ([]*proto.UpdateRatesResult, error) {
+func (l *ingestLimits) UpdateRates(ctx context.Context, tenant string, streams []SegmentedStream) ([]*proto.UpdateRatesResult, error) {
 	l.requests.WithLabelValues("UpdateRates").Inc()
 	req, err := newUpdateRatesRequest(tenant, streams)
 	if err != nil {
@@ -209,7 +209,7 @@ func (l *ingestLimits) UpdateRates(ctx context.Context, tenant string, streams [
 	return resp.Results, nil
 }
 
-func newUpdateRatesRequest(tenant string, streams []KeyedStream) (*proto.UpdateRatesRequest, error) {
+func newUpdateRatesRequest(tenant string, streams []SegmentedStream) (*proto.UpdateRatesRequest, error) {
 	// The distributor sends the hashes of all streams in the request to the
 	// limits-frontend. The limits-frontend is responsible for deciding if
 	// the request would exceed the tenants limits, and if so, which streams
@@ -218,7 +218,7 @@ func newUpdateRatesRequest(tenant string, streams []KeyedStream) (*proto.UpdateR
 	for _, stream := range streams {
 		entriesSize, structuredMetadataSize := calculateStreamSizes(stream.Stream)
 		streamMetadata = append(streamMetadata, &proto.StreamMetadata{
-			StreamHash:      stream.HashKeyNoShard,
+			StreamHash:      stream.SegmentationKey.Sum64(),
 			TotalSize:       entriesSize + structuredMetadataSize,
 			IngestionPolicy: stream.Policy,
 		})
diff --git a/pkg/distributor/ingest_limits_test.go b/pkg/distributor/ingest_limits_test.go
@@ -302,7 +302,7 @@ func TestIngestLimits_UpdateRates(t *testing.T) {
 	tests := []struct {
 		name            string
 		tenant          string
-		streams         []KeyedStream
+		streams         []SegmentedStream
 		expectedRequest *proto.UpdateRatesRequest
 		response        *proto.UpdateRatesResponse
 		responseErr     error
@@ -311,7 +311,7 @@ func TestIngestLimits_UpdateRates(t *testing.T) {
 	}{{
 		name:   "error should be returned if rates cannot be updated",
 		tenant: "test",
-		streams: []KeyedStream{{
+		streams: []SegmentedStream{{
 			HashKeyNoShard: 1,
 		}},
 		expectedRequest: &proto.UpdateRatesRequest{
@@ -325,7 +325,7 @@ func TestIngestLimits_UpdateRates(t *testing.T) {
 	}, {
 		name:   "updates rates",
 		tenant: "test",
-		streams: []KeyedStream{{
+		streams: []SegmentedStream{{
 			HashKeyNoShard: 1,
 		}},
 		expectedRequest: &proto.UpdateRatesRequest{
diff --git a/pkg/limits/frontend/frontend.go b/pkg/limits/frontend/frontend.go
@@ -137,10 +137,7 @@ func (f *Frontend) ExceedsLimits(ctx context.Context, req *proto.ExceedsLimitsRe
 	return &proto.ExceedsLimitsResponse{Results: results}, nil
 }
 
-func (f *Frontend) UpdateRates(
-	ctx context.Context,
-	req *proto.UpdateRatesRequest,
-) (*proto.UpdateRatesResponse, error) {
+func (f *Frontend) UpdateRates(ctx context.Context, req *proto.UpdateRatesRequest) (*proto.UpdateRatesResponse, error) {
 	results := make([]*proto.UpdateRatesResult, 0, len(req.Streams))
 	resps, err := f.limitsClient.UpdateRates(ctx, req)
 	if err != nil {
diff --git a/pkg/limits/frontend/ring.go b/pkg/limits/frontend/ring.go
@@ -73,10 +73,7 @@ func newRingLimitsClient(
 }
 
 // ExceedsLimits implements the [exceedsLimitsGatherer] interface.
-func (r *ringLimitsClient) ExceedsLimits(
-	ctx context.Context,
-	req *proto.ExceedsLimitsRequest,
-) ([]*proto.ExceedsLimitsResponse, error) {
+func (r *ringLimitsClient) ExceedsLimits(ctx context.Context, req *proto.ExceedsLimitsRequest) ([]*proto.ExceedsLimitsResponse, error) {
 	if len(req.Streams) == 0 {
 		return nil, nil
 	}
@@ -101,10 +98,7 @@ func (r *ringLimitsClient) ExceedsLimits(
 }
 
 // UpdateRates implements the [exceedsLimitsGatherer] interface.
-func (r *ringLimitsClient) UpdateRates(
-	ctx context.Context,
-	req *proto.UpdateRatesRequest,
-) ([]*proto.UpdateRatesResponse, error) {
+func (r *ringLimitsClient) UpdateRates(ctx context.Context, req *proto.UpdateRatesRequest) ([]*proto.UpdateRatesResponse, error) {
 	if len(req.Streams) == 0 {
 		return nil, nil
 	}
@@ -238,12 +232,7 @@ type doRPCsFunc func(
 
 // exhaustAllZones queries all zones, one at a time, until either all streams
 // have been answered or all zones have been exhausted.
-func (r *ringLimitsClient) exhaustAllZones(
-	ctx context.Context,
-	tenant string,
-	streams []*proto.StreamMetadata,
-	doRPCs doRPCsFunc,
-) ([]*proto.StreamMetadata, error) {
+func (r *ringLimitsClient) exhaustAllZones(ctx context.Context, tenant string, streams []*proto.StreamMetadata, doRPCs doRPCsFunc) ([]*proto.StreamMetadata, error) {
 	zonesIter, err := r.allZones(ctx)
 	if err != nil {
 		return nil, err
diff --git a/pkg/limits/frontend/ring_test.go b/pkg/limits/frontend/ring_test.go
@@ -15,7 +15,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/limits/proto"
 )
 
-func TestRingGatherer_ExceedsLimits(t *testing.T) {
+func TestRingLimitsClient_ExceedsLimits(t *testing.T) {
 	tests := []struct {
 		name    string
 		request *proto.ExceedsLimitsRequest
@@ -442,7 +442,7 @@ func TestRingGatherer_ExceedsLimits(t *testing.T) {
 	}
 }
 
-func TestRingStreamUsageGatherer_GetZoneAwarePartitionConsumers(t *testing.T) {
+func TestRingLimitsClient_GetZoneAwarePartitionConsumers(t *testing.T) {
 	tests := []struct {
 		name                              string
 		instances                         []ring.InstanceDesc
@@ -603,7 +603,7 @@ func TestRingStreamUsageGatherer_GetZoneAwarePartitionConsumers(t *testing.T) {
 	}
 }
 
-func TestRingStreamUsageGatherer_GetPartitionConsumers(t *testing.T) {
+func TestRingLimitsClient_GetPartitionConsumers(t *testing.T) {
 	tests := []struct {
 		name string
 		// Instances contains the complete set of instances that should be mocked.
@@ -740,7 +740,7 @@ func TestRingStreamUsageGatherer_GetPartitionConsumers(t *testing.T) {
 	}
 }
 
-func TestRingStreamUsageGatherer_GetPartitionConsumers_Caching(t *testing.T) {
+func TestRingLimitsClient_GetPartitionConsumers_Caching(t *testing.T) {
 	// Set up the mock clients.
 	req0 := proto.GetAssignedPartitionsResponse{
 		AssignedPartitions: map[int32]int64{

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ import (`
`15`	`15`	`"github.com/grafana/loki/v3/pkg/limits/proto"`
`16`	`16`	`)`
`17`	`17`
`18`		`-func TestRingGatherer_ExceedsLimits(t *testing.T) {`
	`18`	`+func TestRingLimitsClient_ExceedsLimits(t *testing.T) {`
`19`	`19`	`tests := []struct {`
`20`	`20`	`name string`
`21`	`21`	`request *proto.ExceedsLimitsRequest`
`@@ -442,7 +442,7 @@ func TestRingGatherer_ExceedsLimits(t *testing.T) {`
`442`	`442`	`}`
`443`	`443`	`}`
`444`	`444`
`445`		`-func TestRingStreamUsageGatherer_GetZoneAwarePartitionConsumers(t *testing.T) {`
	`445`	`+func TestRingLimitsClient_GetZoneAwarePartitionConsumers(t *testing.T) {`
`446`	`446`	`tests := []struct {`
`447`	`447`	`name string`
`448`	`448`	`instances []ring.InstanceDesc`
`@@ -603,7 +603,7 @@ func TestRingStreamUsageGatherer_GetZoneAwarePartitionConsumers(t *testing.T) {`
`603`	`603`	`}`
`604`	`604`	`}`
`605`	`605`
`606`		`-func TestRingStreamUsageGatherer_GetPartitionConsumers(t *testing.T) {`
	`606`	`+func TestRingLimitsClient_GetPartitionConsumers(t *testing.T) {`
`607`	`607`	`tests := []struct {`
`608`	`608`	`name string`
`609`	`609`	`// Instances contains the complete set of instances that should be mocked.`
`@@ -740,7 +740,7 @@ func TestRingStreamUsageGatherer_GetPartitionConsumers(t *testing.T) {`
`740`	`740`	`}`
`741`	`741`	`}`
`742`	`742`
`743`		`-func TestRingStreamUsageGatherer_GetPartitionConsumers_Caching(t *testing.T) {`
	`743`	`+func TestRingLimitsClient_GetPartitionConsumers_Caching(t *testing.T) {`
`744`	`744`	`// Set up the mock clients.`
`745`	`745`	`req0 := proto.GetAssignedPartitionsResponse{`
`746`	`746`	`AssignedPartitions: map[int32]int64{`