From 0f24c0d7ef557110137853bf08efa852c0b6092b Mon Sep 17 00:00:00 2001 From: naiba Date: Fri, 29 Jan 2021 11:11:39 +0800 Subject: [PATCH] =?UTF-8?q?[agent=20v0.3.6]=20=F0=9F=90=9B=20fix:=20kill?= =?UTF-8?q?=20process=20group?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/agent/main.go | 18 +++++++++-- cmd/playground/main.go | 18 +++++++++-- go.mod | 1 + pkg/utils/proccess_group_darwin.go | 20 ++++++++++++ pkg/utils/proccess_group_linux.go | 20 ++++++++++++ pkg/utils/proccess_group_windows.go | 49 +++++++++++++++++++++++++++++ 6 files changed, 121 insertions(+), 5 deletions(-) create mode 100644 pkg/utils/proccess_group_darwin.go create mode 100644 pkg/utils/proccess_group_linux.go create mode 100644 pkg/utils/proccess_group_windows.go diff --git a/cmd/agent/main.go b/cmd/agent/main.go index 3d57b65..02e3fc7 100644 --- a/cmd/agent/main.go +++ b/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: } diff --git a/cmd/playground/main.go b/cmd/playground/main.go index c1a76b1..56c8c18 100644 --- a/cmd/playground/main.go +++ b/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 { diff --git a/go.mod b/go.mod index 51ea4d3..ece9295 100644 --- a/go.mod +++ b/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 diff --git a/pkg/utils/proccess_group_darwin.go b/pkg/utils/proccess_group_darwin.go new file mode 100644 index 0000000..04029bc --- /dev/null +++ b/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") +} diff --git a/pkg/utils/proccess_group_linux.go b/pkg/utils/proccess_group_linux.go new file mode 100644 index 0000000..04029bc --- /dev/null +++ b/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") +} diff --git a/pkg/utils/proccess_group_windows.go b/pkg/utils/proccess_group_windows.go new file mode 100644 index 0000000..0355ec7 --- /dev/null +++ b/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)) +}