diff --git a/pkg/aws/awsclient.go b/pkg/aws/awsclient.go index 57d944d..797fd14 100644 --- a/pkg/aws/awsclient.go +++ b/pkg/aws/awsclient.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "os" - "sync" "time" "github.com/aws/aws-sdk-go-v2/aws" @@ -91,6 +90,11 @@ func (c *AWSClient) ListRoles(ctx context.Context) ([]Role, error) { } } + // Nothing to hydrate; skip the progress bar and worker pool entirely. + if len(roles) == 0 { + return roles, nil + } + // Create progress bar bar := progressbar.NewOptions(len(roles), progressbar.OptionEnableColorCodes(true), @@ -108,89 +112,14 @@ func (c *AWSClient) ListRoles(ctx context.Context) ([]Role, error) { }), ) - // Use worker pool pattern for better efficiency - type roleResult struct { - index int - lastUsed *time.Time - err error - } - - // Define our channels - const maxConcurrency = 55 - jobs := make(chan int, len(roles)) - results := make(chan roleResult, len(roles)) - - // Create a cancellable context - ctx, cancel := context.WithCancel(ctx) - defer cancel() - - // Create worker pool - var wg sync.WaitGroup - for w := 0; w < maxConcurrency; w++ { - wg.Add(1) - go func() { - defer wg.Done() - - for i := range jobs { - // Check for cancellation - select { - case <-ctx.Done(): - return - default: - // Continue processing - } - - lastUsed, err := c.GetRoleLastUsed(ctx, roles[i].Name) - select { - case <-ctx.Done(): - return - case results <- roleResult{ - index: i, - lastUsed: lastUsed, - err: err, - }: - // Result sent - } - } - }() - } - - // Send jobs to the workers - for i := range roles { - jobs <- i - } - close(jobs) // No more jobs to send - - // Process results in a separate goroutine - done := make(chan struct{}) - go func() { - defer close(done) - for i := 0; i < len(roles); i++ { - result := <-results - if result.err != nil { - // Log the error but continue processing other roles - fmt.Fprintf(os.Stderr, "Warning: Failed to get last used info for role %s: %v\n", roles[result.index].Name, result.err) - // Continue processing - don't cancel everything for one role's error - roles[result.index].LastUsed = nil - } else { - roles[result.index].LastUsed = result.lastUsed - } - if err := bar.Add(1); err != nil { - // Log the error but continue processing - fmt.Fprintf(os.Stderr, "Failed to update progress bar: %v\n", err) - } + // Hydrate each role's last-used timestamp concurrently with a bounded + // worker pool that writes results directly into the roles slice. + hydrateLastUsed(ctx, roles, defaultListConcurrency, c.GetRoleLastUsed, func() { + if err := bar.Add(1); err != nil { + fmt.Fprintf(os.Stderr, "Failed to update progress bar: %v\n", err) } - }() - - // Wait for workers to finish - go func() { - wg.Wait() - // Close results channel after all workers are done - close(results) - }() + }) - // Wait for either completion or error - <-done return roles, nil } diff --git a/pkg/aws/lastused.go b/pkg/aws/lastused.go new file mode 100644 index 0000000..c89c9d6 --- /dev/null +++ b/pkg/aws/lastused.go @@ -0,0 +1,95 @@ +package aws + +import ( + "context" + "fmt" + "os" + "sync" + "sync/atomic" + "time" +) + +// defaultListConcurrency bounds how many GetRole calls run in parallel while +// hydrating last-used data. IAM's read APIs are rate limited, so this caps the +// burst we send while still keeping the wall-clock time dominated by a single +// wave of requests for typical accounts. +const defaultListConcurrency = 55 + +// lastUsedFetcher resolves the last-used timestamp for a single role. +type lastUsedFetcher func(ctx context.Context, roleName string) (*time.Time, error) + +// hydrateLastUsed populates the LastUsed field of every role in place using a +// bounded worker pool. +// +// Work is handed out through an atomic index counter rather than a jobs +// channel, and because each worker owns a distinct slice index there is no +// results channel or collector goroutine: workers write straight into +// roles[i]. This keeps allocations flat (no per-role channel traffic) and +// avoids spawning more goroutines than there are roles. onProgress, when +// non-nil, is invoked exactly once per completed role under a mutex so +// non-thread-safe sinks (such as a progress bar) stay safe. +func hydrateLastUsed( + ctx context.Context, + roles []Role, + concurrency int, + fetch lastUsedFetcher, + onProgress func(), +) { + if len(roles) == 0 { + return + } + + if concurrency < 1 { + concurrency = 1 + } + if concurrency > len(roles) { + concurrency = len(roles) + } + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + var ( + nextIndex int64 = -1 + progressMu sync.Mutex + wg sync.WaitGroup + ) + + wg.Add(concurrency) + for w := 0; w < concurrency; w++ { + go func() { + defer wg.Done() + + for { + i := int(atomic.AddInt64(&nextIndex, 1)) + if i >= len(roles) { + return + } + + select { + case <-ctx.Done(): + return + default: + } + + lastUsed, err := fetch(ctx, roles[i].Name) + if err != nil { + // Log and continue: one role's failure should not abort + // usage data for the rest of the account. + fmt.Fprintf(os.Stderr, "Warning: Failed to get last used info for role %s: %v\n", roles[i].Name, err) + roles[i].LastUsed = nil + } else { + roles[i].LastUsed = lastUsed + } + + if onProgress != nil { + progressMu.Lock() + onProgress() + progressMu.Unlock() + } + } + }() + } + + wg.Wait() +} diff --git a/pkg/aws/lastused_test.go b/pkg/aws/lastused_test.go new file mode 100644 index 0000000..5da08d8 --- /dev/null +++ b/pkg/aws/lastused_test.go @@ -0,0 +1,190 @@ +package aws + +import ( + "context" + "fmt" + "sync" + "testing" + "time" +) + +// makeRoles builds n roles with deterministic names for benchmarking/testing. +func makeRoles(n int) []Role { + roles := make([]Role, n) + for i := 0; i < n; i++ { + roles[i] = Role{ + Name: fmt.Sprintf("Role-%d", i), + Arn: fmt.Sprintf("arn:aws:iam::123456789012:role/Role-%d", i), + } + } + return roles +} + +// stubFetch returns a fetcher that optionally sleeps to emulate API latency. +func stubFetch(delay time.Duration) lastUsedFetcher { + return func(ctx context.Context, roleName string) (*time.Time, error) { + if delay > 0 { + time.Sleep(delay) + } + t := time.Now() + return &t, nil + } +} + +// hydrateLastUsedChannels is the previous channel-based implementation, kept +// here only so the benchmark can quantify the orchestration overhead removed by +// hydrateLastUsed. It mirrors the old ListRoles worker pool exactly. +func hydrateLastUsedChannels( + ctx context.Context, + roles []Role, + concurrency int, + fetch lastUsedFetcher, + onProgress func(), +) { + type roleResult struct { + index int + lastUsed *time.Time + err error + } + + jobs := make(chan int, len(roles)) + results := make(chan roleResult, len(roles)) + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + var wg sync.WaitGroup + for w := 0; w < concurrency; w++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := range jobs { + select { + case <-ctx.Done(): + return + default: + } + lastUsed, err := fetch(ctx, roles[i].Name) + select { + case <-ctx.Done(): + return + case results <- roleResult{index: i, lastUsed: lastUsed, err: err}: + } + } + }() + } + + for i := range roles { + jobs <- i + } + close(jobs) + + done := make(chan struct{}) + go func() { + defer close(done) + for i := 0; i < len(roles); i++ { + result := <-results + if result.err != nil { + roles[result.index].LastUsed = nil + } else { + roles[result.index].LastUsed = result.lastUsed + } + if onProgress != nil { + onProgress() + } + } + }() + + go func() { + wg.Wait() + close(results) + }() + + <-done +} + +func TestHydrateLastUsedPopulatesEveryRole(t *testing.T) { + roles := makeRoles(200) + var progress int64 + var mu sync.Mutex + + hydrateLastUsed(context.Background(), roles, defaultListConcurrency, stubFetch(0), func() { + mu.Lock() + progress++ + mu.Unlock() + }) + + for i := range roles { + if roles[i].LastUsed == nil { + t.Fatalf("role %d was not hydrated", i) + } + } + if progress != int64(len(roles)) { + t.Fatalf("expected %d progress callbacks, got %d", len(roles), progress) + } +} + +func TestHydrateLastUsedEmptyAndSmall(t *testing.T) { + // Empty input must not panic or spawn work. + hydrateLastUsed(context.Background(), nil, defaultListConcurrency, stubFetch(0), nil) + + // Fewer roles than the concurrency cap must still hydrate every role. + roles := makeRoles(3) + hydrateLastUsed(context.Background(), roles, defaultListConcurrency, stubFetch(0), nil) + for i := range roles { + if roles[i].LastUsed == nil { + t.Fatalf("small batch: role %d not hydrated", i) + } + } +} + +func TestHydrateLastUsedRecordsErrorsAsNil(t *testing.T) { + roles := makeRoles(10) + failing := func(ctx context.Context, roleName string) (*time.Time, error) { + return nil, fmt.Errorf("boom") + } + hydrateLastUsed(context.Background(), roles, defaultListConcurrency, failing, nil) + for i := range roles { + if roles[i].LastUsed != nil { + t.Fatalf("role %d should be nil after fetch error", i) + } + } +} + +func benchmarkHydrate(b *testing.B, fn func(context.Context, []Role, int, lastUsedFetcher, func()), n int, delay time.Duration) { + ctx := context.Background() + roles := makeRoles(n) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + fn(ctx, roles, defaultListConcurrency, stubFetch(delay), nil) + } +} + +// BenchmarkHydrateLastUsed measures the new atomic-dispatch implementation. +func BenchmarkHydrateLastUsed(b *testing.B) { + for _, n := range []int{10, 100, 1000} { + b.Run(fmt.Sprintf("atomic-nodelay-%d", n), func(b *testing.B) { + benchmarkHydrate(b, hydrateLastUsed, n, 0) + }) + } + for _, n := range []int{10, 100, 1000} { + b.Run(fmt.Sprintf("channels-nodelay-%d", n), func(b *testing.B) { + benchmarkHydrate(b, hydrateLastUsedChannels, n, 0) + }) + } +} + +// BenchmarkHydrateLastUsedWithLatency mirrors real API latency so the wall-clock +// throughput of the two pools can be compared head to head. +func BenchmarkHydrateLastUsedWithLatency(b *testing.B) { + const delay = 1 * time.Millisecond + for _, n := range []int{100, 1000} { + b.Run(fmt.Sprintf("atomic-%d", n), func(b *testing.B) { + benchmarkHydrate(b, hydrateLastUsed, n, delay) + }) + b.Run(fmt.Sprintf("channels-%d", n), func(b *testing.B) { + benchmarkHydrate(b, hydrateLastUsedChannels, n, delay) + }) + } +}