From 890605067ffbb4364625859c26ecff1d8f92213c Mon Sep 17 00:00:00 2001 From: Cristian Magherusan-Stanciu Date: Fri, 16 Oct 2026 20:55:27 +0200 Subject: [PATCH] refactor(aws): page Cost Explorer through one capped helper and cap the unbounded loops Add paginateCE, a shared NextPageToken loop that checks ctx.Err() (returned raw) on each iteration, enforces a page cap, and ends on a nil or empty token. Hitting the cap is an error, never partial results reported as complete. Migrate the three loops that had no cap: fetchCoveragePaged (100 pages), GetRIUtilization (100 pages) and fetchDailyCoverage (maxRecommendationPages, 20, and now with a per-page ctx check). The 100-page caps are runaway guards, not size limits: no Cost Explorer input sets MaxResults, so CE picks the page size. The already-capped sites move over separately. Refs #35 Co-Authored-By: Claude Opus 5.5 --- providers/aws/recommendations/ce_paginate.go | 38 ++++ .../aws/recommendations/ce_paginate_test.go | 165 ++++++++++++++++++ providers/aws/recommendations/coverage.go | 16 +- .../aws/recommendations/usage_history.go | 13 +- providers/aws/recommendations/utilization.go | 19 +- 5 files changed, 221 insertions(+), 30 deletions(-) create mode 100644 providers/aws/recommendations/ce_paginate.go create mode 100644 providers/aws/recommendations/ce_paginate_test.go diff --git a/providers/aws/recommendations/ce_paginate.go b/providers/aws/recommendations/ce_paginate.go new file mode 100644 index 00000000..5debf795 --- /dev/null +++ b/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 + } +} diff --git a/providers/aws/recommendations/ce_paginate_test.go b/providers/aws/recommendations/ce_paginate_test.go new file mode 100644 index 00000000..b96b1ad9 --- /dev/null +++ b/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()) +} diff --git a/providers/aws/recommendations/coverage.go b/providers/aws/recommendations/coverage.go index 654e3def..53ad2b41 100644 --- a/providers/aws/recommendations/coverage.go +++ b/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 diff --git a/providers/aws/recommendations/usage_history.go b/providers/aws/recommendations/usage_history.go index 8f7f721b..d223e1dc 100644 --- a/providers/aws/recommendations/usage_history.go +++ b/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 } diff --git a/providers/aws/recommendations/utilization.go b/providers/aws/recommendations/utilization.go index 2f6b600a..e9839f91 100644 --- a/providers/aws/recommendations/utilization.go +++ b/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