refactoring loadtester to support external testing platform

cmd/loadtester/main.go

@@ -23,8 +23,7 @@ var (
 func init() {
 	flag.StringVar(&logLevel, "log-level", "debug", "Log level can be: debug, info, warning, error.")
 	flag.StringVar(&port, "port", "9090", "Port to listen on.")
-	flag.DurationVar(&timeout, "timeout", time.Hour, "Command exec timeout.")
-	flag.BoolVar(&logCmdOutput, "log-cmd-output", true, "Log command output to stderr")
+	flag.DurationVar(&timeout, "timeout", time.Hour, "Load test exec timeout.")
 	flag.BoolVar(&zapReplaceGlobals, "zap-replace-globals", false, "Whether to change the logging level of the global zap logger.")
 	flag.StringVar(&zapEncoding, "zap-encoding", "json", "Zap logger encoding.")
 }
@@ -44,7 +43,7 @@ func main() {
 
 	stopCh := signals.SetupSignalHandler()
 
-	taskRunner := loadtester.NewTaskRunner(logger, timeout, logCmdOutput)
+	taskRunner := loadtester.NewTaskRunner(logger, timeout)
 
 	go taskRunner.Start(100*time.Millisecond, stopCh)
 

pkg/loadtester/runner.go

@@ -2,11 +2,7 @@ package loadtester
 
 import (
 	"context"
-	"encoding/hex"
-	"fmt"
 	"go.uber.org/zap"
-	"hash/fnv"
-	"os/exec"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -21,24 +17,12 @@ type TaskRunner struct {
 	logCmdOutput bool
 }
 
-type Task struct {
-	Canary  string
-	Command string
-}
-
-func (t Task) Hash() string {
-	fnvHash := fnv.New32()
-	fnvBytes := fnvHash.Sum([]byte(t.Canary + t.Command))
-	return hex.EncodeToString(fnvBytes[:])
-}
-
-func NewTaskRunner(logger *zap.SugaredLogger, timeout time.Duration, logCmdOutput bool) *TaskRunner {
+func NewTaskRunner(logger *zap.SugaredLogger, timeout time.Duration) *TaskRunner {
 	return &TaskRunner{
 		logger:       logger,
 		todoTasks:    new(sync.Map),
 		runningTasks: new(sync.Map),
 		timeout:      timeout,
-		logCmdOutput: logCmdOutput,
 	}
 }
 
@@ -69,24 +53,15 @@ func (tr *TaskRunner) runAll() {
 				// increment the total exec counter
 				atomic.AddUint64(&tr.totalExecs, 1)
 
-				tr.logger.With("canary", t.Canary).Infof("command starting %s", t.Command)
-				cmd := exec.CommandContext(ctx, "sh", "-c", t.Command)
+				tr.logger.With("canary", t.Canary()).Infof("task starting %s", t)
 
-				// execute task
-				out, err := cmd.CombinedOutput()
-				if err != nil {
-					tr.logger.With("canary", t.Canary).Errorf("command failed %s %v %s", t.Command, err, out)
-				} else {
-					if tr.logCmdOutput {
-						fmt.Printf("%s\n", out)
-					}
-					tr.logger.With("canary", t.Canary).Infof("command finished %s", t.Command)
-				}
+				// run task with the timeout context
+				t.Run(ctx)
 
 				// remove task from the running list
 				tr.runningTasks.Delete(t.Hash())
 			} else {
-				tr.logger.With("canary", t.Canary).Infof("command skipped %s is already running", t.Command)
+				tr.logger.With("canary", t.Canary()).Infof("command skipped %s is already running", t)
 			}
 		}(task)
 		return true

pkg/loadtester/runner_test.go

@@ -9,18 +9,13 @@ import (
 func TestTaskRunner_Start(t *testing.T) {
 	stop := make(chan struct{})
 	logger, _ := logging.NewLogger("debug")
-	tr := NewTaskRunner(logger, time.Hour, false)
+	tr := NewTaskRunner(logger, time.Hour)
 
 	go tr.Start(10*time.Millisecond, stop)
 
-	task1 := Task{
-		Canary:  "podinfo.default",
-		Command: "sleep 0.6",
-	}
-	task2 := Task{
-		Canary:  "podinfo.default",
-		Command: "sleep 0.7",
-	}
+	taskFactory, _ := GetTaskFactory(TaskTypeShell)
+	task1, _ := taskFactory(map[string]string{"cmd": "sleep 0.6"}, "podinfo.default", logger)
+	task2, _ := taskFactory(map[string]string{"cmd": "sleep 0.7"}, "podinfo.default", logger)
 
 	tr.Add(task1)
 	tr.Add(task2)

pkg/loadtester/server.go

@@ -39,16 +39,25 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
 		}
 
 		if len(payload.Metadata) > 0 {
-			if cmd, ok := payload.Metadata["cmd"]; ok {
-				taskRunner.Add(Task{
-					Canary:  fmt.Sprintf("%s.%s", payload.Name, payload.Namespace),
-					Command: cmd,
-				})
-			} else {
+			metadata := payload.Metadata
+			var typ, ok = metadata["type"]
+			if !ok {
+				typ = TaskTypeShell
+			}
+			taskFactory, ok := GetTaskFactory(typ)
+			if !ok {
 				w.WriteHeader(http.StatusBadRequest)
-				w.Write([]byte("cmd not found in metadata"))
+				w.Write([]byte(fmt.Sprintf("unknown task type %s", typ)))
 				return
 			}
+			canary := fmt.Sprintf("%s.%s", payload.Name, payload.Namespace)
+			task, err := taskFactory(metadata, canary, logger)
+			if err != nil {
+				w.WriteHeader(http.StatusBadRequest)
+				w.Write([]byte(err.Error()))
+				return
+			}
+			taskRunner.Add(task)
 		} else {
 			w.WriteHeader(http.StatusBadRequest)
 			w.Write([]byte("metadata not found in payload"))

pkg/loadtester/task.go

@@ -0,0 +1,40 @@
+package loadtester
+
+import (
+	"context"
+	"encoding/hex"
+	"go.uber.org/zap"
+	"hash/fnv"
+	"sync"
+)
+
+// Modeling a loadtester task
+type Task interface {
+	Hash() string
+	Run(ctx context.Context) bool
+	String() string
+	Canary() string
+}
+
+type TaskBase struct {
+	canary string
+	logger *zap.SugaredLogger
+}
+
+func (task *TaskBase) Canary() string {
+	return task.canary
+}
+func hash(str string) string {
+	fnvHash := fnv.New32()
+	fnvBytes := fnvHash.Sum([]byte(str))
+	return hex.EncodeToString(fnvBytes[:])
+}
+
+var taskFactories = new(sync.Map)
+
+type TaskFactory = func(metadata map[string]string, canary string, logger *zap.SugaredLogger) (Task, error)
+
+func GetTaskFactory(typ string) (TaskFactory, bool) {
+	factory, ok := taskFactories.Load(typ)
+	return factory.(TaskFactory), ok
+}