loadtester: add return cmd output option

go.sum

@@ -49,9 +49,11 @@ github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMo
 github.com/ghodss/yaml v0.0.0-20150909031657-73d445a93680/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
 github.com/go-kit/kit v0.8.0 h1:Wz+5lgoB0kkuqLEc6NVmwRknTKP6dTGbSqvhZtBI/j0=
 github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
+github.com/go-kit/kit v0.9.0 h1:wDJmvq38kDhkVxi50ni9ykkdUr1PKgqKOoi01fa0Mdk=
 github.com/go-kit/kit v0.9.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
 github.com/go-logfmt/logfmt v0.3.0 h1:8HUsc87TaSWLKwrnumgC8/YconD2fJQsRJAsWaPg2ic=
 github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE=
+github.com/go-logfmt/logfmt v0.4.0 h1:MP4Eh7ZCb31lleYCFuwm0oe4/YGak+5l1vA2NOE80nA=
 github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V4qmtdjCk=
 github.com/go-openapi/jsonpointer v0.0.0-20160704185906-46af16f9f7b1/go.mod h1:+35s3my2LFTysnkMfxsJBAMHj/DoqoB9knIWoYG/Vk0=
 github.com/go-openapi/jsonpointer v0.19.2/go.mod h1:3akKfEdA7DF1sugOqz1dVQHBcuDBPKZGEoHC/NkiQRg=
@@ -197,6 +199,7 @@ github.com/stefanprodan/klog v0.0.0-20190418165334-9cbb78b20423 h1:qTtUiiNM+iq4I
 github.com/stefanprodan/klog v0.0.0-20190418165334-9cbb78b20423/go.mod h1:TYstY5LQfzxFVm9MiiMg7kZ39sc5cue/6CFoY5KgXn8=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.2.0 h1:Hbg2NidpLE8veEBkEZTL3CvlkUIVzuU9jDplZO54c48=
 github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE=
 github.com/stretchr/testify v0.0.0-20151208002404-e3a8ff8ce365/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
@@ -371,4 +374,4 @@ modernc.org/strutil v1.0.0/go.mod h1:lstksw84oURvj9y3tn8lGvRxyRC1S2+g5uuIzNfIOBs
 modernc.org/xc v1.0.0/go.mod h1:mRNCo0bvLjGhHO9WsyuKVU4q0ceiDDDoEeWDJHrNx8I=
 sigs.k8s.io/structured-merge-diff v0.0.0-20190525122527-15d366b2352e/go.mod h1:wWxsB5ozmmv/SG7nM11ayaAW51xMvak/t1r0CSlcokI=
 sigs.k8s.io/yaml v1.1.0 h1:4A07+ZFc2wgJwo8YNlQpr1rVlgUDlxXHhPJciaPY5gs=
-sigs.k8s.io/yaml v1.1.0/go.mod h1:UJmg0vDUVViEyp3mgSv9WPwZCDxu4rQW1olrI1uml+o=
+sigs.k8s.io/yaml v1.1.0/go.mod h1:UJmg0vDUVViEyp3mgSv9WPwZCDxu4rQW1olrI1uml+o=

pkg/loadtester/bash.go

@@ -18,20 +18,20 @@ func (task *BashTask) Hash() string {
  return hash(task.canary + task.command)
 }
 
-func (task *BashTask) Run(ctx context.Context) (bool, error) {
+func (task *BashTask) Run(ctx context.Context) (*TaskRunResult, error) {
  cmd := exec.CommandContext(ctx, "bash", "-c", task.command)
  out, err := cmd.CombinedOutput()
 
  if err != nil {
  task.logger.With("canary", task.canary).Errorf("command failed %s %v %s", task.command, err, out)
- return false, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
+ return &TaskRunResult{false, out}, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
  } else {
  if task.logCmdOutput {
  fmt.Printf("%s\n", out)
  }
  task.logger.With("canary", task.canary).Infof("command finished %s", task.command)
  }
- return true, nil
+ return &TaskRunResult{true, out}, nil
 }
 
 func (task *BashTask) String() string {

pkg/loadtester/concord.go

@@ -121,14 +121,15 @@ func (task *ConcordTask) String() string {
  return fmt.Sprintf("%s %s %s %s", task.Org, task.Project, task.Repo, task.Entrypoint)
 }
 
-func (task *ConcordTask) Run(ctx context.Context) (bool, error) {
+func (task *ConcordTask) Run(ctx context.Context) (*TaskRunResult, error) {
  instance, err := task.startProcess()
  if err != nil {
  task.logger.Errorf("failed to start process: %s", err.Error())
- return false, err
+ return &TaskRunResult{false, nil}, err
  }
 
- return task.checkStatus(ctx, instance, task.PollInterval)
+ ok, err := task.checkStatus(ctx, instance, task.PollInterval)
+ return &TaskRunResult{ok, nil}, err
 }
 
 type concordProcess struct {

pkg/loadtester/helm.go

@@ -19,22 +19,22 @@ func (task *HelmTask) Hash() string {
  return hash(task.canary + task.command)
 }
 
-func (task *HelmTask) Run(ctx context.Context) (bool, error) {
+func (task *HelmTask) Run(ctx context.Context) (*TaskRunResult, error) {
  helmCmd := fmt.Sprintf("%s %s", TaskTypeHelm, task.command)
  task.logger.With("canary", task.canary).Infof("running command %v", helmCmd)
 
  cmd := exec.CommandContext(ctx, TaskTypeHelm, strings.Fields(task.command)...)
  out, err := cmd.CombinedOutput()
  if err != nil {
  task.logger.With("canary", task.canary).Errorf("command failed %s %v %s", task.command, err, out)
- return false, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
+ return &TaskRunResult{false, out}, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
  } else {
  if task.logCmdOutput {
  fmt.Printf("%s\n", out)
  }
  task.logger.With("canary", task.canary).Infof("command finished %v", helmCmd)
  }
- return true, nil
+ return &TaskRunResult{true, out}, nil
 }
 
 func (task *HelmTask) String() string {

pkg/loadtester/helmv3.go

@@ -19,22 +19,22 @@ func (task *HelmTaskv3) Hash() string {
  return hash(task.canary + task.command)
 }
 
-func (task *HelmTaskv3) Run(ctx context.Context) (bool, error) {
+func (task *HelmTaskv3) Run(ctx context.Context) (*TaskRunResult, error) {
  helmCmd := fmt.Sprintf("%s %s", TaskTypeHelmv3, task.command)
  task.logger.With("canary", task.canary).Infof("running command %v", helmCmd)
 
  cmd := exec.CommandContext(ctx, TaskTypeHelmv3, strings.Fields(task.command)...)
  out, err := cmd.CombinedOutput()
  if err != nil {
  task.logger.With("canary", task.canary).Errorf("command failed %s %v %s", task.command, err, out)
- return false, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
+ return &TaskRunResult{false, out}, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
  } else {
  if task.logCmdOutput {
  fmt.Printf("%s\n", out)
  }
  task.logger.With("canary", task.canary).Infof("command finished %v", helmCmd)
  }
- return true, nil
+ return &TaskRunResult{true, out}, nil
 }
 
 func (task *HelmTaskv3) String() string {

pkg/loadtester/runner.go

@@ -9,6 +9,13 @@ import (
  "go.uber.org/zap"
 )
 
+type TaskRunnerInterface interface {
+ Add(task Task)
+ GetTotalExecs() uint64
+ Start(interval time.Duration, stopCh <-chan struct{})
+ Timeout() time.Duration
+}
+
 type TaskRunner struct {
  logger *zap.SugaredLogger
  timeout time.Duration
@@ -81,3 +88,7 @@ func (tr *TaskRunner) Start(interval time.Duration, stopCh <-chan struct{}) {
  }
  }
 }
+
+func (tr *TaskRunner) Timeout() time.Duration {
+ return tr.timeout
+}

pkg/loadtester/server.go

@@ -6,6 +6,7 @@ import (
  "fmt"
  "io/ioutil"
  "net/http"
+ "strconv"
  "time"
 
  "github.com/prometheus/client_golang/prometheus/promhttp"
@@ -17,10 +18,7 @@ import (
 func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogger, taskRunner *TaskRunner, gate *GateStorage, stopCh <-chan struct{}) {
  mux := http.DefaultServeMux
  mux.Handle("/metrics", promhttp.Handler())
- mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
- w.WriteHeader(http.StatusOK)
- w.Write([]byte("OK"))
- })
+ mux.HandleFunc("/healthz", HandleHealthz)
  mux.HandleFunc("/gate/approve", func(w http.ResponseWriter, r *http.Request) {
  w.WriteHeader(http.StatusOK)
  w.Write([]byte("OK"))
@@ -187,7 +185,43 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
  logger.Infof("%s rollback closed", canaryName)
  })
 
- mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+ mux.HandleFunc("/", HandleNewTask(logger, taskRunner))
+ srv := &http.Server{
+ Addr: ":" + port,
+ Handler: mux,
+ ReadTimeout: 5 * time.Second,
+ WriteTimeout: timeout,
+ IdleTimeout: 15 * time.Second,
+ }
+
+ // run server in background
+ go func() {
+ if err := srv.ListenAndServe(); err != http.ErrServerClosed {
+ logger.Fatalf("HTTP server crashed %v", err)
+ }
+ }()
+
+ // wait for SIGTERM or SIGINT
+ <-stopCh
+ ctx, cancel := context.WithTimeout(context.Background(), timeout)
+ defer cancel()
+
+ if err := srv.Shutdown(ctx); err != nil {
+ logger.Errorf("HTTP server graceful shutdown failed %v", err)
+ } else {
+ logger.Info("HTTP server stopped")
+ }
+}
+
+// HandleHealthz handles heath check requests
+func HandleHealthz(w http.ResponseWriter, r *http.Request) {
+ w.WriteHeader(http.StatusOK)
+ w.Write([]byte("OK"))
+}
+
+// HandleNewTask handles task creation requests
+func HandleNewTask(logger *zap.SugaredLogger, taskRunner TaskRunnerInterface) func(w http.ResponseWriter, r *http.Request) {
+ return func(w http.ResponseWriter, r *http.Request) {
  body, err := ioutil.ReadAll(r.Body)
  if err != nil {
  logger.Error("reading the request body failed", zap.Error(err))
@@ -211,6 +245,11 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
  typ = TaskTypeShell
  }
 
+ rtnCmdOutput := false
+ if rtn, ok := metadata["returnCmdOutput"]; ok {
+ rtnCmdOutput, err = strconv.ParseBool(rtn)
+ }
+
  // run bats command (blocking task)
  if typ == TaskTypeBash {
  logger.With("canary", payload.Name).Infof("bats command %s", payload.Metadata["cmd"])
@@ -224,17 +263,20 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
  },
  }
 
- ctx, cancel := context.WithTimeout(context.Background(), taskRunner.timeout)
+ ctx, cancel := context.WithTimeout(context.Background(), taskRunner.Timeout())
  defer cancel()
 
- ok, err := bats.Run(ctx)
- if !ok {
+ result, err := bats.Run(ctx)
+ if !result.ok {
  w.WriteHeader(http.StatusInternalServerError)
  w.Write([]byte(err.Error()))
  return
  }
 
  w.WriteHeader(http.StatusOK)
+ if rtnCmdOutput {
+ w.Write(result.out)
+ }
  return
  }
 
@@ -249,17 +291,20 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
  },
  }
 
- ctx, cancel := context.WithTimeout(context.Background(), taskRunner.timeout)
+ ctx, cancel := context.WithTimeout(context.Background(), taskRunner.Timeout())
  defer cancel()
 
- ok, err := helm.Run(ctx)
- if !ok {
+ result, err := helm.Run(ctx)
+ if !result.ok {
  w.WriteHeader(http.StatusInternalServerError)
  w.Write([]byte(err.Error()))
  return
  }
 
  w.WriteHeader(http.StatusOK)
+ if rtnCmdOutput {
+ w.Write(result.out)
+ }
  return
  }
 
@@ -274,17 +319,20 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
  },
  }
 
- ctx, cancel := context.WithTimeout(context.Background(), taskRunner.timeout)
+ ctx, cancel := context.WithTimeout(context.Background(), taskRunner.Timeout())
  defer cancel()
 
- ok, err := helm.Run(ctx)
- if !ok {
+ result, err := helm.Run(ctx)
+ if !result.ok {
  w.WriteHeader(http.StatusInternalServerError)
  w.Write([]byte(err.Error()))
  return
  }
 
  w.WriteHeader(http.StatusOK)
+ if rtnCmdOutput {
+ w.Write(result.out)
+ }
  return
  }
 
@@ -299,11 +347,11 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
  return
  }
 
- ctx, cancel := context.WithTimeout(context.Background(), taskRunner.timeout)
+ ctx, cancel := context.WithTimeout(context.Background(), taskRunner.Timeout())
  defer cancel()
 
- ok, err := concord.Run(ctx)
- if !ok {
+ result, err := concord.Run(ctx)
+ if !result.ok {
  if err != nil {
  logger.With("canary", payload.Name).Errorf("concord task error: %s", err)
  }
@@ -313,6 +361,9 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
  }
 
  w.WriteHeader(http.StatusOK)
+ if rtnCmdOutput {
+ w.Write(result.out)
+ }
  return
  }
 
@@ -337,30 +388,5 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
  }
 
  w.WriteHeader(http.StatusAccepted)
- })
- srv := &http.Server{
- Addr: ":" + port,
- Handler: mux,
- ReadTimeout: 5 * time.Second,
- WriteTimeout: timeout,
- IdleTimeout: 15 * time.Second,
- }
-
- // run server in background
- go func() {
- if err := srv.ListenAndServe(); err != http.ErrServerClosed {
- logger.Fatalf("HTTP server crashed %v", err)
- }
- }()
-
- // wait for SIGTERM or SIGINT
- <-stopCh
- ctx, cancel := context.WithTimeout(context.Background(), timeout)
- defer cancel()
-
- if err := srv.Shutdown(ctx); err != nil {
- logger.Errorf("HTTP server graceful shutdown failed %v", err)
- } else {
- logger.Info("HTTP server stopped")
  }
 }

pkg/loadtester/server_fixture_test.go

@@ -0,0 +1,46 @@
+package loadtester
+
+import (
+ "net/http/httptest"
+ "time"
+
+ "github.com/weaveworks/flagger/pkg/logger"
+ "go.uber.org/zap"
+)
+
+type serverFixture struct {
+ taskRunner *MockTaskRunner
+ resp *httptest.ResponseRecorder
+ logger *zap.SugaredLogger
+}
+
+func newServerFixture() serverFixture {
+ taskRunner := &MockTaskRunner{}
+ resp := httptest.NewRecorder()
+ logger, _ := logger.NewLogger("info")
+
+ return serverFixture{
+ taskRunner: taskRunner,
+ resp: resp,
+ logger: logger,
+ }
+}
+
+type MockTaskRunner struct {
+}
+
+func (m *MockTaskRunner) Add(task Task) {
+
+}
+
+func (m *MockTaskRunner) GetTotalExecs() uint64 {
+ return 0
+}
+
+func (m *MockTaskRunner) Start(interval time.Duration, stopCh <-chan struct{}) {
+
+}
+
+func (m *MockTaskRunner) Timeout() time.Duration {
+ return time.Hour
+}