Fix GetUsersCloseToLimit too

Signed-off-by: Oleg Zaytsev <mail@olegzaytsev.com>

Author: colega

pkg/usagetracker/tracker.go

@@ -665,15 +665,12 @@ func (t *UsageTracker) stop(_ error) error {
 
 // TrackSeries implements usagetrackerpb.UsageTrackerServer.
 func (t *UsageTracker) TrackSeries(_ context.Context, req *usagetrackerpb.TrackSeriesRequest) (*usagetrackerpb.TrackSeriesResponse, error) {
-	t.partitionsMtx.RLock()
-	p, ok := t.partitions[req.Partition]
-	t.partitionsMtx.RUnlock()
-	if !ok {
-		return nil, fmt.Errorf("partition handler %d not found", req.Partition)
-	}
-	if p.State() != services.Running {
-		return nil, fmt.Errorf("partition handler %d is not running (state: %s)", req.Partition, p.State())
+	partition := req.Partition
+	p, err := t.runningPartition(partition)
+	if err != nil {
+		return nil, err
 	}
+
 	rejected, err := p.store.trackSeries(context.Background(), req.UserID, req.SeriesHashes, time.Now())
 	if err != nil {
 		return nil, err
@@ -684,12 +681,9 @@ func (t *UsageTracker) TrackSeries(_ context.Context, req *usagetrackerpb.TrackS
 // GetUsersCloseToLimit implements usagetrackerpb.UsageTrackerServer.
 func (t *UsageTracker) GetUsersCloseToLimit(_ context.Context, req *usagetrackerpb.GetUsersCloseToLimitRequest) (*usagetrackerpb.GetUsersCloseToLimitResponse, error) {
 	partition := req.Partition
-
-	t.partitionsMtx.RLock()
-	p, ok := t.partitions[partition]
-	t.partitionsMtx.RUnlock()
-	if !ok {
-		return nil, fmt.Errorf("partition handler %d not found", partition)
+	p, err := t.runningPartition(partition)
+	if err != nil {
+		return nil, err
 	}
 
 	userIDs := p.store.getSortedUsersCloseToLimit()
@@ -699,6 +693,19 @@ func (t *UsageTracker) GetUsersCloseToLimit(_ context.Context, req *usagetracker
 	}, nil
 }
 
+func (t *UsageTracker) runningPartition(partition int32) (*partitionHandler, error) {
+	t.partitionsMtx.RLock()
+	p, ok := t.partitions[partition]
+	t.partitionsMtx.RUnlock()
+	if !ok {
+		return nil, fmt.Errorf("partition handler %d not found", partition)
+	}
+	if p.State() != services.Running {
+		return nil, fmt.Errorf("partition handler %d is not running (state: %s)", partition, p.State())
+	}
+	return p, nil
+}
+
 // CheckReady performs a readiness check.
 // An instance is ready when it has instantiated all the partitions that should belong to it according to the ring.
 func (t *UsageTracker) CheckReady(_ context.Context) error {

pkg/usagetracker/tracker_test.go

@@ -552,44 +552,61 @@ func TestUsageTracker_PartitionAssignment(t *testing.T) {
 }
 
 func TestUsageTracker_GetUsersCloseToLimit(t *testing.T) {
-	makeSeries := func(n int) []uint64 {
-		series := make([]uint64, n)
-		for i := range series {
-			series[i] = uint64(i)
+	t.Run("happy case", func(t *testing.T) {
+		makeSeries := func(n int) []uint64 {
+			series := make([]uint64, n)
+			for i := range series {
+				series[i] = uint64(i)
+			}
+			return series
 		}
-		return series
-	}
-
-	tracker := newReadyTestUsageTracker(t, map[string]*validation.Limits{
-		"a": {MaxActiveSeriesPerUser: 1000 * testPartitionsCount},
-		"b": {MaxActiveSeriesPerUser: 2000 * testPartitionsCount},
-		"c": {MaxActiveSeriesPerUser: 1000 * testPartitionsCount},
-		"d": {MaxActiveSeriesPerUser: 2000 * testPartitionsCount},
-		"e": {MaxActiveSeriesPerUser: 1000 * testPartitionsCount},
-	})
 
-	for _, tenant := range []string{"a", "b", "c", "d", "e"} {
-		resp, err := tracker.TrackSeries(t.Context(), &usagetrackerpb.TrackSeriesRequest{
-			UserID:       tenant,
-			Partition:    0,
-			SeriesHashes: makeSeries(900),
+		tracker := newReadyTestUsageTracker(t, map[string]*validation.Limits{
+			"a": {MaxActiveSeriesPerUser: 1000 * testPartitionsCount},
+			"b": {MaxActiveSeriesPerUser: 2000 * testPartitionsCount},
+			"c": {MaxActiveSeriesPerUser: 1000 * testPartitionsCount},
+			"d": {MaxActiveSeriesPerUser: 2000 * testPartitionsCount},
+			"e": {MaxActiveSeriesPerUser: 1000 * testPartitionsCount},
 		})
-		require.NoError(t, err)
-		require.Empty(t, resp.RejectedSeriesHashes)
-	}
 
-	// Call updateLimits (on all partitions, although we only need partition 0.
-	withRLock(&tracker.partitionsMtx, func() {
-		for _, p := range tracker.partitions {
-			done := make(chan struct{})
-			p.forceUpdateLimitsForTests <- done
-			<-done
+		for _, tenant := range []string{"a", "b", "c", "d", "e"} {
+			resp, err := tracker.TrackSeries(t.Context(), &usagetrackerpb.TrackSeriesRequest{
+				UserID:       tenant,
+				Partition:    0,
+				SeriesHashes: makeSeries(900),
+			})
+			require.NoError(t, err)
+			require.Empty(t, resp.RejectedSeriesHashes)
 		}
+
+		// Call updateLimits (on all partitions, although we only need partition 0.
+		withRLock(&tracker.partitionsMtx, func() {
+			for _, p := range tracker.partitions {
+				done := make(chan struct{})
+				p.forceUpdateLimitsForTests <- done
+				<-done
+			}
+		})
+
+		resp, err := tracker.GetUsersCloseToLimit(t.Context(), &usagetrackerpb.GetUsersCloseToLimitRequest{Partition: 0})
+		require.NoError(t, err)
+		require.Equal(t, []string{"a", "c", "e"}, resp.SortedUserIds, "List of users close to the limit should be sorted lexicographically")
 	})
 
-	resp, err := tracker.GetUsersCloseToLimit(t.Context(), &usagetrackerpb.GetUsersCloseToLimitRequest{Partition: 0})
-	require.NoError(t, err)
-	require.Equal(t, []string{"a", "c", "e"}, resp.SortedUserIds, "List of users close to the limit should be sorted lexicographically")
+	t.Run("partition handler is not running", func(t *testing.T) {
+		tracker := newReadyTestUsageTracker(t, map[string]*validation.Limits{
+			"a": {MaxActiveSeriesPerUser: 1000 * testPartitionsCount},
+		})
+
+		// Call updateLimits (on all partitions, although we only need partition 0.
+		withRLock(&tracker.partitionsMtx, func() {
+			require.NoError(t, services.StopAndAwaitTerminated(context.Background(), tracker.partitions[0]))
+		})
+
+		_, err := tracker.GetUsersCloseToLimit(t.Context(), &usagetrackerpb.GetUsersCloseToLimitRequest{Partition: 0})
+		require.Error(t, err)
+		require.ErrorContains(t, err, "partition handler 0 is not running (state: Terminated)")
+	})
 }
 
 func callPrepareDownscaleEndpoint(t *testing.T, ut *UsageTracker, method string) {