feat(workloads): replace Roles() with RoleDistribution() on MultiVMWorkload

internal/orchestrator/orchestrator.go

@@ -278,44 +278,48 @@ func (ro *RunOrchestrator) planVMs(
 				vmNames = append(vmNames, vmName)
 			}
 		} else {
-			roles := multiVM.Roles()
-			perRole := vmCount / len(roles)
-			if remainder := vmCount % len(roles); remainder != 0 {
-				ro.logger.Warn("VM count not evenly divisible by roles; remainder VMs will not be created",
-					slog.String("workload", name),
-					slog.Int("requested", vmCount),
-					slog.Int("per_role", perRole),
-					slog.Int("dropped", remainder))
-			}
-			for _, role := range roles {
-				userdata, err := multiVM.UserdataForRole(role, cfg.Namespace)
+			for _, rs := range multiVM.RoleDistribution() {
+				userdata, err := multiVM.UserdataForRole(rs.Role, cfg.Namespace)
 				if err != nil {
-					return nil, nil, nil, nil, fmt.Errorf("generating cloud-init for %q role %q: %w", name, role, err)
+					return nil, nil, nil, nil, fmt.Errorf(
+						"generating cloud-init for %q role %q: %w",
+						name,
+						rs.Role,
+						err,
+					)
 				}
 
-				for i := range perRole {
-					vmName := fmt.Sprintf("virtwork-%s-%s-%d", name, role, i)
+				for i := range rs.VMCount {
+					vmName := fmt.Sprintf("virtwork-%s-%s-%d", name, rs.Role, i)
 					labels := map[string]string{
 						constants.LabelAppName:   fmt.Sprintf("virtwork-%s", name),
 						constants.LabelManagedBy: constants.ManagedByValue,
 						constants.LabelComponent: name,
 						constants.LabelRunID:     runID,
-						"virtwork/role":          role,
+						"virtwork/role":          rs.Role,
 					}
 
 					dvts, err := w.DataVolumeTemplates()
 					if err != nil {
 						return nil, nil, nil, nil, fmt.Errorf(
-							"building data volume templates for %q role %q vm %s: %w", name, role, vmName, err,
+							"building data volume templates for %q role %q vm %s: %w",
+							name,
+							rs.Role,
+							vmName,
+							err,
 						)
 					}
-					dvTemplates, extraVols := NamespaceDataVolumes(dvts, w.ExtraVolumes(), vmName)
+					dvTemplates, extraVols := NamespaceDataVolumes(
+						dvts,
+						w.ExtraVolumes(),
+						vmName,
+					)
 
 					plans = append(plans, VMPlan{
 						WorkloadName: name,
 						Component:    name,
 						VMName:       vmName,
-						Role:         role,
+						Role:         rs.Role,
 						VMSpec: &VMSpecInput{
 							Name:                vmName,
 							Namespace:           cfg.Namespace,

internal/workloads/network.go

@@ -63,17 +63,23 @@ func (w *NetworkWorkload) Name() string {
 	return "network"
 }
 
-// VMCount returns the total VM count — one server and one client per
-// configured vm-count.
-func (w *NetworkWorkload) VMCount() int {
-	count := w.Config.VMCount
-	count = max(1, count)
-	return count * 2
+// RoleDistribution returns per-role VM counts — one server and one client
+// per configured vm-count.
+func (w *NetworkWorkload) RoleDistribution() []RoleSpec {
+	perRole := max(1, w.Config.VMCount)
+	return []RoleSpec{
+		{Role: "server", VMCount: perRole},
+		{Role: "client", VMCount: perRole},
+	}
 }
 
-// Roles returns the supported VM roles for this multi-VM workload.
-func (w *NetworkWorkload) Roles() []string {
-	return []string{"server", "client"}
+// VMCount returns the total VM count across all roles.
+func (w *NetworkWorkload) VMCount() int {
+	total := 0
+	for _, rs := range w.RoleDistribution() {
+		total += rs.VMCount
+	}
+	return total
 }
 
 // RequiresService returns true — the client needs a ClusterIP Service to reach

internal/workloads/network_test.go

@@ -35,9 +35,20 @@ var _ = Describe("NetworkWorkload", func() {
 		Expect(w.RequiresService()).To(BeTrue())
 	})
 
-	It("should return server and client roles", func() {
-		roles := w.Roles()
-		Expect(roles).To(Equal([]string{"server", "client"}))
+	It("should return per-role VM counts via RoleDistribution", func() {
+		dist := w.RoleDistribution()
+		Expect(dist).To(HaveLen(2))
+		Expect(dist[0]).To(Equal(workloads.RoleSpec{Role: "server", VMCount: 2}))
+		Expect(dist[1]).To(Equal(workloads.RoleSpec{Role: "client", VMCount: 2}))
+	})
+
+	It("should have VMCount equal to sum of RoleDistribution counts", func() {
+		dist := w.RoleDistribution()
+		total := 0
+		for _, rs := range dist {
+			total += rs.VMCount
+		}
+		Expect(w.VMCount()).To(Equal(total))
 	})
 
 	It("should produce server userdata with iperf3 -s", func() {

internal/workloads/tps.go

@@ -82,17 +82,23 @@ func (w *TPSWorkload) Name() string {
 	return "tps"
 }
 
-// VMCount returns the total VM count — one server and one client per
-// configured vm-count.
-func (w *TPSWorkload) VMCount() int {
-	count := w.Config.VMCount
-	count = max(1, count)
-	return count * 2
+// RoleDistribution returns per-role VM counts — one server and one client
+// per configured vm-count.
+func (w *TPSWorkload) RoleDistribution() []RoleSpec {
+	perRole := max(1, w.Config.VMCount)
+	return []RoleSpec{
+		{Role: "server", VMCount: perRole},
+		{Role: "client", VMCount: perRole},
+	}
 }
 
-// Roles returns the supported VM roles for this multi-VM workload.
-func (w *TPSWorkload) Roles() []string {
-	return []string{"server", "client"}
+// VMCount returns the total VM count across all roles.
+func (w *TPSWorkload) VMCount() int {
+	total := 0
+	for _, rs := range w.RoleDistribution() {
+		total += rs.VMCount
+	}
+	return total
 }
 
 // RequiresService returns true — the client needs a ClusterIP Service to reach

internal/workloads/tps_test.go

@@ -37,6 +37,22 @@ var _ = Describe("TPSWorkload", func() {
 		Expect(w.RequiresService()).To(BeTrue())
 	})
 
+	It("should return per-role VM counts via RoleDistribution", func() {
+		dist := w.RoleDistribution()
+		Expect(dist).To(HaveLen(2))
+		Expect(dist[0]).To(Equal(workloads.RoleSpec{Role: "server", VMCount: 2}))
+		Expect(dist[1]).To(Equal(workloads.RoleSpec{Role: "client", VMCount: 2}))
+	})
+
+	It("should have VMCount equal to sum of RoleDistribution counts", func() {
+		dist := w.RoleDistribution()
+		total := 0
+		for _, rs := range dist {
+			total += rs.VMCount
+		}
+		Expect(w.VMCount()).To(Equal(total))
+	})
+
 	It("should produce server userdata with netserver and HTTP server", func() {
 		result, err := w.UserdataForRole("server", "virtwork")
 		Expect(err).NotTo(HaveOccurred())

internal/workloads/workload.go

@@ -53,12 +53,18 @@ type Workload interface {
 	VMCount() int
 }
 
+// RoleSpec declares the number of VMs for a single role within a multi-VM workload.
+type RoleSpec struct {
+	Role    string
+	VMCount int
+}
+
 // MultiVMWorkload extends Workload for workloads that need per-role userdata.
-// The orchestration layer checks VMCount() > 1 and type-asserts to this
-// interface to call UserdataForRole() for each VM role.
+// The orchestration layer type-asserts to this interface and iterates
+// RoleDistribution() to create the correct number of VMs per role.
 type MultiVMWorkload interface {
 	Workload
-	Roles() []string
+	RoleDistribution() []RoleSpec
 	UserdataForRole(role string, namespace string) (string, error)
 }