chore: add PartitionLifecycler to manage assignment and revocation (#17709)

grobinson-grafana · web-flow · commit a580d52f6545 · 2025-05-16T11:01:59.000+01:00
diff --git a/pkg/limits/partition_lifecycler.go b/pkg/limits/partition_lifecycler.go
@@ -0,0 +1,45 @@
+package limits
+
+import (
+	"context"
+
+	"github.com/go-kit/log"
+	"github.com/twmb/franz-go/pkg/kgo"
+)
+
+// PartitionLifecycler manages assignment and revocation of partitions.
+type PartitionLifecycler struct {
+	partitionManager *PartitionManager
+	usage            *UsageStore
+	logger           log.Logger
+}
+
+// NewPartitionLifecycler returns a new PartitionLifecycler.
+func NewPartitionLifecycler(partitionManager *PartitionManager, usage *UsageStore, logger log.Logger) *PartitionLifecycler {
+	return &PartitionLifecycler{
+		partitionManager: partitionManager,
+		usage:            usage,
+		logger:           log.With(logger, "component", "limits.PartitionLifecycler"),
+	}
+}
+
+// Assign implements kgo.OnPartitionsAssigned.
+func (l *PartitionLifecycler) Assign(ctx context.Context, _ *kgo.Client, topics map[string][]int32) {
+	// We expect the client to just consume one topic.
+	// TODO(grobinson): Figure out what to do if this is not the case.
+	for _, partitions := range topics {
+		l.partitionManager.Assign(ctx, partitions)
+		return
+	}
+}
+
+// Revoke implements kgo.OnPartitionsRevoked.
+func (l *PartitionLifecycler) Revoke(ctx context.Context, _ *kgo.Client, topics map[string][]int32) {
+	// We expect the client to just consume one topic.
+	// TODO(grobinson): Figure out what to do if this is not the case.
+	for _, partitions := range topics {
+		l.partitionManager.Revoke(ctx, partitions)
+		l.usage.EvictPartitions(partitions)
+		return
+	}
+}
diff --git a/pkg/limits/partition_manager.go b/pkg/limits/partition_manager.go
@@ -6,13 +6,11 @@ import (
 
 	"github.com/coder/quartz"
 	"github.com/go-kit/log"
-	"github.com/twmb/franz-go/pkg/kgo"
 )
 
-// PartitionManager keeps track of the partitions assigned and for
-// each partition a timestamp of when it was last updated.
+// PartitionManager keeps track of the partitions assigned and the timestamp
+// of when it was assigned.
 type PartitionManager struct {
-	// partitions maps partitionID to last updated (unix nanoseconds).
 	partitions map[int32]int64
 	mtx        sync.Mutex
 	logger     log.Logger
@@ -30,15 +28,12 @@ func NewPartitionManager(logger log.Logger) *PartitionManager {
 	}
 }
 
-// Assign assigns the partitions and sets the last updated timestamp for each
-// partition to the current time.
-func (m *PartitionManager) Assign(_ context.Context, _ *kgo.Client, topicPartitions map[string][]int32) {
+// Assign assigns the partitions.
+func (m *PartitionManager) Assign(_ context.Context, partitions []int32) {
 	m.mtx.Lock()
 	defer m.mtx.Unlock()
-	for _, partitions := range topicPartitions {
-		for _, partition := range partitions {
-			m.partitions[partition] = m.clock.Now().UnixNano()
-		}
+	for _, partition := range partitions {
+		m.partitions[partition] = m.clock.Now().UnixNano()
 	}
 }
 
@@ -50,23 +45,23 @@ func (m *PartitionManager) Has(partition int32) bool {
 	return ok
 }
 
-// List returns a map of all assigned partitions and their last updated timestamps.
+// List returns a map of all assigned partitions and the timestamp of when
+// each partition was assigned.
 func (m *PartitionManager) List() map[int32]int64 {
 	m.mtx.Lock()
 	defer m.mtx.Unlock()
-	v := make(map[int32]int64)
+	result := make(map[int32]int64)
 	for partition, lastUpdated := range m.partitions {
-		v[partition] = lastUpdated
+		result[partition] = lastUpdated
 	}
-	return v
+	return result
 }
 
-func (m *PartitionManager) Remove(_ context.Context, _ *kgo.Client, topicPartitions map[string][]int32) {
+// Revoke revokes the partitions.
+func (m *PartitionManager) Revoke(_ context.Context, partitions []int32) {
 	m.mtx.Lock()
 	defer m.mtx.Unlock()
-	for _, partitions := range topicPartitions {
-		for _, partition := range partitions {
-			delete(m.partitions, partition)
-		}
+	for _, partition := range partitions {
+		delete(m.partitions, partition)
 	}
 }
diff --git a/pkg/limits/partition_manager_test.go b/pkg/limits/partition_manager_test.go
@@ -17,9 +17,7 @@ func TestPartitionManager_Assign(t *testing.T) {
 	// Advance the clock so we compare with a time that is not the default
 	// value.
 	c.Advance(1)
-	m.Assign(context.Background(), nil, map[string][]int32{
-		"foo": {1, 2, 3},
-	})
+	m.Assign(context.Background(), []int32{1, 2, 3})
 	// Assert that the partitions were assigned and the timestamps are set to
 	// the current time.
 	now := c.Now().UnixNano()
@@ -32,9 +30,7 @@ func TestPartitionManager_Assign(t *testing.T) {
 	// partition #4. We expect the updated timestamp is equal to the advanced
 	// time.
 	c.Advance(1)
-	m.Assign(context.Background(), nil, map[string][]int32{
-		"foo": {3, 4},
-	})
+	m.Assign(context.Background(), []int32{3, 4})
 	later := c.Now().UnixNano()
 	require.Equal(t, map[int32]int64{
 		1: now,
@@ -48,9 +44,7 @@ func TestPartitionManager_Has(t *testing.T) {
 	m := NewPartitionManager(log.NewNopLogger())
 	c := quartz.NewMock(t)
 	m.clock = c
-	m.Assign(context.Background(), nil, map[string][]int32{
-		"foo": {1, 2, 3},
-	})
+	m.Assign(context.Background(), []int32{1, 2, 3})
 	require.True(t, m.Has(1))
 	require.True(t, m.Has(2))
 	require.True(t, m.Has(3))
@@ -64,9 +58,7 @@ func TestPartitionManager_List(t *testing.T) {
 	// Advance the clock so we compare with a time that is not the default
 	// value.
 	c.Advance(1)
-	m.Assign(context.Background(), nil, map[string][]int32{
-		"foo": {1, 2, 3},
-	})
+	m.Assign(context.Background(), []int32{1, 2, 3})
 	now := c.Now().UnixNano()
 	result := m.List()
 	require.Equal(t, map[int32]int64{
@@ -81,13 +73,11 @@ func TestPartitionManager_List(t *testing.T) {
 	require.NotEqual(t, p1, p2)
 }
 
-func TestPartitionManager_Remove(t *testing.T) {
+func TestPartitionManager_Revoke(t *testing.T) {
 	m := NewPartitionManager(log.NewNopLogger())
 	c := quartz.NewMock(t)
 	m.clock = c
-	m.Assign(context.Background(), nil, map[string][]int32{
-		"foo": {1, 2, 3},
-	})
+	m.Assign(context.Background(), []int32{1, 2, 3})
 	// Assert that the partitions were assigned and the timestamps are set to
 	// the current time.
 	now := c.Now().UnixNano()
@@ -96,10 +86,8 @@ func TestPartitionManager_Remove(t *testing.T) {
 		2: now,
 		3: now,
 	}, m.partitions)
-	// Remove partitions 2 and 3.
-	m.Remove(context.Background(), nil, map[string][]int32{
-		"foo": {2, 3},
-	})
+	// Revoke partitions 2 and 3.
+	m.Revoke(context.Background(), []int32{2, 3})
 	require.Equal(t, map[int32]int64{
 		1: now,
 	}, m.partitions)
diff --git a/pkg/limits/service.go b/pkg/limits/service.go
@@ -110,7 +110,8 @@ type IngestLimits struct {
 	wal   WAL
 
 	// Track partition assignments
-	partitionManager *PartitionManager
+	partitionManager    *PartitionManager
+	partitionLifecycler *PartitionLifecycler
 
 	// Used for tests.
 	clock quartz.Clock
@@ -138,6 +139,8 @@ func NewIngestLimits(cfg Config, lims Limits, logger log.Logger, reg prometheus.
 		clock:            quartz.NewReal(),
 	}
 
+	s.partitionLifecycler = NewPartitionLifecycler(s.partitionManager, s.usage, logger)
+
 	// Initialize internal metadata metrics
 	if err := reg.Register(s); err != nil {
 		return nil, fmt.Errorf("failed to register ingest limits internal metadata metrics: %w", err)
@@ -165,8 +168,8 @@ func NewIngestLimits(cfg Config, lims Limits, logger log.Logger, reg prometheus.
 		kgo.Balancers(kgo.CooperativeStickyBalancer()),
 		kgo.ConsumeResetOffset(kgo.NewOffset().AfterMilli(s.clock.Now().Add(-s.cfg.WindowSize).UnixMilli())),
 		kgo.DisableAutoCommit(),
-		kgo.OnPartitionsAssigned(s.onPartitionsAssigned),
-		kgo.OnPartitionsRevoked(s.onPartitionsRevoked),
+		kgo.OnPartitionsAssigned(s.partitionLifecycler.Assign),
+		kgo.OnPartitionsRevoked(s.partitionLifecycler.Revoke),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create kafka client: %w", err)
@@ -225,18 +228,6 @@ func (s *IngestLimits) Collect(m chan<- prometheus.Metric) {
 	)
 }
 
-func (s *IngestLimits) onPartitionsAssigned(ctx context.Context, client *kgo.Client, partitions map[string][]int32) {
-	s.partitionManager.Assign(ctx, client, partitions)
-}
-
-func (s *IngestLimits) onPartitionsRevoked(ctx context.Context, client *kgo.Client, partitions map[string][]int32) {
-	s.partitionManager.Remove(ctx, client, partitions)
-
-	for _, ids := range partitions {
-		s.usage.EvictPartitions(ids)
-	}
-}
-
 func (s *IngestLimits) CheckReady(ctx context.Context) error {
 	if s.State() != services.Running {
 		return fmt.Errorf("service is not running: %v", s.State())
diff --git a/pkg/limits/service_test.go b/pkg/limits/service_test.go
@@ -338,10 +338,7 @@ func TestIngestLimits_ExceedsLimits(t *testing.T) {
 			}
 
 			// Assign the Partition IDs.
-			partitions := make(map[string][]int32)
-			partitions["test"] = make([]int32, 0, len(tt.assignedPartitions))
-			partitions["test"] = append(partitions["test"], tt.assignedPartitions...)
-			s.partitionManager.Assign(context.Background(), nil, partitions)
+			s.partitionManager.Assign(context.Background(), tt.assignedPartitions)
 
 			// Call ExceedsLimits.
 			req := &proto.ExceedsLimitsRequest{
@@ -429,8 +426,7 @@ func TestIngestLimits_ExceedsLimits_Concurrent(t *testing.T) {
 	}
 
 	// Assign the Partition IDs.
-	partitions := map[string][]int32{"tenant1": {0}}
-	s.partitionManager.Assign(context.Background(), nil, partitions)
+	s.partitionManager.Assign(context.Background(), []int32{0})
 
 	// Run concurrent requests
 	concurrency := 10
