[agent v0.3.6] 🐛 fix: kill process group

cmd/agent/main.go

@@ -234,18 +234,32 @@ func doTask(task *pb.Task) {
 		startedAt := time.Now()
 		var cmd *exec.Cmd
 		var endCh = make(chan struct{})
+		var pg utils.ProcessExitGroup
 		timeout := time.NewTimer(time.Hour * 2)
 		if utils.IsWindows() {
+			var err error
+			pg, err = utils.NewProcessExitGroup()
+			if err != nil {
+				// Windows 进程组创建失败，直接退出
+				result.Data = err.Error()
+				client.ReportTask(ctx, &result)
+				return
+			}
 			cmd = exec.Command("cmd", "/c", task.GetData())
+			pg.AddProcess(cmd.Process)
 		} else {
 			cmd = exec.Command("sh", "-c", task.GetData())
+			cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 		}
-		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 		go func() {
 			select {
 			case <-timeout.C:
 				result.Data = "任务执行超时\n"
-				cmd.Process.Kill()
+				if utils.IsWindows() {
+					pg.Dispose()
+				} else {
+					cmd.Process.Kill()
+				}
 				close(endCh)
 			case <-endCh:
 			}

cmd/playground/main.go

@@ -80,13 +80,19 @@ func cmdExec() {
 		panic(err)
 	}
 	var cmd *exec.Cmd
+	var pg utils.ProcessExitGroup
 	if utils.IsWindows() {
+		pg, err = utils.NewProcessExitGroup()
+		if err != nil {
+			panic(err)
+		}
 		cmd = exec.Command("cmd", "/c", execFrom+"/cmd/playground/example.sh hello asd")
+		pg.AddProcess(cmd.Process)
 	} else {
 		cmd = exec.Command("sh", "-c", execFrom+`/cmd/playground/example.sh hello && \
 echo world!`)
+		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 	}
-	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 	var endCh = make(chan struct{})
 	go func() {
 		output, err := cmd.Output()
@@ -97,8 +103,14 @@ echo world!`)
 	go func() {
 		time.Sleep(time.Second * 2)
 		fmt.Println("killed")
-		if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
-			panic(err)
+		if utils.IsWindows() {
+			if err := pg.Dispose(); err != nil {
+				panic(err)
+			}
+		} else {
+			if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
+				panic(err)
+			}
 		}
 	}()
 	select {

go.mod

@@ -22,6 +22,7 @@ require (
 	github.com/spf13/viper v1.7.1
 	github.com/stretchr/testify v1.6.1
 	golang.org/x/oauth2 v0.0.0-20200902213428-5d25da1a8d43
+	golang.org/x/sys v0.0.0-20201024232916-9f70ab9862d5
 	google.golang.org/grpc v1.33.1
 	google.golang.org/protobuf v1.25.0
 	gopkg.in/yaml.v2 v2.2.8

pkg/utils/proccess_group_darwin.go

@@ -0,0 +1,20 @@
+package utils
+
+import (
+	"errors"
+	"os"
+)
+
+type ProcessExitGroup struct{}
+
+func NewProcessExitGroup() (ProcessExitGroup, error) {
+	return ProcessExitGroup{}, errors.New("not implement")
+}
+
+func (g ProcessExitGroup) Dispose() error {
+	return errors.New("not implement")
+}
+
+func (g ProcessExitGroup) AddProcess(p *os.Process) error {
+	return errors.New("not implement")
+}

pkg/utils/proccess_group_linux.go

@@ -0,0 +1,20 @@
+package utils
+
+import (
+	"errors"
+	"os"
+)
+
+type ProcessExitGroup struct{}
+
+func NewProcessExitGroup() (ProcessExitGroup, error) {
+	return ProcessExitGroup{}, errors.New("not implement")
+}
+
+func (g ProcessExitGroup) Dispose() error {
+	return errors.New("not implement")
+}
+
+func (g ProcessExitGroup) AddProcess(p *os.Process) error {
+	return errors.New("not implement")
+}

pkg/utils/proccess_group_windows.go

@@ -0,0 +1,49 @@
+package utils
+
+import (
+	"os"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+// We use this struct to retreive process handle(which is unexported)
+// from os.Process using unsafe operation.
+type process struct {
+	Pid    int
+	Handle uintptr
+}
+
+type ProcessExitGroup windows.Handle
+
+func NewProcessExitGroup() (ProcessExitGroup, error) {
+	handle, err := windows.CreateJobObject(nil, nil)
+	if err != nil {
+		return 0, err
+	}
+
+	info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{
+		BasicLimitInformation: windows.JOBOBJECT_BASIC_LIMIT_INFORMATION{
+			LimitFlags: windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+		},
+	}
+	if _, err := windows.SetInformationJobObject(
+		handle,
+		windows.JobObjectExtendedLimitInformation,
+		uintptr(unsafe.Pointer(&info)),
+		uint32(unsafe.Sizeof(info))); err != nil {
+		return 0, err
+	}
+
+	return ProcessExitGroup(handle), nil
+}
+
+func (g ProcessExitGroup) Dispose() error {
+	return windows.CloseHandle(windows.Handle(g))
+}
+
+func (g ProcessExitGroup) AddProcess(p *os.Process) error {
+	return windows.AssignProcessToJobObject(
+		windows.Handle(g),
+		windows.Handle((*process)(unsafe.Pointer(p)).Handle))
+}