68 lines
1.4 KiB
Go
68 lines
1.4 KiB
Go
package procgroup
|
|
|
|
import (
|
|
"context"
|
|
"os/exec"
|
|
"strconv"
|
|
"strings"
|
|
"syscall"
|
|
)
|
|
|
|
type Cmd struct {
|
|
*exec.Cmd
|
|
}
|
|
|
|
func Command(name string, args ...string) *Cmd {
|
|
return group(exec.Command(name, args...))
|
|
}
|
|
|
|
func CommandContext(ctx context.Context, name string, args ...string) *Cmd {
|
|
cmd := group(exec.CommandContext(ctx, name, args...))
|
|
cmd.Cancel = cmd.Kill
|
|
return cmd
|
|
}
|
|
|
|
func group(cmd *exec.Cmd) *Cmd {
|
|
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
|
return &Cmd{Cmd: cmd}
|
|
}
|
|
|
|
// Kill also kills descendants that left the group: agents run their commands in sessions of their own.
|
|
func (cmd *Cmd) Kill() error {
|
|
pids := descendants(cmd.Process.Pid)
|
|
err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
|
|
for _, pid := range pids {
|
|
syscall.Kill(pid, syscall.SIGKILL)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func descendants(root int) []int {
|
|
out, err := exec.Command("ps", "-A", "-o", "pid=,ppid=").Output()
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
children := map[int][]int{}
|
|
for _, line := range strings.Split(string(out), "\n") {
|
|
f := strings.Fields(line)
|
|
if len(f) != 2 {
|
|
continue
|
|
}
|
|
pid, err1 := strconv.Atoi(f[0])
|
|
ppid, err2 := strconv.Atoi(f[1])
|
|
if err1 != nil || err2 != nil {
|
|
continue
|
|
}
|
|
children[ppid] = append(children[ppid], pid)
|
|
}
|
|
pids := []int{}
|
|
queue := children[root]
|
|
for len(queue) > 0 {
|
|
pid := queue[0]
|
|
queue = queue[1:]
|
|
pids = append(pids, pid)
|
|
queue = append(queue, children[pid]...)
|
|
}
|
|
return pids
|
|
}
|