feat: validate workload names upfront and add VM creation concurrency limit

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),

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() {

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

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

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 {

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) {

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