feat: add performance testing to run on the page

Author: zhufuyi

cmd/sponge/server/handler.go

@@ -10,7 +10,9 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"sort"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -22,6 +24,7 @@ import (
 	"github.com/go-dev-frame/sponge/pkg/gofile"
 	"github.com/go-dev-frame/sponge/pkg/krand"
 	"github.com/go-dev-frame/sponge/pkg/mgo"
+	"github.com/go-dev-frame/sponge/pkg/process"
 	"github.com/go-dev-frame/sponge/pkg/sgorm"
 	"github.com/go-dev-frame/sponge/pkg/sgorm/mysql"
 	"github.com/go-dev-frame/sponge/pkg/sgorm/postgresql"
@@ -307,7 +310,7 @@ func HandleAssistant(c *gin.Context) {
 	args := strings.Split(form.Arg, " ")
 	params := parseCommandArgs(args)
 
-	ctx, _ := context.WithTimeout(context.Background(), time.Minute*30) // nolint
+	ctx, _ := context.WithTimeout(context.Background(), time.Minute*60) // nolint
 	result := gobash.Run(ctx, "sponge", args...)
 	resultInfo := ""
 	count := 0
@@ -331,6 +334,45 @@ func HandleAssistant(c *gin.Context) {
 	response.Success(c, resultInfo)
 }
 
+var processMap = sync.Map{}
+
+func getCommand(args []string) string {
+	if len(args) < 4 {
+		return ""
+	}
+	commandArgs := []string{"sponge", args[0], args[1]}
+	for _, arg := range args {
+		if strings.Contains(arg, "--url") {
+			commandArgs = append(commandArgs, arg)
+		}
+		if strings.Contains(arg, "--method") {
+			commandArgs = append(commandArgs, arg)
+		}
+	}
+	sort.Strings(commandArgs)
+	return strings.Join(commandArgs, "&")
+}
+
+func addProcess(key string, pid int) {
+	processMap.Store(key, pid)
+}
+
+func getPid(key string) (int, bool) {
+	value, ok := processMap.Load(key)
+	if !ok {
+		return -1, false
+	}
+	valueInt, ok := value.(int)
+	if !ok {
+		return -1, false
+	}
+	return valueInt, true
+}
+
+func removeProcess(key string) {
+	processMap.Delete(key)
+}
+
 // HandlePerformanceTest handle performance test
 func HandlePerformanceTest(c *gin.Context) {
 	form := &GenerateCodeForm{}
@@ -358,13 +400,22 @@ func HandlePerformanceTest(c *gin.Context) {
 		params.PushType = "custom"
 	}
 
-	ctx, _ := context.WithTimeout(context.Background(), time.Minute*30) // nolint
+	ctx, _ := context.WithTimeout(context.Background(), time.Hour*24) // nolint
 	result := gobash.Run(ctx, "sponge", args...)
+	key := ""
+	pid := 0
 	resultInfo := ""
 	count := 0
 	for v := range result.StdOut {
 		count++
-		if count == 1 { // first line is the command
+		if count == 1 { // first line is the command and pid value
+			pid = gobash.ParsePid(v)
+			if pid > 0 {
+				key = getCommand(args)
+				if result.Pid > 0 {
+					addProcess(key, result.Pid)
+				}
+			}
 			continue
 		}
 		if strings.Contains(v, "Waiting for assistant responses") {
@@ -379,12 +430,48 @@ func HandlePerformanceTest(c *gin.Context) {
 
 	recordObj().set(c.ClientIP(), form.Path, params)
 
-	resultInfo = splitString(resultInfo, "Performance Test Report ==========")
+	if result.Pid > 0 {
+		removeProcess(key)
+	}
 
-	response.Success(c, resultInfo)
+	lineContent, reportStr := splitString(resultInfo, "Performance Test Report ==========")
+	command := "sponge " + strings.Join(args, " ")
+	insertStr := fmt.Sprintf(`
+
+[Overview]
+  • %-19s%s
+  • %-19s%s`, "Command:", command,
+		"End Time:", time.Now().Format(time.DateTime))
+	reportStr = strings.ReplaceAll(reportStr, lineContent, lineContent+insertStr)
+
+	response.Success(c, reportStr)
+}
+
+// HandleStopPerformanceTest handle stop performance test
+func HandleStopPerformanceTest(c *gin.Context) {
+	form := &GenerateCodeForm{}
+	err := c.ShouldBindJSON(form)
+	if err != nil {
+		responseErr(c, err, errcode.InvalidParams)
+		return
+	}
+
+	args := strings.Split(form.Arg, " ")
+	key := getCommand(args)
+	pid, _ := getPid(key)
+	if pid > 0 {
+		err = process.Kill(pid)
+		if err != nil {
+			responseErr(c, err, errcode.InternalServerError)
+			return
+		}
+		removeProcess(key)
+	}
+
+	response.Success(c)
 }
 
-func splitString(str string, sep string) string {
+func splitString(str string, sep string) (lineContent string, out string) {
 	lines := strings.Split(str, "\n")
 	startIndex := 0
 	isFound := false
@@ -393,14 +480,15 @@ func splitString(str string, sep string) string {
 		if strings.Contains(line, sep) {
 			isFound = true
 			startIndex = i
+			lineContent = line
 			break
 		}
 	}
 
 	if isFound {
-		return strings.Join(lines[startIndex:], "\n")
+		return lineContent, strings.Join(lines[startIndex:], "\n")
 	}
-	return str
+	return "", str
 }
 
 // GetRecord generate run command record

cmd/sponge/server/http.go

@@ -59,6 +59,7 @@ func NewRouter(spongeAddr string, isLog bool) *gin.Engine {
 	apiV1.POST("/getTemplateInfo", GetTemplateInfo)
 	apiV1.POST("/assistant", HandleAssistant)
 	apiV1.POST("/performanceTest", HandlePerformanceTest)
+	apiV1.POST("/performanceTest/stop", HandleStopPerformanceTest)
 	apiV1.POST("/uploadFiles", UploadFiles)
 	apiV1.POST("/listTables", ListTables)
 	apiV1.GET("/listDrivers", ListDbDrivers)

cmd/sponge/server/record.go

@@ -38,12 +38,13 @@ type parameters struct {
 	DepProtoDir string `json:"depProtoDir"`
 	OnlyPrint   bool   `json:"onlyPrint"`
 
-	LLMType           string `json:"llmType"`
-	LLMModel          string `json:"llmModel"`
-	APIKey            string `json:"apiKey"`
-	GoDir             string `json:"goDir"`
-	GoFile            string `json:"goFile"`
-	IsSpecifiedGoFile bool   `json:"isSpecifiedGoFile"`
+	LLMType       string `json:"llmType"`
+	LLMModel      string `json:"llmModel"`
+	APIKey        string `json:"apiKey"`
+	GoDir         string `json:"goDir"`
+	GoFile        string `json:"goFile"`
+	TargetType    string `json:"targetType"`
+	IsCleanAICode bool   `json:"isCleanAICode"`
 
 	Protocol      string   `json:"protocol"`
 	URL           string   `json:"url"`
@@ -143,6 +144,9 @@ func parseCommandArgs(args []string) *parameters {
 			if ss[0] == "--only-print" {
 				params.OnlyPrint = true
 			}
+			if ss[0] == "--is-clean" {
+				params.IsCleanAICode = true
+			}
 		} else {
 			val := ss[1]
 			switch ss[0] {
@@ -197,10 +201,12 @@ func parseCommandArgs(args []string) *parameters {
 				params.APIKey = val
 			case "--dir":
 				params.GoDir = val
-				params.IsSpecifiedGoFile = false
+				params.TargetType = "goDir"
 			case "--file":
 				params.GoFile = val
-				params.IsSpecifiedGoFile = true
+				params.TargetType = "goFile"
+			case "--is-clean":
+				params.IsCleanAICode = val == "true"
 
 			case "--url":
 				params.URL = val

cmd/sponge/server/static/appConfig.js

@@ -1,3 +1,3 @@
-var appConfig = {
+const appConfig = {
   spongeServiceAddr: "http://localhost:24631/api/v1",
 };

cmd/sponge/server/static/css/app.161971ad6f41b518cab65fa120506c99.css

@@ -1,2 +1,2 @@
-<!DOCTYPE html><html><head><meta charset=utf-8><meta name=viewport content="width=device-width,initial-scale=1"><title>Sponge Generate Code</title><link rel=icon type=image/png sizes=32x32 href="/static/favicon.png?v=1.0"><script type=text/javascript src=/static/appConfig.js async></script><link href=/static/css/app.81ae100c53d45c1085ec370cf96a8ce1.css rel=stylesheet></head><body style="margin: 0px; padding: 0px;"><style>.el-tooltip__popper {box-shadow: 3px 3px 10px 5px #d3d3d6;border-width: 0px !important;}
-      .el-tooltip__popper[x-placement^="top"] .popper__arrow:after {border-top-color: #dcdfe6 !important;}</style><div id=app></div><script type=text/javascript src=/static/js/manifest.2ae2e69a05c33dfc65f8.js></script><script type=text/javascript src=/static/js/vendor.d8cdb3748af43d2ae51a.js></script><script type=text/javascript src=/static/js/app.abed82fc254ec468ebb1.js></script></body></html>
+<!DOCTYPE html><html><head><meta charset=utf-8><meta name=viewport content="width=device-width,initial-scale=1"><title>Sponge Generate Code</title><link rel=icon type=image/png sizes=32x32 href="/static/favicon.png?v=1.0"><script type=text/javascript src=/static/appConfig.js async></script><link href=/static/css/app.161971ad6f41b518cab65fa120506c99.css rel=stylesheet></head><body style="margin: 0px; padding: 0px;"><style>.el-tooltip__popper {box-shadow: 3px 3px 10px 5px #d3d3d6;border-width: 0px !important;}
+      .el-tooltip__popper[x-placement^="top"] .popper__arrow:after {border-top-color: #dcdfe6 !important;}</style><div id=app></div><script type=text/javascript src=/static/js/manifest.2ae2e69a05c33dfc65f8.js></script><script type=text/javascript src=/static/js/vendor.81645738747fb507a683.js></script><script type=text/javascript src=/static/js/app.7d14490af8ba07678235.js></script></body></html>