poll ngrinder task status

pkg/loadtester/task_ngrinder.go

@@ -40,18 +40,19 @@ func init() {
 		}
 		return &NGrinderTask{
 			TaskBase{canary, logger},
-			baseUrl, cloneId, username, string(passwdDecoded), -1,
+			baseUrl, cloneId, username, string(passwdDecoded), -1, 5,
 		}, nil
 	})
 }
 
 type NGrinderTask struct {
 	TaskBase
-	baseUrl  *url.URL
-	cloneId  int
-	username string
-	passwd   string
-	testId   int
+	baseUrl      *url.URL
+	cloneId      int
+	username     string
+	passwd       string
+	testId       int
+	pollInterval time.Duration
 }
 
 func (task *NGrinderTask) Hash() string {
@@ -62,6 +63,10 @@ func (task *NGrinderTask) CloneAndStartEndpoint() *url.URL {
 	path, _ := url.Parse(fmt.Sprintf("perftest/api/%d/clone_and_start", task.cloneId))
 	return task.baseUrl.ResolveReference(path)
 }
+func (task *NGrinderTask) StatusEndpoint() *url.URL {
+	path, _ := url.Parse(fmt.Sprintf("perftest/api/%d/status", task.testId))
+	return task.baseUrl.ResolveReference(path)
+}
 func (task *NGrinderTask) StopEndpoint() *url.URL {
 	path, _ := url.Parse(fmt.Sprintf("perftest/api/%d?action=stop", task.testId))
 	return task.baseUrl.ResolveReference(path)
@@ -75,13 +80,8 @@ func (task *NGrinderTask) Run(ctx context.Context) bool {
 		return false
 	}
 	id := result["id"]
-	testId, ok := id.(int)
-	if !ok {
-		return false
-	} else {
-		task.testId = testId
-		return task.PollStatus(ctx)
-	}
+	task.testId = int(id.(float64))
+	return task.PollStatus(ctx)
 }
 
 func (task *NGrinderTask) String() string {
@@ -89,13 +89,28 @@ func (task *NGrinderTask) String() string {
 }
 
 func (task *NGrinderTask) PollStatus(ctx context.Context) bool {
-	// wait until ngrinder test completed or timedout
-	tickChan := time.NewTicker(time.Second * 15).C
+	// wait until ngrinder test finished/canceled or timedout
+	tickChan := time.NewTicker(time.Second * task.pollInterval).C
 	for {
 		select {
 		case <-tickChan:
-
+			result, err := task.request("GET", task.StatusEndpoint().String(), ctx)
+			if err == nil {
+				statusArray, ok := result["status"].([]interface{})
+				if ok && len(statusArray) > 0 {
+					status := statusArray[0].(map[string]interface{})
+					statusId := status["status_id"]
+					task.logger.Debugf("status of ngrinder task %d is %s", task.testId, statusId)
+					if statusId == "FINISHED" {
+						return true
+					} else if statusId == "STOP_BY_ERROR" || statusId == "CANCELED" || statusId == "UNKNOWN" {
+						return false
+					}
+				}
+			}
 		case <-ctx.Done():
+			task.logger.Warnf("context timedout, top ngrinder task %d forcibly", task.testId)
+			task.request("PUT", task.StopEndpoint().String(), nil)
 			return false
 		}
 	}
@@ -104,16 +119,20 @@ func (task *NGrinderTask) PollStatus(ctx context.Context) bool {
 func (task *NGrinderTask) request(method, url string, ctx context.Context) (map[string]interface{}, error) {
 	req, _ := http.NewRequest(method, url, nil)
 	req.SetBasicAuth(task.username, task.passwd)
-	resp, err := http.DefaultClient.Do(req.WithContext(ctx))
+	if ctx != nil {
+		req = req.WithContext(ctx)
+	}
+	resp, err := http.DefaultClient.Do(req)
 	defer resp.Body.Close()
 	if err != nil {
 		return nil, err
 	}
 	respBytes, err := ioutil.ReadAll(resp.Body)
 	res := make(map[string]interface{})
-	json.Unmarshal(respBytes, res)
-	err = nil
-	if success, ok := res["success"]; ok && success == false {
+	err = json.Unmarshal(respBytes, &res)
+	if err != nil {
+		task.logger.Errorf("bad response, %s ,json expected:\n %s", err.Error(), string(respBytes))
+	} else if success, ok := res["success"]; ok && success == false {
 		err = errors.New(res["message"].(string))
 	}
 	return res, err