diff --git a/cmd/virtwork/main.go b/cmd/virtwork/main.go index 8082955..f3fbc73 100644 --- a/cmd/virtwork/main.go +++ b/cmd/virtwork/main.go @@ -125,6 +125,11 @@ func runE(cmd *cobra.Command, args []string) error { return fmt.Errorf("loading config: %w", err) } + workloadNames, _ := cmd.Flags().GetStringSlice("workloads") + if err := workloads.ValidateWorkloadNames(workloadNames); err != nil { + return err + } + verbose, _ := cmd.Flags().GetBool("verbose") logger := logging.NewLogger(cmd.OutOrStdout(), verbose) @@ -179,7 +184,6 @@ func runE(cmd *cobra.Command, args []string) error { logger.Warn("audit record failed", slog.String("op", "RecordEvent"), slog.String("error", auditErr.Error())) } - workloadNames, _ := cmd.Flags().GetStringSlice("workloads") vmCountFlag, _ := cmd.Flags().GetInt("vm-count") ro := orchestrator.NewRunOrchestrator(logger, c, cfg, auditor, cmd.OutOrStdout()) @@ -365,7 +369,11 @@ func printSummary(logger *slog.Logger, result *orchestrator.RunResult, cfg *conf slog.String("container_image", cfg.ContainerDiskImage)) } -func printCleanupPreview(logger *slog.Logger, preview *cleanup.CleanupPreview, namespace, runID string) { +func printCleanupPreview( + logger *slog.Logger, + preview *cleanup.CleanupPreview, + namespace, runID string, +) { attrs := []slog.Attr{ slog.String("namespace", namespace), slog.Int("vms", preview.VMCount), diff --git a/cmd/virtwork/main_test.go b/cmd/virtwork/main_test.go index 3cf9464..22aa672 100644 --- a/cmd/virtwork/main_test.go +++ b/cmd/virtwork/main_test.go @@ -229,6 +229,26 @@ var _ = Describe("Run command flags", func() { Expect(val).To(HaveLen(1)) Expect(val[0]).To(Equal("/home/user/.ssh/id_rsa.pub")) }) + + It("should accept vm-concurrency flag", func() { + rootCmd.SetArgs([]string{"run", "--vm-concurrency", "5"}) + Expect(rootCmd.Execute()).To(Succeed()) + + runCmd, _, _ := rootCmd.Find([]string{"run"}) + val, err := runCmd.Flags().GetInt("vm-concurrency") + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal(5)) + }) + + It("should default vm-concurrency to 10", func() { + rootCmd.SetArgs([]string{"run"}) + Expect(rootCmd.Execute()).To(Succeed()) + + runCmd, _, _ := rootCmd.Find([]string{"run"}) + val, err := runCmd.Flags().GetInt("vm-concurrency") + Expect(err).NotTo(HaveOccurred()) + Expect(val).To(Equal(10)) + }) }) var _ = Describe("Cleanup command flags", func() { diff --git a/internal/config/config.go b/internal/config/config.go index b464900..2b6fdca 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -77,6 +77,7 @@ type Config struct { SSHAuthorizedKeys []string `mapstructure:"ssh-authorized-keys"` AuditEnabled bool `mapstructure:"audit"` AuditDBPath string `mapstructure:"audit-db"` + VMConcurrency int `mapstructure:"vm-concurrency"` } // SetDefaults registers Viper defaults. @@ -96,6 +97,7 @@ func SetDefaults(v *viper.Viper) { v.SetDefault("cleanup-mode", "") v.SetDefault("audit", true) v.SetDefault("audit-db", constants.DefaultAuditDBPath) + v.SetDefault("vm-concurrency", constants.DefaultVMConcurrency) } // BindPersistentFlags registers persistent flags shared across all subcommands. @@ -127,6 +129,7 @@ func BindRunFlags(cmd *cobra.Command, defaultWorkloads []string) { f.String("ssh-password", "", "SSH password for VMs") f.StringSlice("ssh-key", nil, "SSH authorized key (repeatable)") f.StringSlice("ssh-key-file", nil, "SSH key file path (repeatable)") + f.Int("vm-concurrency", constants.DefaultVMConcurrency, "Max concurrent VM creation operations") } // BindCleanupFlags registers flags specific to the "cleanup" subcommand. @@ -189,6 +192,10 @@ func LoadConfig(cmd *cobra.Command) (*Config, error) { val, _ := cmd.Flags().GetBool("no-wait") v.Set("wait-for-ready", !val) } + if cmd.Flags().Changed("vm-concurrency") { + val, _ := cmd.Flags().GetInt("vm-concurrency") + v.Set("vm-concurrency", val) + } // Build the Config struct cfg := &Config{} @@ -207,6 +214,7 @@ func LoadConfig(cmd *cobra.Command) (*Config, error) { cfg.SSHPassword = v.GetString("ssh-password") cfg.AuditEnabled = v.GetBool("audit") cfg.AuditDBPath = v.GetString("audit-db") + cfg.VMConcurrency = v.GetInt("vm-concurrency") // Handle SSH authorized keys: CLI flags, env var (comma-split), or YAML list sshKeys, err := resolveSSHKeys(v, cmd) @@ -253,7 +261,9 @@ func resolveSSHKeys(v *viper.Viper, cmd *cobra.Command) ([]string, error) { if cmd.Flags().Changed("ssh-key-file") { paths, _ := cmd.Flags().GetStringSlice("ssh-key-file") for _, p := range paths { - data, err := os.ReadFile(filepath.Clean(p)) //nolint:gosec // CLI user supplies the path intentionally + data, err := os.ReadFile( + filepath.Clean(p), + ) //nolint:gosec // CLI user supplies the path intentionally if err != nil { return nil, fmt.Errorf("reading SSH key file %s: %w", p, err) } diff --git a/internal/constants/constants.go b/internal/constants/constants.go index 7a21da6..41d79ee 100644 --- a/internal/constants/constants.go +++ b/internal/constants/constants.go @@ -70,3 +70,6 @@ const ( DefaultReadyTimeout = 600 * time.Second DefaultPollInterval = 15 * time.Second ) + +// Concurrency defaults for VM creation. +const DefaultVMConcurrency = 10 diff --git a/internal/orchestrator/orchestrator.go b/internal/orchestrator/orchestrator.go index 345bead..f5bec01 100644 --- a/internal/orchestrator/orchestrator.go +++ b/internal/orchestrator/orchestrator.go @@ -460,6 +460,9 @@ func (ro *RunOrchestrator) createVMs( auditWorkloadIDs map[string]int64, ) error { g, gctx := errgroup.WithContext(ctx) + if cfg.VMConcurrency > 0 { + g.SetLimit(cfg.VMConcurrency) + } for _, plan := range plans { p := plan g.Go(func() error { diff --git a/internal/workloads/registry.go b/internal/workloads/registry.go index d1924d0..1635056 100644 --- a/internal/workloads/registry.go +++ b/internal/workloads/registry.go @@ -63,13 +63,29 @@ func AllWorkloadNames() []string { func DefaultRegistry() Registry { return Registry{ "chaos-process": func(cfg config.WorkloadConfig, opts *RegistryOpts) Workload { - return NewChaosProcessWorkload(cfg, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) + return NewChaosProcessWorkload( + cfg, + opts.SSHUser, + opts.SSHPassword, + opts.SSHAuthorizedKeys, + ) }, "chaos-network": func(cfg config.WorkloadConfig, opts *RegistryOpts) Workload { - return NewChaosNetworkWorkload(cfg, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) + return NewChaosNetworkWorkload( + cfg, + opts.SSHUser, + opts.SSHPassword, + opts.SSHAuthorizedKeys, + ) }, "chaos-disk": func(cfg config.WorkloadConfig, opts *RegistryOpts) Workload { - return NewChaosDiskWorkload(cfg, opts.DataDiskSize, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) + return NewChaosDiskWorkload( + cfg, + opts.DataDiskSize, + opts.SSHUser, + opts.SSHPassword, + opts.SSHAuthorizedKeys, + ) }, "cpu": func(cfg config.WorkloadConfig, opts *RegistryOpts) Workload { return NewCPUWorkload(cfg, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) @@ -78,20 +94,118 @@ func DefaultRegistry() Registry { return NewMemoryWorkload(cfg, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) }, "disk": func(cfg config.WorkloadConfig, opts *RegistryOpts) Workload { - return NewDiskWorkload(cfg, opts.DataDiskSize, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) + return NewDiskWorkload( + cfg, + opts.DataDiskSize, + opts.SSHUser, + opts.SSHPassword, + opts.SSHAuthorizedKeys, + ) }, "database": func(cfg config.WorkloadConfig, opts *RegistryOpts) Workload { - return NewDatabaseWorkload(cfg, opts.DataDiskSize, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) + return NewDatabaseWorkload( + cfg, + opts.DataDiskSize, + opts.SSHUser, + opts.SSHPassword, + opts.SSHAuthorizedKeys, + ) }, "network": func(cfg config.WorkloadConfig, opts *RegistryOpts) Workload { - return NewNetworkWorkload(cfg, opts.Namespace, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) + return NewNetworkWorkload( + cfg, + opts.Namespace, + opts.SSHUser, + opts.SSHPassword, + opts.SSHAuthorizedKeys, + ) }, "tps": func(cfg config.WorkloadConfig, opts *RegistryOpts) Workload { - return NewTPSWorkload(cfg, opts.Namespace, opts.SSHUser, opts.SSHPassword, opts.SSHAuthorizedKeys) + return NewTPSWorkload( + cfg, + opts.Namespace, + opts.SSHUser, + opts.SSHPassword, + opts.SSHAuthorizedKeys, + ) }, } } +// ValidateWorkloadNames checks that all names are valid workload names from the +// default registry. Returns an error listing invalid names with "did you mean?" +// suggestions for close matches. +func ValidateWorkloadNames(names []string) error { + if len(names) == 0 { + return fmt.Errorf("no workloads specified; available: %s; %w", + strings.Join(AllWorkloadNames(), ", "), ErrWorkloadUnknown) + } + + valid := DefaultRegistry() + var invalid []string + for _, name := range names { + if _, ok := valid[name]; !ok { + suggestion := closestMatch(name, valid.List()) + if suggestion != "" { + invalid = append(invalid, fmt.Sprintf("%q (did you mean %q?)", name, suggestion)) + } else { + invalid = append(invalid, fmt.Sprintf("%q", name)) + } + } + } + if len(invalid) > 0 { + return fmt.Errorf( + "unknown workload(s): %s; available: %s; %w", + strings.Join(invalid, ", "), + strings.Join(valid.List(), ", "), + ErrWorkloadUnknown, + ) + } + return nil +} + +func closestMatch(input string, candidates []string) string { + best := "" + bestDist := len(input)/2 + 1 + for _, c := range candidates { + d := levenshtein(input, c) + if d < bestDist { + bestDist = d + best = c + } + } + return best +} + +func levenshtein(a, b string) int { + la, lb := len(a), len(b) + if la == 0 { + return lb + } + if lb == 0 { + return la + } + + prev := make([]int, lb+1) + curr := make([]int, lb+1) + for j := range prev { + prev[j] = j + } + + for i := 1; i <= la; i++ { + curr[0] = i + for j := 1; j <= lb; j++ { + cost := 1 + if a[i-1] == b[j-1] { + cost = 0 + } + curr[j] = min(curr[j-1]+1, min(prev[j]+1, prev[j-1]+cost)) + } + prev, curr = curr, prev + } + return prev[lb] +} + // Get retrieves a workload by name, constructing it with the given config and options. // Returns an error listing available names if the workload is not found. func (r Registry) Get(name string, cfg config.WorkloadConfig, opts ...Option) (Workload, error) { diff --git a/internal/workloads/registry_test.go b/internal/workloads/registry_test.go index 80c930d..d6bfd63 100644 --- a/internal/workloads/registry_test.go +++ b/internal/workloads/registry_test.go @@ -4,6 +4,8 @@ package workloads_test import ( + "errors" + . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -221,3 +223,61 @@ var _ = Describe("AllWorkloadNames", func() { Expect(workloads.AllWorkloadNames()).To(Equal(workloads.DefaultRegistry().List())) }) }) + +var _ = Describe("ValidateWorkloadNames", func() { + It("should accept all valid workload names", func() { + err := workloads.ValidateWorkloadNames(workloads.AllWorkloadNames()) + Expect(err).NotTo(HaveOccurred()) + }) + + It("should accept a subset of valid names", func() { + err := workloads.ValidateWorkloadNames([]string{"cpu", "memory", "disk"}) + Expect(err).NotTo(HaveOccurred()) + }) + + It("should return error for a single invalid name", func() { + err := workloads.ValidateWorkloadNames([]string{"cppu"}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("cppu")) + Expect(err.Error()).To(ContainSubstring("available:")) + }) + + It("should return error listing all invalid names", func() { + err := workloads.ValidateWorkloadNames([]string{"cpu", "memorry", "diks"}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("memorry")) + Expect(err.Error()).To(ContainSubstring("diks")) + Expect(err.Error()).NotTo(ContainSubstring(`"cpu"`)) + }) + + It("should suggest similar names for typos", func() { + err := workloads.ValidateWorkloadNames([]string{"cppu"}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("did you mean")) + Expect(err.Error()).To(ContainSubstring("cpu")) + }) + + It("should suggest similar names for close misspellings", func() { + err := workloads.ValidateWorkloadNames([]string{"memori"}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("did you mean")) + Expect(err.Error()).To(ContainSubstring("memory")) + }) + + It("should not suggest when nothing is close", func() { + err := workloads.ValidateWorkloadNames([]string{"zzzzzzz"}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).NotTo(ContainSubstring("did you mean")) + }) + + It("should return error for empty list", func() { + err := workloads.ValidateWorkloadNames([]string{}) + Expect(err).To(HaveOccurred()) + }) + + It("should wrap ErrWorkloadUnknown", func() { + err := workloads.ValidateWorkloadNames([]string{"bogus"}) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, workloads.ErrWorkloadUnknown)).To(BeTrue()) + }) +})