diff --git a/internal/orchestrator/orchestrator.go b/internal/orchestrator/orchestrator.go index f5bec01..d0509eb 100644 --- a/internal/orchestrator/orchestrator.go +++ b/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, diff --git a/internal/workloads/network.go b/internal/workloads/network.go index dc55fba..0b2fd2e 100644 --- a/internal/workloads/network.go +++ b/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 diff --git a/internal/workloads/network_test.go b/internal/workloads/network_test.go index 9eb2fc6..1ea8b87 100644 --- a/internal/workloads/network_test.go +++ b/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() { diff --git a/internal/workloads/tps.go b/internal/workloads/tps.go index de552a8..01a81c5 100644 --- a/internal/workloads/tps.go +++ b/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 diff --git a/internal/workloads/tps_test.go b/internal/workloads/tps_test.go index 98e9826..d97cf8c 100644 --- a/internal/workloads/tps_test.go +++ b/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()) diff --git a/internal/workloads/workload.go b/internal/workloads/workload.go index bbbe043..4918297 100644 --- a/internal/workloads/workload.go +++ b/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) }