feat: configurable params for built-in workloads

internal/workloads/cpu.go

@@ -4,16 +4,18 @@
 package workloads
 
 import (
+	"fmt"
+
 	"github.com/opdev/virtwork/internal/config"
 )
 
-const cpuSystemdUnit = `[Unit]
+const cpuSystemdUnitTemplate = `[Unit]
 Description=Virtwork CPU stress workload
 After=network.target
 
 [Service]
 Type=simple
-ExecStart=/usr/bin/stress-ng --cpu 0 --cpu-method all --timeout 0
+ExecStart=/usr/bin/stress-ng --cpu 0 --cpu-load %s --cpu-method %s --timeout 0
 Restart=always
 RestartSec=10
 
@@ -28,7 +30,11 @@ type CPUWorkload struct {
 }
 
 // NewCPUWorkload creates a CPUWorkload with the given configuration and SSH credentials.
-func NewCPUWorkload(cfg config.WorkloadConfig, sshUser, sshPassword string, sshKeys []string) *CPUWorkload {
+func NewCPUWorkload(
+	cfg config.WorkloadConfig,
+	sshUser, sshPassword string,
+	sshKeys []string,
+) *CPUWorkload {
 	return &CPUWorkload{
 		BaseWorkload: BaseWorkload{
 			Config:            cfg,
@@ -39,6 +45,24 @@ func NewCPUWorkload(cfg config.WorkloadConfig, sshUser, sshPassword string, sshK
 	}
 }
 
+func (w *CPUWorkload) cpuLoadPercent() string {
+	if w.Config.Params != nil {
+		if val, ok := w.Config.Params["cpu-load-percent"]; ok && val != "" {
+			return val
+		}
+	}
+	return "100"
+}
+
+func (w *CPUWorkload) cpuMethod() string {
+	if w.Config.Params != nil {
+		if val, ok := w.Config.Params["cpu-method"]; ok && val != "" {
+			return val
+		}
+	}
+	return "all"
+}
+
 // Name returns "cpu".
 func (w *CPUWorkload) Name() string {
 	return "cpu"
@@ -47,12 +71,13 @@ func (w *CPUWorkload) Name() string {
 // CloudInitUserdata returns cloud-init YAML that installs stress-ng and runs a
 // continuous CPU stress workload via systemd.
 func (w *CPUWorkload) CloudInitUserdata() (string, error) {
+	unit := fmt.Sprintf(cpuSystemdUnitTemplate, w.cpuLoadPercent(), w.cpuMethod())
 	return w.BuildCloudConfig(CloudConfigOpts{
 		Packages: []string{"stress-ng"},
 		WriteFiles: []WriteFile{
 			{
 				Path:        "/etc/systemd/system/virtwork-cpu.service",
-				Content:     cpuSystemdUnit,
+				Content:     unit,
 				Permissions: "0644",
 			},
 		},

internal/workloads/cpu_test.go

@@ -52,6 +52,7 @@ var _ = Describe("CPUWorkload", func() {
 		content := file["content"].(string)
 		Expect(content).To(ContainSubstring("stress-ng"))
 		Expect(content).To(ContainSubstring("--cpu 0"))
+		Expect(content).To(ContainSubstring("--cpu-load 100"))
 		Expect(content).To(ContainSubstring("--cpu-method all"))
 		Expect(content).To(ContainSubstring("--timeout 0"))
 	})
@@ -79,4 +80,66 @@ var _ = Describe("CPUWorkload", func() {
 		Expect(res.CPUCores).To(Equal(2))
 		Expect(res.Memory).To(Equal("2Gi"))
 	})
+
+	Context("param wiring", func() {
+		It("should use default param values when Params is nil", func() {
+			result, err := w.CloudInitUserdata()
+			Expect(err).NotTo(HaveOccurred())
+
+			parsed := parseYAML(result)
+			files := parsed["write_files"].([]interface{})
+			file := files[0].(map[string]interface{})
+			content := file["content"].(string)
+
+			Expect(content).To(ContainSubstring("--cpu-load 100"))
+			Expect(content).To(ContainSubstring("--cpu-method all"))
+		})
+
+		It("should wire custom params from WorkloadConfig.Params", func() {
+			custom := workloads.NewCPUWorkload(config.WorkloadConfig{
+				Enabled:  new(true),
+				VMCount:  1,
+				CPUCores: 2,
+				Memory:   "2Gi",
+				Params: map[string]string{
+					"cpu-load-percent": "50",
+					"cpu-method":       "matrixprod",
+				},
+			}, "virtwork", "", nil)
+
+			result, err := custom.CloudInitUserdata()
+			Expect(err).NotTo(HaveOccurred())
+
+			parsed := parseYAML(result)
+			files := parsed["write_files"].([]interface{})
+			file := files[0].(map[string]interface{})
+			content := file["content"].(string)
+
+			Expect(content).To(ContainSubstring("--cpu-load 50"))
+			Expect(content).To(ContainSubstring("--cpu-method matrixprod"))
+		})
+
+		It("should use defaults for missing individual params", func() {
+			partial := workloads.NewCPUWorkload(config.WorkloadConfig{
+				Enabled:  new(true),
+				VMCount:  1,
+				CPUCores: 2,
+				Memory:   "2Gi",
+				Params: map[string]string{
+					"cpu-load-percent": "75",
+				},
+			}, "virtwork", "", nil)
+
+			result, err := partial.CloudInitUserdata()
+			Expect(err).NotTo(HaveOccurred())
+
+			parsed := parseYAML(result)
+			files := parsed["write_files"].([]interface{})
+			file := files[0].(map[string]interface{})
+			content := file["content"].(string)
+
+			Expect(content).To(ContainSubstring("--cpu-load 75"))
+			Expect(content).To(ContainSubstring("--cpu-method all"))
+		})
+	})
 })

internal/workloads/database.go

@@ -4,14 +4,16 @@
 package workloads
 
 import (
+	"fmt"
+
 	kubevirtv1 "kubevirt.io/api/core/v1"
 
 	"github.com/opdev/virtwork/internal/config"
 	"github.com/opdev/virtwork/internal/constants"
 	"github.com/opdev/virtwork/internal/vm"
 )
 
-const dbSetupScript = `#!/bin/bash
+const dbSetupScriptTemplate = `#!/bin/bash
 set -euo pipefail
 
 DATA_DIR="/var/lib/pgsql/data"
@@ -55,9 +57,9 @@ postgresql-setup --initdb
 # Start PostgreSQL temporarily for pgbench init
 systemctl start postgresql
 
-# Create pgbench database with scale factor 50
+# Create pgbench database with configured scale factor
 sudo -u postgres createdb pgbench
-sudo -u postgres pgbench -i -s 50 pgbench
+sudo -u postgres pgbench -i -s %s pgbench
 
 # Stop PostgreSQL (systemd will manage it)
 systemctl stop postgresql
@@ -67,7 +69,7 @@ touch "${MARKER}"
 chown postgres:postgres "${MARKER}"
 `
 
-const dbSystemdUnit = `[Unit]
+const dbSystemdUnitTemplate = `[Unit]
 Description=Virtwork database benchmark workload
 After=network.target local-fs.target postgresql.service
 Requires=postgresql.service
@@ -76,7 +78,7 @@ Requires=postgresql.service
 Type=simple
 User=postgres
 ExecStartPre=/usr/local/bin/virtwork-db-setup.sh
-ExecStart=/bin/bash -c 'while true; do pgbench -c 10 -j 2 -T 300 pgbench; sleep 10; done'
+ExecStart=/bin/bash -c 'while true; do pgbench -c %s -j 2 -T %s pgbench; sleep 10; done'
 Restart=always
 RestartSec=10
 
@@ -111,6 +113,33 @@ func NewDatabaseWorkload(
 	}
 }
 
+func (w *DatabaseWorkload) scaleFactor() string {
+	if w.Config.Params != nil {
+		if val, ok := w.Config.Params["scale-factor"]; ok && val != "" {
+			return val
+		}
+	}
+	return "50"
+}
+
+func (w *DatabaseWorkload) clients() string {
+	if w.Config.Params != nil {
+		if val, ok := w.Config.Params["clients"]; ok && val != "" {
+			return val
+		}
+	}
+	return "10"
+}
+
+func (w *DatabaseWorkload) duration() string {
+	if w.Config.Params != nil {
+		if val, ok := w.Config.Params["duration"]; ok && val != "" {
+			return val
+		}
+	}
+	return "300"
+}
+
 // Name returns "database".
 func (w *DatabaseWorkload) Name() string {
 	return "database"
@@ -120,17 +149,19 @@ func (w *DatabaseWorkload) Name() string {
 // a setup script for one-time database initialization, and creates a systemd
 // service that runs continuous pgbench benchmarks.
 func (w *DatabaseWorkload) CloudInitUserdata() (string, error) {
+	setupScript := fmt.Sprintf(dbSetupScriptTemplate, w.scaleFactor())
+	serviceUnit := fmt.Sprintf(dbSystemdUnitTemplate, w.clients(), w.duration())
 	return w.BuildCloudConfig(CloudConfigOpts{
 		Packages: []string{"postgresql-server"},
 		WriteFiles: []WriteFile{
 			{
 				Path:        "/usr/local/bin/virtwork-db-setup.sh",
-				Content:     dbSetupScript,
+				Content:     setupScript,
 				Permissions: "0755",
 			},
 			{
 				Path:        "/etc/systemd/system/virtwork-database.service",
-				Content:     dbSystemdUnit,
+				Content:     serviceUnit,
 				Permissions: "0644",
 			},
 		},

internal/workloads/database_test.go

@@ -165,4 +165,91 @@ var _ = Describe("DatabaseWorkload", func() {
 		Expect(res.CPUCores).To(Equal(2))
 		Expect(res.Memory).To(Equal("4Gi"))
 	})
+
+	Context("param wiring", func() {
+		It("should use default param values when Params is nil", func() {
+			result, err := w.CloudInitUserdata()
+			Expect(err).NotTo(HaveOccurred())
+
+			parsed := parseYAML(result)
+			files := parsed["write_files"].([]interface{})
+
+			filesByPath := map[string]string{}
+			for _, f := range files {
+				fm := f.(map[string]interface{})
+				filesByPath[fm["path"].(string)] = fm["content"].(string)
+			}
+
+			setupScript := filesByPath["/usr/local/bin/virtwork-db-setup.sh"]
+			Expect(setupScript).To(ContainSubstring("pgbench -i -s 50"))
+
+			serviceUnit := filesByPath["/etc/systemd/system/virtwork-database.service"]
+			Expect(serviceUnit).To(ContainSubstring("-c 10"))
+			Expect(serviceUnit).To(ContainSubstring("-T 300"))
+		})
+
+		It("should wire custom params from WorkloadConfig.Params", func() {
+			custom := workloads.NewDatabaseWorkload(config.WorkloadConfig{
+				Enabled:  new(true),
+				VMCount:  1,
+				CPUCores: 2,
+				Memory:   "4Gi",
+				Params: map[string]string{
+					"scale-factor": "100",
+					"clients":      "20",
+					"duration":     "600",
+				},
+			}, "10Gi", "virtwork", "", nil)
+
+			result, err := custom.CloudInitUserdata()
+			Expect(err).NotTo(HaveOccurred())
+
+			parsed := parseYAML(result)
+			files := parsed["write_files"].([]interface{})
+
+			filesByPath := map[string]string{}
+			for _, f := range files {
+				fm := f.(map[string]interface{})
+				filesByPath[fm["path"].(string)] = fm["content"].(string)
+			}
+
+			setupScript := filesByPath["/usr/local/bin/virtwork-db-setup.sh"]
+			Expect(setupScript).To(ContainSubstring("pgbench -i -s 100"))
+
+			serviceUnit := filesByPath["/etc/systemd/system/virtwork-database.service"]
+			Expect(serviceUnit).To(ContainSubstring("-c 20"))
+			Expect(serviceUnit).To(ContainSubstring("-T 600"))
+		})
+
+		It("should use defaults for missing individual params", func() {
+			partial := workloads.NewDatabaseWorkload(config.WorkloadConfig{
+				Enabled:  new(true),
+				VMCount:  1,
+				CPUCores: 2,
+				Memory:   "4Gi",
+				Params: map[string]string{
+					"clients": "16",
+				},
+			}, "10Gi", "virtwork", "", nil)
+
+			result, err := partial.CloudInitUserdata()
+			Expect(err).NotTo(HaveOccurred())
+
+			parsed := parseYAML(result)
+			files := parsed["write_files"].([]interface{})
+
+			filesByPath := map[string]string{}
+			for _, f := range files {
+				fm := f.(map[string]interface{})
+				filesByPath[fm["path"].(string)] = fm["content"].(string)
+			}
+
+			setupScript := filesByPath["/usr/local/bin/virtwork-db-setup.sh"]
+			Expect(setupScript).To(ContainSubstring("pgbench -i -s 50"))
+
+			serviceUnit := filesByPath["/etc/systemd/system/virtwork-database.service"]
+			Expect(serviceUnit).To(ContainSubstring("-c 16"))
+			Expect(serviceUnit).To(ContainSubstring("-T 300"))
+		})
+	})
 })