perf(aws): replace channel-based last-used fetch with atomic worker pool

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
 }
 

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()
+}