From c609713390214d64bd15483a9a03536041bcb56b Mon Sep 17 00:00:00 2001 From: Tim Schindler Date: Thu, 19 Feb 2026 08:19:52 +0100 Subject: [PATCH] perf: parallelize grant status API calls Fire sessions + all-CSP eligibility requests concurrently via fetchStatusData(), reducing wall-clock time from sum to max of individual calls (~2s saving in practice). --- CHANGELOG.md | 4 + cmd/helpers.go | 87 +++++++++++++ cmd/helpers_test.go | 294 ++++++++++++++++++++++++++++++++++++++++++++ cmd/status.go | 15 +-- cmd/status_test.go | 23 ---- 5 files changed, 391 insertions(+), 32 deletions(-) create mode 100644 cmd/helpers_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 72aa571..6a3bfb3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,10 @@ All notable changes to this project will be documented in this file. - `--provider`/`-p` flag on `grant revoke --all` and interactive mode to filter by cloud provider - Session ID displayed in `grant status` output for easy reference with `grant revoke` +### Changed + +- `grant status` now fetches sessions and eligibility data concurrently, reducing wall-clock time by ~2s + ### Fixed - `grant revoke` now rejects `--provider` in direct mode (session IDs are already explicit) diff --git a/cmd/helpers.go b/cmd/helpers.go index 7299ca5..04b5712 100644 --- a/cmd/helpers.go +++ b/cmd/helpers.go @@ -4,10 +4,97 @@ import ( "context" "fmt" "io" + "sync" scamodels "github.com/aaearon/grant-cli/internal/sca/models" ) +// statusData holds the results of concurrent sessions + eligibility fetches. +type statusData struct { + sessions *scamodels.SessionsResponse + nameMap map[string]string +} + +// fetchStatusData fires sessions and all-CSP eligibility calls concurrently, +// then joins results. A sessions error is fatal; eligibility errors are +// gracefully degraded (empty nameMap entry, verbose warning). +func fetchStatusData( + ctx context.Context, + sessionLister sessionLister, + eligLister eligibilityLister, + cspFilter *scamodels.CSP, + errWriter io.Writer, +) (*statusData, error) { + type eligResult struct { + csp scamodels.CSP + targets []scamodels.EligibleTarget + err error + } + + var ( + sessions *scamodels.SessionsResponse + sessionsErr error + wg sync.WaitGroup + ) + + // Goroutine 1: fetch sessions + wg.Add(1) + go func() { + defer wg.Done() + sessions, sessionsErr = sessionLister.ListSessions(ctx, cspFilter) + }() + + // Determine which CSPs to query for eligibility + cspsToQuery := supportedCSPs + if cspFilter != nil { + cspsToQuery = []scamodels.CSP{*cspFilter} + } + + // Goroutines 2..N: fetch eligibility for each CSP + eligResults := make(chan eligResult, len(cspsToQuery)) + for _, csp := range cspsToQuery { + wg.Add(1) + go func(csp scamodels.CSP) { + defer wg.Done() + resp, err := eligLister.ListEligibility(ctx, csp) + if err != nil || resp == nil { + eligResults <- eligResult{csp: csp, err: err} + return + } + eligResults <- eligResult{csp: csp, targets: resp.Response} + }(csp) + } + + // Close channel after all goroutines finish + go func() { + wg.Wait() + close(eligResults) + }() + + // Build nameMap from eligibility results + nameMap := make(map[string]string) + for r := range eligResults { + if r.err != nil { + if verbose { + fmt.Fprintf(errWriter, "Warning: failed to fetch names for %s: %v\n", r.csp, r.err) + } + continue + } + for _, t := range r.targets { + if t.WorkspaceName != "" { + nameMap[t.WorkspaceID] = t.WorkspaceName + } + } + } + + // Check sessions result (goroutine has finished since channel is drained after wg.Wait) + if sessionsErr != nil { + return nil, fmt.Errorf("failed to list sessions: %w", sessionsErr) + } + + return &statusData{sessions: sessions, nameMap: nameMap}, nil +} + // buildWorkspaceNameMap fetches eligibility for each unique CSP in sessions // and builds a workspaceID -> workspaceName map. Errors are silently ignored // (graceful degradation — the raw workspace ID is shown as fallback). diff --git a/cmd/helpers_test.go b/cmd/helpers_test.go new file mode 100644 index 0000000..11675c0 --- /dev/null +++ b/cmd/helpers_test.go @@ -0,0 +1,294 @@ +package cmd + +import ( + "bytes" + "context" + "errors" + "strings" + "testing" + + scamodels "github.com/aaearon/grant-cli/internal/sca/models" +) + +func TestFetchStatusData(t *testing.T) { + tests := []struct { + name string + setupSessions func() *mockSessionLister + setupEligibility func() *mockEligibilityLister + cspFilter *scamodels.CSP + wantErr bool + wantErrContain string + wantSessions int + wantNameMapKeys []string + }{ + { + name: "both sessions and eligibility succeed", + setupSessions: func() *mockSessionLister { + return &mockSessionLister{ + sessions: &scamodels.SessionsResponse{ + Response: []scamodels.SessionInfo{ + {SessionID: "s1", CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + {SessionID: "s2", CSP: scamodels.CSPAWS, WorkspaceID: "arn:aws:iam::123:role/Admin"}, + }, + Total: 2, + }, + } + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listFunc: func(ctx context.Context, csp scamodels.CSP) (*scamodels.EligibilityResponse, error) { + switch csp { + case scamodels.CSPAzure: + return &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "/subscriptions/sub-1", WorkspaceName: "Dev Sub"}, + }, + }, nil + case scamodels.CSPAWS: + return &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "arn:aws:iam::123:role/Admin", WorkspaceName: "Admin Account"}, + }, + }, nil + } + return &scamodels.EligibilityResponse{}, nil + }, + } + }, + wantSessions: 2, + wantNameMapKeys: []string{"/subscriptions/sub-1", "arn:aws:iam::123:role/Admin"}, + }, + { + name: "sessions error propagated", + setupSessions: func() *mockSessionLister { + return &mockSessionLister{ + listErr: errors.New("API error: service unavailable"), + } + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{} + }, + wantErr: true, + wantErrContain: "failed to list sessions", + }, + { + name: "single eligibility error - graceful degradation", + setupSessions: func() *mockSessionLister { + return &mockSessionLister{ + sessions: &scamodels.SessionsResponse{ + Response: []scamodels.SessionInfo{ + {SessionID: "s1", CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + }, + Total: 1, + }, + } + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listFunc: func(ctx context.Context, csp scamodels.CSP) (*scamodels.EligibilityResponse, error) { + if csp == scamodels.CSPAWS { + return nil, errors.New("AWS unavailable") + } + return &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "/subscriptions/sub-1", WorkspaceName: "Dev Sub"}, + }, + }, nil + }, + } + }, + wantSessions: 1, + wantNameMapKeys: []string{"/subscriptions/sub-1"}, + }, + { + name: "all eligibility errors - empty name map, sessions still returned", + setupSessions: func() *mockSessionLister { + return &mockSessionLister{ + sessions: &scamodels.SessionsResponse{ + Response: []scamodels.SessionInfo{ + {SessionID: "s1", CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + }, + Total: 1, + }, + } + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listErr: errors.New("eligibility unavailable"), + } + }, + wantSessions: 1, + wantNameMapKeys: nil, + }, + { + name: "no sessions - empty name map", + setupSessions: func() *mockSessionLister { + return &mockSessionLister{ + sessions: &scamodels.SessionsResponse{ + Response: []scamodels.SessionInfo{}, + Total: 0, + }, + } + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + response: &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "/subscriptions/sub-1", WorkspaceName: "Dev Sub"}, + }, + }, + } + }, + wantSessions: 0, + }, + { + name: "with provider filter - only filtered CSP eligibility called", + setupSessions: func() *mockSessionLister { + cspAzure := scamodels.CSPAzure + return &mockSessionLister{ + listFunc: func(ctx context.Context, csp *scamodels.CSP) (*scamodels.SessionsResponse, error) { + if csp == nil || *csp != cspAzure { + return nil, errors.New("expected azure filter") + } + return &scamodels.SessionsResponse{ + Response: []scamodels.SessionInfo{ + {SessionID: "s1", CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + }, + Total: 1, + }, nil + }, + } + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listFunc: func(ctx context.Context, csp scamodels.CSP) (*scamodels.EligibilityResponse, error) { + if csp != scamodels.CSPAzure { + return nil, errors.New("should not call non-Azure eligibility") + } + return &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "/subscriptions/sub-1", WorkspaceName: "Dev Sub"}, + }, + }, nil + }, + } + }, + cspFilter: func() *scamodels.CSP { c := scamodels.CSPAzure; return &c }(), + wantSessions: 1, + wantNameMapKeys: []string{"/subscriptions/sub-1"}, + }, + { + name: "context cancellation - no hang", + setupSessions: func() *mockSessionLister { + return &mockSessionLister{ + listFunc: func(ctx context.Context, csp *scamodels.CSP) (*scamodels.SessionsResponse, error) { + return nil, ctx.Err() + }, + } + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{} + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + if tt.name == "context cancellation - no hang" { + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(ctx) + cancel() + } + + sessionLister := tt.setupSessions() + eligLister := tt.setupEligibility() + var errBuf bytes.Buffer + + data, err := fetchStatusData(ctx, sessionLister, eligLister, tt.cspFilter, &errBuf) + + if tt.wantErr { + if err == nil { + t.Fatal("expected error but got none") + } + if tt.wantErrContain != "" && !strings.Contains(err.Error(), tt.wantErrContain) { + t.Errorf("error %q should contain %q", err.Error(), tt.wantErrContain) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(data.sessions.Response) != tt.wantSessions { + t.Errorf("got %d sessions, want %d", len(data.sessions.Response), tt.wantSessions) + } + + for _, key := range tt.wantNameMapKeys { + if _, ok := data.nameMap[key]; !ok { + t.Errorf("nameMap missing key %q, got: %v", key, data.nameMap) + } + } + }) + } +} + +func TestFetchStatusData_VerboseWarning(t *testing.T) { + oldVerbose := verbose + verbose = true + defer func() { verbose = oldVerbose }() + + ctx := context.Background() + sessionLister := &mockSessionLister{ + sessions: &scamodels.SessionsResponse{ + Response: []scamodels.SessionInfo{ + {SessionID: "s1", CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + }, + Total: 1, + }, + } + eligLister := &mockEligibilityLister{ + listErr: errors.New("eligibility API unavailable"), + } + + var errBuf bytes.Buffer + data, err := fetchStatusData(ctx, sessionLister, eligLister, nil, &errBuf) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(data.sessions.Response) != 1 { + t.Errorf("expected 1 session, got %d", len(data.sessions.Response)) + } + if len(data.nameMap) != 0 { + t.Errorf("expected empty nameMap, got %v", data.nameMap) + } + + // Verify verbose warnings were written for both CSPs + errOutput := errBuf.String() + if !strings.Contains(errOutput, "Warning:") { + t.Errorf("expected verbose warning, got: %q", errOutput) + } +} + +func TestBuildWorkspaceNameMap_VerboseWarning(t *testing.T) { + oldVerbose := verbose + verbose = true + defer func() { verbose = oldVerbose }() + + ctx := context.Background() + eligLister := &mockEligibilityLister{ + listErr: errors.New("eligibility API unavailable"), + } + + sessions := []scamodels.SessionInfo{ + {CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + } + + var buf bytes.Buffer + _ = buildWorkspaceNameMap(ctx, eligLister, sessions, &buf) + + if !strings.Contains(buf.String(), "Warning: failed to fetch names for AZURE") { + t.Errorf("expected verbose warning, got: %q", buf.String()) + } +} diff --git a/cmd/status.go b/cmd/status.go index 854c948..f07daa9 100644 --- a/cmd/status.go +++ b/cmd/status.go @@ -67,32 +67,29 @@ func runStatus(cmd *cobra.Command, authLoader authLoader, sessionLister sessionL cspFilter = &csp } - // List sessions + // Fetch sessions and eligibility concurrently ctx, cancel := context.WithTimeout(context.Background(), apiTimeout) defer cancel() - sessions, err := sessionLister.ListSessions(ctx, cspFilter) + data, err := fetchStatusData(ctx, sessionLister, eligLister, cspFilter, cmd.ErrOrStderr()) if err != nil { - return fmt.Errorf("failed to list sessions: %w", err) + return err } // Display sessions - if len(sessions.Response) == 0 { + if len(data.sessions.Response) == 0 { fmt.Fprintf(cmd.OutOrStdout(), "\nNo active sessions.\n") return nil } - // Build workspace name map from eligibility data - nameMap := buildWorkspaceNameMap(ctx, eligLister, sessions.Response, cmd.ErrOrStderr()) - // Group sessions by provider - sessionsByProvider := groupSessionsByProvider(sessions.Response) + sessionsByProvider := groupSessionsByProvider(data.sessions.Response) // Display grouped sessions fmt.Fprintf(cmd.OutOrStdout(), "\n") for _, p := range sortedProviders(sessionsByProvider) { fmt.Fprintf(cmd.OutOrStdout(), "%s sessions:\n", formatProviderName(p)) for _, session := range sessionsByProvider[p] { - fmt.Fprintf(cmd.OutOrStdout(), " %s\n", ui.FormatSessionOption(session, nameMap)) + fmt.Fprintf(cmd.OutOrStdout(), " %s\n", ui.FormatSessionOption(session, data.nameMap)) } } diff --git a/cmd/status_test.go b/cmd/status_test.go index 5a89487..4457e76 100644 --- a/cmd/status_test.go +++ b/cmd/status_test.go @@ -1,7 +1,6 @@ package cmd import ( - "bytes" "context" "errors" "strings" @@ -550,28 +549,6 @@ func TestStatusCommandIntegration(t *testing.T) { } } -func TestBuildWorkspaceNameMap_VerboseWarning(t *testing.T) { - // Set verbose to true to trigger the warning - oldVerbose := verbose - verbose = true - defer func() { verbose = oldVerbose }() - - ctx := context.Background() - eligLister := &mockEligibilityLister{ - listErr: errors.New("eligibility API unavailable"), - } - - sessions := []scamodels.SessionInfo{ - {CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, - } - - var buf bytes.Buffer - _ = buildWorkspaceNameMap(ctx, eligLister, sessions, &buf) - - if !strings.Contains(buf.String(), "Warning: failed to fetch names for AZURE") { - t.Errorf("expected verbose warning, got: %q", buf.String()) - } -} func TestStatusCommandUsage(t *testing.T) { cmd := NewStatusCommand()