feat: 增加--pull-retry参数设置重试拉取任务次数 #45

analysis-tool-sdk-golang/api/bkrepo.go

@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"os"
 	"strconv"
+	"time"
 )
 
 // analystTemporaryPrefix 制品分析服务接口前缀
@@ -187,7 +188,23 @@ func (c *BkRepoClient) fetchToolInput() (*object.ToolInput, error) {
 func (c *BkRepoClient) pullToolInput() (*object.ToolInput, error) {
 	url := c.Args.Url + analystTemporaryPrefix + "/scan/subtask/input?executionCluster=" + c.Args.ExecutionCluster +
 		"&token=" + c.Args.Token
-	return c.doFetchToolInput(url)
+
+	var toolInput *object.ToolInput = nil
+	var err error = nil
+	var pullRetry = c.Args.PullRetry
+	for (toolInput == nil || toolInput.TaskId == "") && pullRetry != 0 {
+		if err != nil {
+			return nil, err
+		}
+		util.Info("try to pull subtask...")
+		toolInput, err = c.doFetchToolInput(url)
+		pullRetry--
+		if toolInput == nil || toolInput.TaskId == "" {
+			time.Sleep(5 * time.Second)
+		}
+	}
+
+	return toolInput, err
 }
 
 func (c *BkRepoClient) doFetchToolInput(url string) (*object.ToolInput, error) {

analysis-tool-sdk-golang/api/bkrepo_test.go

@@ -0,0 +1,26 @@
+package api
+
+import (
+	"fmt"
+	"github.com/TencentBlueKing/ci-repoAnalysis/analysis-tool-sdk-golang/object"
+	"os"
+	"testing"
+)
+
+func TestPullTask(t *testing.T) {
+	client := BkRepoClient{}
+	client.Args = &object.Arguments{
+		Url:              os.Getenv("URL"),
+		Token:            os.Getenv("TOKEN"),
+		ExecutionCluster: os.Getenv("EXECUTION_CLUSTER"),
+		PullRetry:        -1,
+	}
+	toolInput, err := client.pullToolInput()
+	if err != nil {
+		fmt.Println(err.Error())
+	}
+
+	if toolInput != nil {
+		fmt.Println(toolInput.TaskId)
+	}
+}

analysis-tool-sdk-golang/framework/executor.go

@@ -26,6 +26,10 @@ func doAnalyze(executor Executor, arguments *object.Arguments) {
 	if err != nil {
 		panic("Start analyze failed: " + err.Error())
 	}
+	if input == nil || input.TaskId == "" {
+		util.Info("no subtask found, exit")
+		os.Exit(0)
+	}
 	file, err := client.GenerateInputFile()
 	defer file.Close()
 	if err != nil {

analysis-tool-sdk-golang/object/argument.go

@@ -11,6 +11,7 @@ type Arguments struct {
 	Token            string
 	TaskId           string
 	ExecutionCluster string
+	PullRetry        int
 	InputFilePath    string
 	OutputFilePath   string
 }
@@ -23,13 +24,22 @@ func newArguments() *Arguments {
 	flag.StringVar(&args.Token, "token", "", "制品库临时令牌")
 	flag.StringVar(&args.TaskId, "task-id", "", "扫描任务Id")
 	flag.StringVar(&args.ExecutionCluster, "execution-cluster", "", "所在扫描执行集群名")
+	flag.IntVar(&args.PullRetry, "pull-retry", 3, "拉取模式下拉取任务的次数，-1表示一直拉取直到拉取到任务")
 	flag.StringVar(&args.InputFilePath, "input", "", "输入文件路径")
 	flag.StringVar(&args.OutputFilePath, "output", "", "输出文件路径")
 
 	flag.Parse()
 
-	fmt.Printf("url: %s, token: %s, taskId: %s, inputFilePath: %s, outputFilePath: %s\n",
-		args.Url, args.Token, args.TaskId, args.InputFilePath, args.OutputFilePath)
+	fmt.Printf(
+		"url: %s, token: %s, taskId: %s, executionCluster: %s, pull-retry: %d, inputFilePath: %s, outputFilePath: %s\n",
+		args.Url,
+		args.Token,
+		args.TaskId,
+		args.ExecutionCluster,
+		args.PullRetry,
+		args.InputFilePath,
+		args.OutputFilePath,
+	)
 	if (args.Offline() || args.Online()) == false {
 		panic("缺少必要输入参数")
 	}