refactor(aws): page Cost Explorer through one capped helper and cap the unbounded loops

providers/aws/recommendations/ce_paginate.go

@@ -0,0 +1,38 @@
+package recommendations
+
+import (
+	"context"
+	"fmt"
+)
+
+// Runaway guards, not size limits. No Cost Explorer input sets MaxResults, so
+// CE picks the page size; these only stop a token loop that never ends. One
+// over-cap pool aborts the whole GetRICoverageMap survey (coverage.go).
+const (
+	maxCoveragePages    = 100
+	maxUtilizationPages = 100
+)
+
+// paginateCE drives a Cost Explorer NextPageToken loop. page fetches one page
+// for the given token and returns the next token; a nil or empty token ends
+// the loop. Each iteration returns ctx.Err() raw, then enforces maxPages, so a
+// runaway loop errors instead of returning partial results as complete.
+func paginateCE(ctx context.Context, maxPages int, what string, page func(token *string) (next *string, err error)) error {
+	var token *string
+	for n := 0; ; n++ {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		if n >= maxPages {
+			return fmt.Errorf("%s: pagination cap reached after %d pages", what, maxPages)
+		}
+		next, err := page(token)
+		if err != nil {
+			return err
+		}
+		if next == nil || *next == "" {
+			return nil
+		}
+		token = next
+	}
+}

providers/aws/recommendations/ce_paginate_test.go

@@ -0,0 +1,165 @@
+package recommendations
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync/atomic"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/costexplorer"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPaginateCE_FollowsTokens(t *testing.T) {
+	var seen []*string
+	err := paginateCE(context.Background(), 10, "x", func(token *string) (*string, error) {
+		seen = append(seen, token)
+		if len(seen) == 3 {
+			return nil, nil
+		}
+		return aws.String(fmt.Sprintf("t%d", len(seen))), nil
+	})
+	require.NoError(t, err)
+	require.Len(t, seen, 3)
+	assert.Nil(t, seen[0])
+	assert.Equal(t, "t1", *seen[1])
+	assert.Equal(t, "t2", *seen[2])
+}
+
+func TestPaginateCE_EmptyTokenEndsLoop(t *testing.T) {
+	calls := 0
+	err := paginateCE(context.Background(), 10, "x", func(*string) (*string, error) {
+		calls++
+		return aws.String(""), nil
+	})
+	require.NoError(t, err)
+	assert.Equal(t, 1, calls)
+}
+
+func TestPaginateCE_CapErrorsAfterExactlyMaxPages(t *testing.T) {
+	calls := 0
+	err := paginateCE(context.Background(), 5, "thing", func(*string) (*string, error) {
+		calls++
+		return aws.String("more"), nil
+	})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "pagination cap reached")
+	assert.Contains(t, err.Error(), "thing: ")
+	assert.Equal(t, 5, calls)
+}
+
+func TestPaginateCE_ExactlyMaxPagesSucceeds(t *testing.T) {
+	calls := 0
+	err := paginateCE(context.Background(), 3, "x", func(*string) (*string, error) {
+		calls++
+		if calls == 3 {
+			return nil, nil
+		}
+		return aws.String("more"), nil
+	})
+	require.NoError(t, err)
+	assert.Equal(t, 3, calls)
+}
+
+func TestPaginateCE_CanceledBeforeFirstCall(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	calls := 0
+	err := paginateCE(ctx, 5, "x", func(*string) (*string, error) {
+		calls++
+		return nil, nil
+	})
+	assert.Equal(t, context.Canceled, err, "ctx error must be returned raw")
+	assert.Zero(t, calls)
+}
+
+func TestPaginateCE_CanceledMidLoop(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	calls := 0
+	err := paginateCE(ctx, 5, "x", func(*string) (*string, error) {
+		calls++
+		cancel()
+		return aws.String("more"), nil
+	})
+	assert.Equal(t, context.Canceled, err)
+	assert.Equal(t, 1, calls)
+}
+
+func TestPaginateCE_CallbackErrorPropagates(t *testing.T) {
+	boom := errors.New("boom")
+	err := paginateCE(context.Background(), 5, "x", func(*string) (*string, error) { return nil, boom })
+	assert.Equal(t, boom, err)
+}
+
+// endlessCeiling stops a regressed cap from spinning until the test timeout.
+const endlessCeiling = 500
+
+// endlessCE returns a non-empty NextPageToken for every coverage and
+// utilization call.
+type endlessCE struct {
+	mockCostExplorerAPI
+	covCalls  atomic.Int32
+	utilCalls atomic.Int32
+}
+
+func (m *endlessCE) GetReservationCoverage(_ context.Context, _ *costexplorer.GetReservationCoverageInput, _ ...func(*costexplorer.Options)) (*costexplorer.GetReservationCoverageOutput, error) {
+	n := m.covCalls.Add(1)
+	if n > endlessCeiling {
+		return nil, errors.New("mock: runaway pagination")
+	}
+	return &costexplorer.GetReservationCoverageOutput{NextPageToken: aws.String(fmt.Sprintf("t%d", n))}, nil
+}
+
+func (m *endlessCE) GetReservationUtilization(_ context.Context, _ *costexplorer.GetReservationUtilizationInput, _ ...func(*costexplorer.Options)) (*costexplorer.GetReservationUtilizationOutput, error) {
+	n := m.utilCalls.Add(1)
+	if n > endlessCeiling {
+		return nil, errors.New("mock: runaway pagination")
+	}
+	return &costexplorer.GetReservationUtilizationOutput{NextPageToken: aws.String(fmt.Sprintf("t%d", n))}, nil
+}
+
+func TestFetchCoveragePaged_CapsEndlessTokens(t *testing.T) {
+	assert.Equal(t, 100, maxCoveragePages)
+	mock := &endlessCE{}
+	client := NewClientWithAPI(mock, "us-east-1")
+	err := client.fetchCoveragePaged(context.Background(), &costexplorer.GetReservationCoverageInput{},
+		func(string, string, PoolCoverage) { t.Fatal("partial results must not be recorded") }, 1)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "pagination cap reached")
+	assert.Equal(t, int32(maxCoveragePages), mock.covCalls.Load())
+}
+
+func TestGetRIUtilization_CapsEndlessTokens(t *testing.T) {
+	assert.Equal(t, 100, maxUtilizationPages)
+	mock := &endlessCE{}
+	client := NewClientWithAPI(mock, "us-east-1")
+	_, err := client.GetRIUtilization(context.Background(), 30, "us-east-1")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "pagination cap reached")
+	assert.Equal(t, int32(maxUtilizationPages), mock.utilCalls.Load())
+}
+
+func TestGetDailyUsagePcts_CapsEndlessTokens(t *testing.T) {
+	assert.Equal(t, 20, maxRecommendationPages)
+	mock := &endlessCE{}
+	client := NewClientWithAPI(mock, "us-east-1")
+	_, err := client.GetDailyUsagePcts(context.Background(), "Amazon Elastic Compute Cloud - Compute", "m5.large", "us-east-1")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to get daily coverage for")
+	assert.Contains(t, err.Error(), "pagination cap reached")
+	assert.Equal(t, int32(maxRecommendationPages), mock.covCalls.Load())
+}
+
+func TestGetDailyUsagePcts_PreCanceledContext(t *testing.T) {
+	mock := &endlessCE{}
+	client := NewClientWithAPI(mock, "us-east-1")
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	_, err := client.GetDailyUsagePcts(ctx, "Amazon Elastic Compute Cloud - Compute", "m5.large", "us-east-1")
+	require.Error(t, err)
+	assert.ErrorIs(t, err, context.Canceled)
+	assert.Zero(t, mock.covCalls.Load())
+}

providers/aws/recommendations/coverage.go

@@ -317,25 +317,21 @@ func (c *Client) fetchCoveragePaged(
 	windowHours float64,
 ) error {
 	acc := make(riCoverageAccumulator)
-	var token *string
-	for {
-		if err := ctx.Err(); err != nil {
-			return fmt.Errorf("coverage: pagination canceled: %w", err)
-		}
+	err := paginateCE(ctx, maxCoveragePages, "coverage", func(token *string) (*string, error) {
 		input.NextPageToken = token
 		result, err := c.fetchCoveragePage(ctx, input)
 		if err != nil {
-			return err
+			return nil, err
 		}
 		for _, period := range result.CoveragesByTime {
 			for _, group := range period.Groups {
 				acc.addGroup(group, windowHours)
 			}
 		}
-		if result.NextPageToken == nil || *result.NextPageToken == "" {
-			break
-		}
-		token = result.NextPageToken
+		return result.NextPageToken, nil
+	})
+	if err != nil {
+		return err
 	}
 	acc.record(record)
 	return nil

providers/aws/recommendations/usage_history.go

@@ -87,20 +87,19 @@ func (c *Client) fetchDailyCoverage(ctx context.Context, serviceFilter, resource
 	}
 
 	anyData := false
-	var token *string
-	for {
+	err := paginateCE(ctx, maxRecommendationPages, "GetReservationCoverage", func(token *string) (*string, error) {
 		input.NextPageToken = token
 		result, err := c.fetchCoveragePage(ctx, input)
 		if err != nil {
-			return false, fmt.Errorf("failed to get daily coverage for %s/%s/%s: %w", serviceFilter, resourceType, region, err)
+			return nil, err
 		}
 		if applyPeriodsToDayMap(result.CoveragesByTime, dayPct) {
 			anyData = true
 		}
-		if result.NextPageToken == nil || *result.NextPageToken == "" {
-			break
-		}
-		token = result.NextPageToken
+		return result.NextPageToken, nil
+	})
+	if err != nil {
+		return false, fmt.Errorf("failed to get daily coverage for %s/%s/%s: %w", serviceFilter, resourceType, region, err)
 	}
 	return anyData, nil
 }

providers/aws/recommendations/utilization.go

@@ -65,28 +65,21 @@ func (c *Client) GetRIUtilization(ctx context.Context, lookbackDays int, region
 
 	agg := make(map[string]*riAccumulator)
 
-	var nextPageToken *string
-	for {
-		if err := ctx.Err(); err != nil {
-			return nil, fmt.Errorf("utilization: pagination canceled: %w", err)
-		}
-		input.NextPageToken = nextPageToken
-
+	err := paginateCE(ctx, maxUtilizationPages, "utilization", func(token *string) (*string, error) {
+		input.NextPageToken = token
 		result, err := c.fetchUtilizationPage(ctx, input)
 		if err != nil {
 			return nil, err
 		}
-
 		for _, group := range result.UtilizationsByTime {
 			for _, detail := range group.Groups {
 				aggregateUtilizationDetail(agg, detail)
 			}
 		}
-
-		if result.NextPageToken == nil || *result.NextPageToken == "" {
-			break
-		}
-		nextPageToken = result.NextPageToken
+		return result.NextPageToken, nil
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	return buildUtilizations(agg), nil