Skip to content

Commit

Permalink
feature: replace ps command to gopsutil
Browse files Browse the repository at this point in the history 
Signed-off-by: xcaspar <[email protected]>
  • Loading branch information
@xcaspar
xcaspar committed Jun 28, 2020
1 parent 0003add commit fa9dc1f
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 80 deletions.
186 changes: 106 additions & 80 deletions channel/local.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"strings"
"time"

"github.com/shirou/gopsutil/process"
"github.com/sirupsen/logrus"

"github.com/chaosblade-io/chaosblade-spec-go/spec"
Expand All @@ -50,78 +51,116 @@ func (l *LocalChannel) GetScriptPath() string {
}

func (l *LocalChannel) GetPidsByProcessCmdName(processName string, ctx context.Context) ([]string, error) {
excludeProcesses := ctx.Value(ExcludeProcessKey)
excludeGrepInfo := ""
if excludeProcesses != nil {
excludeProcessesString := excludeProcesses.(string)
excludeProcessArrays := strings.Split(excludeProcessesString, ",")
for _, excludeProcess := range excludeProcessArrays {
if excludeProcess != "" {
excludeGrepInfo += fmt.Sprintf(`| grep -v -w %s`, excludeProcess)
}
}
processName = strings.TrimSpace(processName)
if processName == "" {
return []string{}, fmt.Errorf("processName is blank")
}
response := l.Run(ctx, "pgrep",
fmt.Sprintf(`-l %s %s | grep -v -w chaos_killprocess | grep -v -w chaos_stopprocess | awk '{print $1}' | tr '\n' ' '`,
processName, excludeGrepInfo))
if !response.Success {
return nil, fmt.Errorf(response.Err)
processes, err := process.Processes()
if err != nil {
return []string{}, err
}
pidString := response.Result.(string)
pids := strings.Fields(strings.TrimSpace(pidString))
currPid := strconv.Itoa(os.Getpid())
for idx, pid := range pids {
if pid == currPid {
return util.Remove(pids, idx), nil
excludeProcesses := getExcludeProcesses(ctx)
pids := make([]string, 0)
for _, p := range processes {
name, err := p.Name()
if err != nil {
logrus.WithField("pid", p.Pid).WithError(err).Debugln("get process name error")
continue
}
if processName != name {
continue
}
if int32(os.Getpid()) == p.Pid {
continue
}
cmdline, _ := p.Cmdline()
containsExcludeProcess := false
logrus.WithFields(logrus.Fields{
"name": name,
"cmdline": cmdline,
"processName": processName,
}).Debugln("process info")
for _, ep := range excludeProcesses {
if strings.Contains(cmdline, strings.TrimSpace(ep)) {
containsExcludeProcess = true
break
}
}
if containsExcludeProcess {
continue
}
pids = append(pids, fmt.Sprintf("%d", p.Pid))
}
return pids, nil
}

func (l *LocalChannel) GetPidsByProcessName(processName string, ctx context.Context) ([]string, error) {
psArgs := l.GetPsArgs()
otherProcess := ctx.Value(ProcessKey)
otherGrepInfo := ""
if otherProcess != nil {
processString := otherProcess.(string)
if processString != "" {
otherGrepInfo = fmt.Sprintf(`| grep "%s"`, processString)
}
}
excludeProcesses := ctx.Value(ExcludeProcessKey)
excludeGrepInfo := ""
if excludeProcesses != nil {
excludeProcessesString := excludeProcesses.(string)
excludeProcessArrays := strings.Split(excludeProcessesString, ",")
for _, excludeProcess := range excludeProcessArrays {
if excludeProcess != "" {
excludeGrepInfo += fmt.Sprintf(`| grep -v -w %s`, excludeProcess)
}
}
processName = strings.TrimSpace(processName)
if processName == "" {
return []string{}, fmt.Errorf("process keyword is blank")
}
if strings.HasPrefix(processName, "-") {
processName = fmt.Sprintf(`\%s`, processName)
}
response := l.Run(ctx, "ps",
fmt.Sprintf(`%s | grep "%s" %s %s | grep -v -w grep | grep -v -w chaos_killprocess | grep -v -w chaos_stopprocess | awk '{print $2}' | tr '\n' ' '`,
psArgs, processName, otherGrepInfo, excludeGrepInfo))
if !response.Success {
return nil, fmt.Errorf(response.Err)
processes, err := process.Processes()
if err != nil {
return []string{}, err
}
pidString := strings.TrimSpace(response.Result.(string))
if pidString == "" {
return make([]string, 0), nil
otherConditionProcessValue := ctx.Value(ProcessKey)
otherConditionProcessName := ""
if otherConditionProcessValue != nil {
otherConditionProcessName = otherConditionProcessValue.(string)
}
pids := strings.Fields(pidString)
currPid := strconv.Itoa(os.Getpid())
for idx, pid := range pids {
if pid == currPid {
return util.Remove(pids, idx), nil
excludeProcesses := getExcludeProcesses(ctx)
pids := make([]string, 0)
for _, p := range processes {
cmdline, err := p.Cmdline()
if err != nil {
logrus.WithField("pid", p.Pid).WithError(err).Debugln("get command line err")
continue
}
if !strings.Contains(cmdline, processName) {
continue
}
logrus.WithFields(logrus.Fields{
"cmdline": cmdline,
"processName": processName,
"otherConditionProcessName": otherConditionProcessName,
"excludeProcesses": excludeProcesses,
}).Debugln("process info")
if otherConditionProcessName != "" && !strings.Contains(cmdline, otherConditionProcessName) {
continue
}
containsExcludeProcess := false
for _, ep := range excludeProcesses {
if strings.Contains(cmdline, ep) {
containsExcludeProcess = true
break
}
}
if containsExcludeProcess {
continue
}
pids = append(pids, fmt.Sprintf("%d", p.Pid))
}
return pids, nil
}

func getExcludeProcesses(ctx context.Context) []string {
excludeProcessValue := ctx.Value(ExcludeProcessKey)
excludeProcesses := make([]string, 0)
if excludeProcessValue != nil {
excludeProcessesString := excludeProcessValue.(string)
processNames := strings.Split(excludeProcessesString, ",")
for _, name := range processNames {
name = strings.TrimSpace(name)
if name == "" {
continue
}
excludeProcesses = append(excludeProcesses, name)
}
}
excludeProcesses = append(excludeProcesses, "chaos_killprocess", "chaos_stopprocess")
return excludeProcesses
}

func (l *LocalChannel) GetPsArgs() string {
var psArgs = "-eo user,pid,ppid,args"
if l.isAlpinePlatform() {
Expand All @@ -147,36 +186,23 @@ func (l *LocalChannel) IsCommandAvailable(commandName string) bool {
}

func (l *LocalChannel) ProcessExists(pid string) (bool, error) {
if l.isAlpinePlatform() {
response := l.Run(context.TODO(), "ps", fmt.Sprintf("-o pid | grep %s", pid))
if !response.Success {
return false, fmt.Errorf(response.Err)
}
if strings.TrimSpace(response.Result.(string)) == "" {
return false, nil
}
return true, nil
p, err := strconv.Atoi(pid)
if err != nil {
return false, err
}
response := l.Run(context.TODO(), "ps", fmt.Sprintf("-p %s", pid))
return response.Success, nil
return process.PidExists(int32(p))
}

func (l *LocalChannel) GetPidUser(pid string) (string, error) {
var response *spec.Response
if l.isAlpinePlatform() {
response = l.Run(context.TODO(), "ps", fmt.Sprintf("-o user,pid | grep %s", pid))

} else {
response = l.Run(context.TODO(), "ps", fmt.Sprintf("-o user,pid -p %s | grep %s", pid, pid))
}
if !response.Success {
return "", fmt.Errorf(response.Err)
p, err := strconv.Atoi(pid)
if err != nil {
return "", err
}
result := strings.TrimSpace(response.Result.(string))
if result == "" {
return "", fmt.Errorf("process user not found by pid")
process, err := process.NewProcess(int32(p))
if err != nil {
return "", err
}
return strings.Fields(result)[0], nil
return process.Username()
}

func (l *LocalChannel) GetPidsByLocalPorts(localPorts []string) ([]string, error) {
Expand Down
1 change: 1 addition & 0 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ go 1.13

require (
github.com/BurntSushi/toml v0.3.1 // indirect
github.com/shirou/gopsutil v2.20.5+incompatible
github.com/sirupsen/logrus v1.4.2
gopkg.in/natefinch/lumberjack.v2 v2.0.0
gopkg.in/yaml.v2 v2.2.4
Expand Down
2 changes: 2 additions & 0 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@ github.com/pmezard/go-difflib v0.0.0-20151028094244-d8ed2627bdf0/go.mod h1:iKH77
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
github.com/shirou/gopsutil v2.20.5+incompatible h1:tYH07UPoQt0OCQdgWWMgYHy3/a9bcxNpBIysykNIP7I=
github.com/shirou/gopsutil v2.20.5+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
github.com/sirupsen/logrus v1.4.2 h1:SPIRibHv4MatM3XXNO2BJeFLZwZ2LvZgfQ5+UNI2im4=
github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
github.com/spf13/pflag v0.0.0-20170130214245-9ff6c6923cff/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4=
Expand Down

0 comments on commit fa9dc1f

Please sign in to comment.