pid.go

package utils

import (
	"fmt"
	"log"
	"os"
	"regexp"
	"strconv"
	"syscall"
	"time"

	"github.com/shirou/gopsutil/v4/process"
)

var (
	RePidNS = regexp.MustCompile(`pid:\[([1-9][0-9]*)]$`) // 匹配pid命名空间的正则表达式 pid:[4026545939]
)

// GetPidNS 获取指定进程的Pid命名空间号
func GetPidNS(pid uint64) (uint64, error) {
	str, err := os.Readlink(fmt.Sprintf("/proc/%d/ns/pid", pid))
	if err != nil {
		return 0, err
	}
	if !RePidNS.MatchString(str) {
		return 0, fmt.Errorf("error matching pid")
	}
	strs := RePidNS.FindStringSubmatch(str)
	if len(strs) < 2 {
		return 0, fmt.Errorf("error matching pid")
	}
	return strconv.ParseUint(strs[1], 10, 64)
}

// GetProcessInfo 获取进程信息
func GetProcessInfo(pid int32) {
	p, err := process.NewProcess(pid)
	if err != nil {
		return
	}
	p.Times()
	cmdStr, err := p.Cmdline()
	if err != nil {
		return
	}
	log.Println(cmdStr)
}

func GetProcessByName(cmdline string) ([]*process.Process, error) {
	p, err := process.Processes()
	if err != nil {
		return nil, err
	}
	result := make([]*process.Process, 0)
	for _, i := range p {
		c, innerErr := i.CmdlineSlice()
		if innerErr != nil || len(c) <= 0 {
			continue
		}
		if c[0] == cmdline {
			result = append(result, i)
		}
	}
	return result, nil
}

// GetProcessCPUUsage 获取进程的CPU使用率
func GetProcessCPUUsage(pid int32) (float64, error) {
	p, err := process.NewProcess(pid)
	if err != nil {
		return 0, err
	}
	return p.CPUPercent()
}

type ProcessInfo struct {
	Ppid  int32
	Pid   int32
	User  string
	Uid   uint32
	CPU   float64
	Mem   float32
	Time  time.Duration
	Cmd   string
	Child map[int32]*ProcessInfo
}

func NewProcessInfo(p *process.Process) *ProcessInfo {
	if p == nil {
		return nil
	}
	result := &ProcessInfo{
		Pid:   p.Pid,
		Child: make(map[int32]*ProcessInfo),
	}
	uids, err := p.Uids()
	if err == nil && len(uids) > 0 {
		result.Uid = uids[0]
	}
	user, err := p.Username()
	if err == nil {
		result.User = user
	}
	cpu, err := p.CPUPercent()
	if err == nil {
		result.CPU = cpu
	}
	mem, err := p.MemoryPercent()
	if err == nil {
		result.Mem = mem
	}
	times, err := p.Times()
	if err == nil {
		result.Time = (time.Duration((times.System + times.User)) * time.Second)
	}
	cmd, err := p.Cmdline()
	if err == nil {
		result.Cmd = cmd
	}

	ppid, err := p.Ppid()
	if err == nil {
		result.Ppid = ppid
	}
	return result
}

func GetPsTree(pids []int32) (map[int32]*ProcessInfo, *ProcessInfo) {
	findedProcess := make(map[int32]*ProcessInfo)
	for _, pid := range pids {
		pp := pid
		plist := make([]*ProcessInfo, 0, 16)
		// 向上追踪
		for {
			p, have := findedProcess[pp]
			if !have {
				pInner, err := process.NewProcess(pp)
				if err != nil {
					break
				}
				p = NewProcessInfo(pInner)
				findedProcess[pInner.Pid] = p
			}
			plist = append(plist, p)
			if p.Pid == 1 {
				break
			}
			pp = p.Ppid
		}
		plistLen := len(plist)
		if plistLen >= 2 {
			for i := 0; i <= plistLen-2; i++ {
				p := plist[i+1]
				p.Child[plist[i].Pid] = plist[i]
			}
		}
		// addChilds(pidProcess, findedProcess)
	}
	return findedProcess, findedProcess[1]
}

func GetPsTree2(pids []int32) map[int32]*ProcessInfo {
	allProcess, err := process.Processes()
	if err != nil {
		return nil
	}
	m := make(map[int32]*process.Process)
	for _, v := range allProcess {
		m[v.Pid] = v
	}
	findedProcess := make(map[int32]*ProcessInfo)
	for _, pid := range pids {
		pp := pid
		plist := make([]*ProcessInfo, 0, 16)
		for {
			p, have := findedProcess[pp]
			if !have {
				pInner, have := m[pp]
				if !have {
					break
				}
				p = NewProcessInfo(pInner)
				findedProcess[pInner.Pid] = p
			}
			plist = append(plist, p)
			if p.Pid == 1 {
				break
			}
			pp = p.Ppid
		}
		plistLen := len(plist)
		if plistLen >= 2 {
			for i := 0; i <= plistLen-2; i++ {
				p := plist[i+1]
				p.Child[plist[i].Pid] = plist[i]
			}
		}
	}
	return findedProcess
}

func SendSignal(pids []int32, sig syscall.Signal) {
	if len(pids) == 0 {
		return
	}
	for _, v := range pids {
		if v == 1 {
			continue
		}
		p, err := process.NewProcess(v)
		if err != nil {
			continue
		}
		_ = p.SendSignal(sig)
	}
}