From 9f12bbcd98935df4b15a270073f1c4c9c5fde591 Mon Sep 17 00:00:00 2001 From: Alex Wong Date: Fri, 8 Mar 2019 15:49:35 +0800 Subject: [PATCH] refactoring loadtester to support external testing platform --- cmd/loadtester/main.go | 5 ++--- pkg/loadtester/runner.go | 35 +++++------------------------- pkg/loadtester/runner_test.go | 13 ++++-------- pkg/loadtester/server.go | 23 ++++++++++++++------ pkg/loadtester/task.go | 40 +++++++++++++++++++++++++++++++++++ 5 files changed, 67 insertions(+), 49 deletions(-) create mode 100644 pkg/loadtester/task.go diff --git a/cmd/loadtester/main.go b/cmd/loadtester/main.go index 35800df5..df463a92 100644 --- a/cmd/loadtester/main.go +++ b/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) diff --git a/pkg/loadtester/runner.go b/pkg/loadtester/runner.go index d70b9b4f..2cba9050 100644 --- a/pkg/loadtester/runner.go +++ b/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 diff --git a/pkg/loadtester/runner_test.go b/pkg/loadtester/runner_test.go index 5bd7b4ae..cb963462 100644 --- a/pkg/loadtester/runner_test.go +++ b/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) diff --git a/pkg/loadtester/server.go b/pkg/loadtester/server.go index 58a50ed4..85d086dc 100644 --- a/pkg/loadtester/server.go +++ b/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")) diff --git a/pkg/loadtester/task.go b/pkg/loadtester/task.go new file mode 100644 index 00000000..2e30228a --- /dev/null +++ b/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 +}