fix dragonfly response when p2p preheat (#15916)

* fix dragonfly response when p2p preheat

Signed-off-by: anyiwei <anyiwei@didichuxing.com>

src/pkg/p2p/preheat/provider/dragonfly.go

@@ -2,13 +2,13 @@ package provider
 
 import (
 	"encoding/json"
-	"errors"
 	"fmt"
 	"net/http"
 	"strings"
 
 	common_http "github.com/goharbor/harbor/src/common/http"
 	"github.com/goharbor/harbor/src/lib"
+	"github.com/goharbor/harbor/src/lib/errors"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/models/provider"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/provider/auth"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/provider/client"
@@ -19,6 +19,7 @@ const (
 	preheatEndpoint     = "/preheats"
 	preheatTaskEndpoint = "/preheats/{task_id}"
 	dragonflyPending    = "WAITING"
+	dragonflyFailed     = "FAILED"
 )
 
 type dragonflyPreheatCreateResp struct {
@@ -29,6 +30,7 @@ type dragonflyPreheatInfo struct {
 	ID         string `json:"ID"`
 	StartTime  string `json:"startTime,omitempty"`
 	FinishTime string `json:"finishTime,omitempty"`
+	ErrorMsg   string `json:"errorMsg"`
 	Status     string
 }
 
@@ -109,6 +111,51 @@ func (dd *DragonflyDriver) Preheat(preheatingImage *PreheatImage) (*PreheatingSt
 
 // CheckProgress implements @Driver.CheckProgress.
 func (dd *DragonflyDriver) CheckProgress(taskID string) (*PreheatingStatus, error) {
+	status, err := dd.getProgressStatus(taskID)
+	if err != nil {
+		return nil, err
+	}
+
+	// If preheat job already exists
+	if strings.Index(status.ErrorMsg, "preheat task already exists, id:") >= 0 {
+		if taskID, err = getTaskExistedFromErrMsg(status.ErrorMsg); err != nil {
+			return nil, err
+		}
+		if status, err = dd.getProgressStatus(taskID); err != nil {
+			return nil, err
+		}
+	}
+
+	if status.Status == dragonflyPending {
+		status.Status = provider.PreheatingStatusPending
+	} else if status.Status == dragonflyFailed {
+		status.Status = provider.PreheatingStatusFail
+	}
+
+	res := &PreheatingStatus{
+		Status: status.Status,
+		TaskID: taskID,
+	}
+	if status.StartTime != "" {
+		res.StartTime = status.StartTime
+	}
+	if status.FinishTime != "" {
+		res.FinishTime = status.FinishTime
+	}
+
+	return res, nil
+}
+
+func getTaskExistedFromErrMsg(msg string) (string, error) {
+	begin := strings.Index(msg, "preheat task already exists, id:") + 32
+	end := strings.LastIndex(msg, "\"}")
+	if end-begin <= 0 {
+		return "", errors.Errorf("can't find existed task id by error msg:%s", msg)
+	}
+	return msg[begin:end], nil
+}
+
+func (dd *DragonflyDriver) getProgressStatus(taskID string) (*dragonflyPreheatInfo, error) {
 	if dd.instance == nil {
 		return nil, errors.New("missing instance metadata")
 	}
@@ -128,23 +175,7 @@ func (dd *DragonflyDriver) CheckProgress(taskID string) (*PreheatingStatus, erro
 	if err := json.Unmarshal(bytes, status); err != nil {
 		return nil, err
 	}
-
-	if status.Status == dragonflyPending {
-		status.Status = provider.PreheatingStatusPending
-	}
-
-	res := &PreheatingStatus{
-		Status: status.Status,
-		TaskID: taskID,
-	}
-	if status.StartTime != "" {
-		res.StartTime = status.StartTime
-	}
-	if status.FinishTime != "" {
-		res.FinishTime = status.FinishTime
-	}
-
-	return res, nil
+	return status, nil
 }
 
 func (dd *DragonflyDriver) getCred() *auth.Credential {

src/pkg/p2p/preheat/provider/dragonfly_test.go

@@ -113,4 +113,30 @@ func (suite *DragonflyTestSuite) TestCheckProgress() {
 	st, err := suite.driver.CheckProgress("dragonfly-id")
 	require.NoError(suite.T(), err, "get preheat status")
 	suite.Equal(provider.PreheatingStatusSuccess, st.Status, "preheat status")
+
+	// preheat job exit but returns no id
+	st, err = suite.driver.CheckProgress("preheat-job-exist-with-no-id")
+	require.Error(suite.T(), err, "get preheat status")
+
+	// preheat job exit returns id but get info with that failed
+	st, err = suite.driver.CheckProgress("preheat-job-exist-with-id-1")
+	require.Error(suite.T(), err, "get preheat status")
+
+	// preheat job normal failed
+	st, err = suite.driver.CheckProgress("preheat-job-normal-failed")
+	require.NoError(suite.T(), err, "get preheat status")
+	suite.Equal(provider.PreheatingStatusFail, st.Status, "preheat status")
+
+	// instance is empty
+	testDriver := &DragonflyDriver{}
+	st, err = testDriver.CheckProgress("")
+	require.Error(suite.T(), err, "get preheat status")
+
+	// preheat job with no task id
+	st, err = suite.driver.CheckProgress("")
+	require.Error(suite.T(), err, "get preheat status")
+
+	// preheat job with err json response
+	st, err = suite.driver.CheckProgress("preheat-job-err-body-json")
+	require.Error(suite.T(), err, "get preheat status")
 }

src/pkg/p2p/preheat/provider/mock.go

@@ -95,6 +95,61 @@ func MockDragonflyProvider() *httptest.Server {
 			}
 			bytes, _ := json.Marshal(status)
 			_, _ = w.Write(bytes)
+		case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-exist-with-no-id", 1):
+			if r.Method != http.MethodGet {
+				w.WriteHeader(http.StatusNotImplemented)
+				return
+			}
+			status := &dragonflyPreheatInfo{
+				ID:         "preheat-exist-with-no-id",
+				StartTime:  time.Now().UTC().String(),
+				FinishTime: time.Now().Add(5 * time.Minute).UTC().String(),
+				Status:     "FAILED",
+				ErrorMsg:   "{\"Code\":208,\"Msg\":\"preheat task already exists, id:\"}",
+			}
+			bytes, _ := json.Marshal(status)
+			_, _ = w.Write(bytes)
+		case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-normal-failed", 1):
+			if r.Method != http.MethodGet {
+				w.WriteHeader(http.StatusNotImplemented)
+				return
+			}
+			status := &dragonflyPreheatInfo{
+				ID:         "preheat-job-exist-with-id-1",
+				StartTime:  time.Now().UTC().String(),
+				FinishTime: time.Now().Add(5 * time.Minute).UTC().String(),
+				Status:     "FAILED",
+				ErrorMsg:   "{\"Code\":208,\"Msg\":\"some msg\"}",
+			}
+			bytes, _ := json.Marshal(status)
+			_, _ = w.Write(bytes)
+		case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-exist-with-id-1", 1):
+			if r.Method != http.MethodGet {
+				w.WriteHeader(http.StatusNotImplemented)
+				return
+			}
+			status := &dragonflyPreheatInfo{
+				ID:         "preheat-job-exist-with-id-1",
+				StartTime:  time.Now().UTC().String(),
+				FinishTime: time.Now().Add(5 * time.Minute).UTC().String(),
+				Status:     "FAILED",
+				ErrorMsg:   "{\"Code\":208,\"Msg\":\"preheat task already exists, id:preheat-job-exist-with-id-1-1\"}",
+			}
+			bytes, _ := json.Marshal(status)
+			_, _ = w.Write(bytes)
+		case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-exist-with-id-1-1", 1):
+			if r.Method != http.MethodGet {
+				w.WriteHeader(http.StatusNotImplemented)
+				return
+			}
+			w.WriteHeader(http.StatusInternalServerError)
+		case strings.Replace(preheatTaskEndpoint, "{task_id}", "preheat-job-err-body-json", 1):
+			if r.Method != http.MethodGet {
+				w.WriteHeader(http.StatusNotImplemented)
+				return
+			}
+			bodyStr := "\"err body\""
+			_, _ = w.Write([]byte(bodyStr))
 		default:
 			w.WriteHeader(http.StatusNotImplemented)
 		}