diff --git a/cmd/allstar/main.go b/cmd/allstar/main.go index b57c7b6a..357a11f0 100644 --- a/cmd/allstar/main.go +++ b/cmd/allstar/main.go @@ -62,6 +62,8 @@ func main() { specificPolicyArg := flag.String("policy", "", fmt.Sprintf("Run a specific policy check. Supported policies: %s", supportedPoliciesMsg)) specificRepoArg := flag.String("repo", "", "Run on a specific \"owner/repo\". For example \"ossf/allstar\"") + numWorkersArg := flag.Int("workers", 5, "maximum number of active goroutines for Allstar scans") + flag.Parse() if *specificPolicyArg != "" { @@ -81,7 +83,7 @@ func main() { } if runOnce { - _, err := enforce.EnforceAll(ctx, ghc, *specificPolicyArg, *specificRepoArg) + _, err := enforce.EnforceAll(ctx, ghc, *specificPolicyArg, *specificRepoArg, *numWorkersArg) if err != nil { log.Fatal(). Err(err). @@ -94,7 +96,7 @@ func main() { go func() { defer wg.Done() log.Info(). - Err(enforce.EnforceJob(ctx, ghc, (5 * time.Minute), *specificPolicyArg, *specificRepoArg)). + Err(enforce.EnforceJob(ctx, ghc, (5 * time.Minute), *specificPolicyArg, *specificRepoArg, *numWorkersArg)). Msg("Enforce job shutting down.") }() sigs := make(chan os.Signal, 1) diff --git a/pkg/enforce/enforce.go b/pkg/enforce/enforce.go index cdd48cc6..8293276c 100644 --- a/pkg/enforce/enforce.go +++ b/pkg/enforce/enforce.go @@ -67,7 +67,7 @@ func init() { // // TBD: determine if this should remain exported, or if it will only be called // from EnforceJob. -func EnforceAll(ctx context.Context, ghc ghclients.GhClientsInterface, specificPolicyArg string, specificRepoArg string) (EnforceAllResults, error) { +func EnforceAll(ctx context.Context, ghc ghclients.GhClientsInterface, specificPolicyArg string, specificRepoArg string, numWorkersArg int) (EnforceAllResults, error) { var repoCount int var enforceAllResults = make(EnforceAllResults) ac, err := ghc.Get(0) @@ -85,7 +85,7 @@ func EnforceAll(ctx context.Context, ghc ghclients.GhClientsInterface, specificP Msg("Enforcing policies on installations.") g, ctx := errgroup.WithContext(ctx) - g.SetLimit(5) + g.SetLimit(numWorkersArg) var mu sync.Mutex for _, i := range insts { @@ -302,9 +302,9 @@ func getAppInstallationReposReal(ctx context.Context, ic *github.Client) ([]*git // EnforceJob is a reconciliation job that enforces policies on all repos every // d duration. It runs forever until the context is done. -func EnforceJob(ctx context.Context, ghc *ghclients.GHClients, d time.Duration, specificPolicyArg string, specificRepoArg string) error { +func EnforceJob(ctx context.Context, ghc *ghclients.GHClients, d time.Duration, specificPolicyArg string, specificRepoArg string, numWorkersArg int) error { for { - _, err := EnforceAll(ctx, ghc, specificPolicyArg, specificRepoArg) + _, err := EnforceAll(ctx, ghc, specificPolicyArg, specificRepoArg, numWorkersArg) if err != nil { log.Error(). Err(err). diff --git a/pkg/enforce/enforce_test.go b/pkg/enforce/enforce_test.go index a45fda66..dc5b46c1 100644 --- a/pkg/enforce/enforce_test.go +++ b/pkg/enforce/enforce_test.go @@ -549,7 +549,8 @@ func TestEnforceAll(t *testing.T) { policy1Results = test.Policy1Results policy2Results = test.Policy2Results - enforceAllResults, err := EnforceAll(context.Background(), mockGhc, "", "") + numWorkers := 1 + enforceAllResults, err := EnforceAll(context.Background(), mockGhc, "", "", numWorkers) if err != nil { t.Fatalf("Unexpected error: %v", err) } @@ -581,7 +582,8 @@ func TestSuspendedEnforce(t *testing.T) { } suspended = false gaicalled = false - if _, err := EnforceAll(context.Background(), &MockGhClients{}, "", ""); err != nil { + numWorkers := 1 + if _, err := EnforceAll(context.Background(), &MockGhClients{}, "", "", numWorkers); err != nil { t.Fatalf("Unexpected error: %v", err) } if !gaicalled { @@ -589,7 +591,7 @@ func TestSuspendedEnforce(t *testing.T) { } suspended = true gaicalled = false - if _, err := EnforceAll(context.Background(), &MockGhClients{}, "", ""); err != nil { + if _, err := EnforceAll(context.Background(), &MockGhClients{}, "", "", numWorkers); err != nil { t.Fatalf("Unexpected error: %v", err) } if gaicalled {