feat: 支持设置扫描超时时间 #70 (#72)

* feat: 心跳失败时终止任务 #71

* feat: 支持设置任务执行超时时间 #70

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

@@ -60,7 +60,7 @@ func GetClient(args *object.Arguments) *BkRepoClient {
 }
 
 // Start 开始分析
-func (c *BkRepoClient) Start(ctx context.Context) (*object.ToolInput, error) {
+func (c *BkRepoClient) Start(ctx context.Context, cancel context.CancelFunc) (*object.ToolInput, error) {
 	if c.ToolInput == nil {
 		if err := c.initToolInput(); err != nil {
 			return nil, err
@@ -73,7 +73,7 @@ func (c *BkRepoClient) Start(ctx context.Context) (*object.ToolInput, error) {
 				return nil, err
 			}
 			if c.Args.Heartbeat > 0 {
-				go c.heartbeat(ctx)
+				go c.heartbeat(ctx, cancel)
 			}
 			util.Info("update subtask status success")
 		}
@@ -188,7 +188,7 @@ func (c *BkRepoClient) updateSubtaskStatus() error {
 	return nil
 }
 
-func (c *BkRepoClient) heartbeat(ctx context.Context) {
+func (c *BkRepoClient) heartbeat(ctx context.Context, cancel context.CancelFunc) {
 	ticker := time.NewTicker(time.Duration(c.Args.Heartbeat) * time.Second)
 	taskId := c.ToolInput.TaskId
 	reqUrl := c.Args.Url + analystTemporaryPrefix + "/scan/subtask/" + taskId + "/heartbeat"
@@ -213,6 +213,7 @@ func (c *BkRepoClient) heartbeat(ctx context.Context) {
 				return
 			}
 			if response.StatusCode != http.StatusOK {
+				cancel()
 				b, _ := io.ReadAll(response.Body)
 				util.Error("heartbeat failed: " + response.Status + ", message: " + string(b))
 			}

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

@@ -3,6 +3,7 @@ package framework
 import (
 	"context"
 	"errors"
+	"fmt"
 	"github.com/TencentBlueKing/ci-repoAnalysis/analysis-tool-sdk-golang/api"
 	"github.com/TencentBlueKing/ci-repoAnalysis/analysis-tool-sdk-golang/object"
 	"github.com/TencentBlueKing/ci-repoAnalysis/analysis-tool-sdk-golang/util"
@@ -13,7 +14,7 @@ import (
 type Executor interface {
 	// Execute 框架会调用该函数执行扫描，传入的参数config为工具相关配置，file为待分析的制品
 	// 扫描成功时返回toolOutput，出错时返回error，工具框架会自动上报或输出结果给制品分析服务
-	Execute(config *object.ToolConfig, file *os.File) (*object.ToolOutput, error)
+	Execute(ctx context.Context, config *object.ToolConfig, file *os.File) (*object.ToolOutput, error)
 }
 
 // Analyze 执行分析
@@ -38,18 +39,19 @@ func Analyze(executor Executor) {
 func doAnalyze(executor Executor, arguments *object.Arguments) {
 	client := api.GetClient(arguments)
 	ctx, cancel := context.WithCancel(context.Background())
-	input, err := client.Start(ctx)
+	defer cancel()
+	input, err := client.Start(ctx, cancel)
 	if err != nil {
 		panic("Start analyze failed: " + err.Error())
 	}
 	if input == nil || input.TaskId == "" {
 		util.Info("no subtask found, exit")
-		os.Exit(0)
+		return
 	}
 	file, err := client.GenerateInputFile()
 	if err != nil {
 		client.Failed(cancel, errors.New("Generate input file failed: "+err.Error()))
-		os.Exit(1)
+		return
 	}
 	// 返回的file为nil时表示文件被忽略，直接返回
 	if file == nil {
@@ -59,9 +61,15 @@ func doAnalyze(executor Executor, arguments *object.Arguments) {
 	}
 	defer file.Close()
 	util.Info("generate input file success")
-	output, err := executor.Execute(&input.ToolConfig, file)
+	execCtx, execCancel := context.WithTimeout(ctx, client.ToolInput.MaxTime())
+	defer execCancel()
+	output, err := executor.Execute(execCtx, &input.ToolConfig, file)
 	if err != nil {
-		client.Failed(cancel, errors.New("Execute analysis failed: "+err.Error()))
+		errMsg := "Execute analysis failed: " + err.Error()
+		if ctx.Err() != nil {
+			errMsg = fmt.Sprintf("%s, ctx err[%s]", errMsg, ctx.Err().Error())
+		}
+		client.Failed(cancel, errors.New(errMsg))
 	} else {
 		client.Finish(cancel, output)
 	}

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

@@ -2,6 +2,7 @@ package object
 
 import (
 	"strconv"
+	"time"
 )
 
 // ToolInput 工具输入
@@ -90,3 +91,9 @@ func (toolInput *ToolInput) FileUrlMap() map[string]FileUrl {
 	}
 	return fileUrlMap
 }
+
+// MaxTime 获取允许执行的最长时间
+func (toolInput *ToolInput) MaxTime() time.Duration {
+	maxTime, _ := toolInput.ToolConfig.GetIntArg("maxTime")
+	return time.Duration(maxTime) * time.Millisecond
+}

analysis-tool-sdk-golang/util/cmd.go

@@ -2,15 +2,16 @@ package util
 
 import (
 	"bufio"
+	"context"
 	"errors"
 	"fmt"
 	"os/exec"
 	"strings"
 )
 
 // ExecAndLog 执行命令并实时输出日志
-func ExecAndLog(name string, args []string, workDir string) error {
-	cmd := exec.Command(name, args...)
+func ExecAndLog(ctx context.Context, name string, args []string, workDir string) error {
+	cmd := exec.CommandContext(ctx, name, args...)
 
 	if len(workDir) > 0 {
 		cmd.Dir = workDir

analysis-tool-sdk-golang/util/cmd_test.go

@@ -1,16 +1,29 @@
 package util
 
 import (
+	"context"
 	"os"
 	"testing"
+	"time"
 )
 
 func TestExecAndLog(t *testing.T) {
 	dir, _ := os.Getwd()
 	cmd := "go"
 	args := []string{"version"}
-	err := ExecAndLog(cmd, args, dir)
+	err := ExecAndLog(context.Background(), cmd, args, dir)
 	if err != nil {
 		t.Fatalf("exec cmd %s failed: %s", cmd, err.Error())
 	}
+
+	cmd = "sleep"
+	args = []string{"5"}
+	// test timeout
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(3)*time.Second)
+	err = ExecAndLog(ctx, cmd, args, dir)
+
+	if err != nil {
+		Info("exec cmd %s failed: %s", cmd, err.Error())
+	}
+	cancel()
 }