fix 优化部分逻辑

cmd/hytop/backend/main.go

@@ -22,19 +22,24 @@ import (
 
 var (
 	MapIdDCU      = sync.Map{} // 记录dcu信息
+	ContainerInfo *docker.ContainersInfo
 	DockerPidInfo *DockerProcessMap = nil
 	User                            = ""
 	HostName                        = ""
 )
 
 func init() {
+	ContainerInfo = docker.NewContainersInfo()
+	ContainerInfo.Update()
 	i, err := docker.ContainerInfo.GetProcessIdInDocker(false)
-	if err != nil || i == nil {
 	DockerPidInfo = &DockerProcessMap{lock: sync.RWMutex{}, pids: make(map[int32]bool)}
-		return
+	if err == nil && i != nil {
+		for _, v := range i {
+			for _, pidInfo := range v {
+				DockerPidInfo.pids[pidInfo.Pid] = true
+			}
+		}
 	}
-	DockerPidInfo = &DockerProcessMap{lock: sync.RWMutex{}, pids: maps.Clone(i)}
-
 	HostName, _ = os.Hostname()
 	uid := os.Getuid()
 	u, err := utils.GetSysUserById(uid)
@@ -48,14 +53,14 @@ func init() {
 type DCUInfo struct {
 	Id               int
 	Name             string // full
-	PerformanceLevel string // full
+	PerformanceLevel string
 	Fan              string // full
 	Temp             float32
 	PwrAvg           float32
-	PwrCap           float32 // full
+	PwrCap           float32
 	BusId            string // full
-	MemTotal         int
-	MemUsed          int
+	MemTotal         int    // full
+	MemUsed          int    // full
 	MemUsedPerent    float32
 	Mig              bool // full
 	DCUUTil          float32
@@ -212,15 +217,21 @@ func (dpm *DockerProcessMap) GetPidInfo() map[int32]bool {
 	return maps.Clone(dpm.pids)
 }
 
-// Update 重新获取数据，这是一个耗时的操作
-func (dpm *DockerProcessMap) Update() map[int32]bool {
-	i, err := docker.ContainerInfo.GetProcessIdInDocker(true)
+func (dpm *DockerProcessMap) Update(dinfo *docker.ContainersInfo) (map[int32]bool, sync.Locker) {
+	dpm.lock.Lock()
+	clear(dpm.pids)
+	dpm.lock.Unlock()
+	rl := dpm.lock.RLocker()
+	i, err := dinfo.GetProcessIdInDocker(false)
 	if err != nil || i == nil {
-		dpm.pids = make(map[int32]bool)
-		return make(map[int32]bool)
+		return dpm.pids, rl
+	}
+	for _, v := range i {
+		for _, pidInfo := range v {
+			DockerPidInfo.pids[pidInfo.Pid] = true
+		}
 	}
-	dpm.pids = maps.Clone(i)
-	return maps.Clone(i)
+	return dpm.pids, rl
 }
 
 type DCUProcessInfo struct {

docker/find.go

 package docker
 
 import (
-	"get-container/utils"
 	"strconv"
 
 	"context"
 	"errors"
-	"fmt"
-	"os"
 	"regexp"
 	"strings"
 	"sync"
@@ -17,36 +14,32 @@ import (
 	"github.com/moby/moby/client"
 )
 
-func init() {
-	_ = initContainerInfo()
-}
-
-/**
-有两种方法获取进程属于哪个容器
-1. 通过查询pid命名空间，仅在没有指定--pid参数时有效
-2. 通过查询进程的cgroup
-3. 使用docker top <container-id>匹配
-*/
-
-type FindCIDMethod string
-
-const (
-	ByCgroup FindCIDMethod = "byCGroup"
-	ByPidNS  FindCIDMethod = "byPidNS"
-	ByTop    FindCIDMethod = "byTop"
-)
-
 var (
 	ReDocker                      = regexp.MustCompile(`^.*docker[-/]([0-9a-z]*)(?:|.*)`)
 	ContainerInfo *ContainersInfo = nil
 )
 
+func NewContainersInfo() *ContainersInfo {
+	ContainerInfo = &ContainersInfo{
+		inspectInfo: make(map[string]container.InspectResponse),
+		listInfo:    make(map[string]container.Summary),
+		topInfo:     make(map[string]container.TopResponse),
+		time:        time.Now(),
+		inspectLock: sync.RWMutex{},
+		topLock:     sync.RWMutex{},
+		listLock:    sync.RWMutex{},
+	}
+	return ContainerInfo
+}
+
 type ContainersInfo struct {
-	lock        sync.RWMutex // 读写锁，防止对Info的并发写
 	time        time.Time // 记录写入Info的时间
 	inspectInfo map[string]container.InspectResponse
+	inspectLock sync.RWMutex
 	listInfo    map[string]container.Summary
+	listLock    sync.RWMutex
 	topInfo     map[string]container.TopResponse
+	topLock     sync.RWMutex
 }
 
 type ContainerPsInfo struct {
@@ -113,9 +106,15 @@ func ParsePsInfo(topInfo map[string]container.TopResponse) (map[string][]Contain
 }
 
 func (info *ContainersInfo) Update() error {
-	info.lock.Lock()
-	defer info.lock.Unlock()
-	i, s, t, err := getContainerInfo()
+	info.listLock.Lock()
+	info.topLock.Lock()
+	info.inspectLock.Lock()
+	defer func() {
+		info.inspectLock.Unlock()
+		info.topLock.Unlock()
+		info.listLock.Unlock()
+	}()
+	i, s, t, err := getRunningContainerInfo()
 	if err != nil {
 		return err
 	}
@@ -126,201 +125,39 @@ func (info *ContainersInfo) Update() error {
 	return nil
 }
 
-func (info *ContainersInfo) Get() (map[string]container.InspectResponse, sync.Locker) {
-	rl := info.lock.RLocker()
-	rl.Lock()
-	return info.inspectInfo, rl
-}
-
-func initContainerInfo() error {
-	inspect, lists, tops, err := getContainerInfo()
-	if err != nil {
-		return err
+func (info *ContainersInfo) GetInspectInfo(update bool) (map[string]container.InspectResponse, sync.Locker) {
+	if update {
+		info.Update()
 	}
-	ContainerInfo = &ContainersInfo{
-		lock:        sync.RWMutex{},
-		time:        time.Now(),
-		inspectInfo: inspect,
-		listInfo:    lists,
-		topInfo:     tops,
+	rl := info.inspectLock.RLocker()
+	rl.Lock()
+	if info.inspectInfo == nil {
+		return make(map[string]container.InspectResponse), rl
 	}
-	return nil
+	return info.inspectInfo, rl
 }
 
 // GetProcessIdInDocker 获取所用容器的进程信息
-func (info *ContainersInfo) GetProcessIdInDocker(update bool) (map[int32]bool, error) {
-	result := make(map[int32]bool)
+func (info *ContainersInfo) GetProcessIdInDocker(update bool) (map[string][]ContainerPsInfo, error) {
 	if update {
 		err := info.Update()
 		if err != nil {
-			return result, err
+			return nil, err
 		}
 	}
-	rl := info.lock.RLocker()
+	rl := info.topLock.RLocker()
 	rl.Lock()
 	i, err := ParsePsInfo(info.topInfo)
 	rl.Unlock()
 	rl = nil
-	if err != nil {
-		return result, err
-	}
-	for _, v := range i {
-		for _, k := range v {
-			result[k.Pid] = true
-		}
-	}
-	return result, nil
-}
-
-// FindContainerIdByPid 根据pid获取该进程属于哪个docker容器，返回容器id，如果为nil，表示找不到容器id
-func FindContainerIdByPid(pid uint64, method FindCIDMethod) (*string, error) {
-	switch method {
-	case ByPidNS:
-		return findContainerIdByNS(pid)
-	case ByCgroup:
-		return findContainerIdByCgroup(pid)
-	default:
-		return nil, fmt.Errorf("unknown method: %s", method)
-	}
-}
-
-func FindContainerIdByPidBatch(pids []uint64, method FindCIDMethod) (map[uint64]string, error) {
-	if len(pids) == 0 {
-		return nil, nil
-	}
-	switch method {
-	case ByPidNS:
-		return findContainerIdByNSBatch(pids)
-	case ByCgroup:
-		return findContainerIdByCgroupBatch(pids)
-	default:
-		return nil, fmt.Errorf("unknown method: %s", method)
-	}
-}
-
-// findContainerIdByPidCgroup 通过cgroup查询docker容器id
-func findContainerIdByCgroup(pid uint64) (*string, error) {
-	content, err := os.ReadFile(fmt.Sprintf("/proc/%d/cgroup", pid))
 	if err != nil {
 		return nil, err
 	}
-	contentStr := strings.Trim(string(content), "\n")
-	if len(contentStr) == 0 {
-		return nil, errors.New("process's cgroup not found")
-	}
-	lines := strings.Split(contentStr, "\n")
-	var target string
-	if len(lines) > 1 {
-		// 如果有多行，解析有pids的行
-		for _, line := range lines {
-			if strings.Contains(line, "pids") {
-				target = strings.TrimSpace(line)
-				break
-			}
-		}
-		if target == "" {
-			return nil, errors.New("process's cgroup not found pids line")
-		}
-	} else {
-		// 如果是单行，直接解析
-		target = strings.TrimSpace(lines[0])
-	}
-	target = strings.TrimSpace(target)
-	if !strings.Contains(target, "docker") {
-		return nil, errors.New("process's cgroup is not create by docker")
-	}
-	if ReDocker.MatchString(target) {
-		fields := ReDocker.FindStringSubmatch(target)
-		if len(fields) < 2 {
-			return nil, errors.New("process's cgroup is not create by docker")
-		}
-		cid := fields[1]
-		return &cid, nil
-	} else {
-		return nil, errors.New("process's cgroup is not create by docker")
-	}
-}
-
-func findContainerIdByCgroupBatch(pids []uint64) (map[uint64]string, error) {
-	results := make(map[uint64]string)
-	for _, pid := range pids {
-		str, err := findContainerIdByCgroup(pid)
-		if err != nil {
-			return nil, err
-		}
-		s := *str
-		results[pid] = s
-	}
-	return results, nil
-}
-
-// findContainerIdByNS 通过pid命名空间查询docker容器id
-func findContainerIdByNS(pid uint64) (*string, error) {
-	ns, err := utils.GetPidNS(pid)
-	if err != nil {
-		return nil, err
-	}
-	if ContainerInfo == nil {
-		innerErr := initContainerInfo()
-		if innerErr != nil {
-			return nil, innerErr
-		}
-	} else {
-		if innerErr := ContainerInfo.Update(); innerErr != nil {
-			return nil, innerErr
-		}
-	}
-	info, lock := ContainerInfo.Get()
-	defer lock.Unlock()
-	for k, v := range info {
-		containerNs, innerErr := utils.GetPidNS(uint64(v.State.Pid))
-		if innerErr != nil {
-			continue
-		}
-		if containerNs == ns {
-			cid := k
-			return &cid, nil
-		}
-	}
-	return nil, nil
-}
-
-func findContainerIdByNSBatch(pids []uint64) (map[uint64]string, error) {
-	if ContainerInfo == nil {
-		innerErr := initContainerInfo()
-		if innerErr != nil {
-			return nil, innerErr
-		}
-	} else {
-		if innerErr := ContainerInfo.Update(); innerErr != nil {
-			return nil, innerErr
-		}
-	}
-	info, lock := ContainerInfo.Get()
-	defer lock.Unlock()
-	results := make(map[uint64]string)
-	ns2cid := make(map[uint64]string)
-	for k, v := range info {
-		containerNs, innerErr := utils.GetPidNS(uint64(v.State.Pid))
-		if innerErr != nil {
-			return nil, innerErr
-		}
-		ns2cid[containerNs] = k
-	}
-	for _, pid := range pids {
-		ns, err := utils.GetPidNS(pid)
-		if err != nil {
-			continue
-		}
-		if cid, ok := ns2cid[ns]; ok {
-			results[pid] = cid
-		}
-	}
-	return results, nil
+	return i, nil
 }
 
-// getContainerInfo 获取所有正在运行的docker容器的详细信息
-func getContainerInfo() (map[string]container.InspectResponse, map[string]container.Summary, map[string]container.TopResponse, error) {
+// getRunningContainerInfo 获取所有正在运行的docker容器的详细信息
+func getRunningContainerInfo() (map[string]container.InspectResponse, map[string]container.Summary, map[string]container.TopResponse, error) {
 	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
 	if err != nil {
 		return nil, nil, nil, err

docker/find_test.go

@@ -91,7 +91,7 @@ func TestDocker(t *testing.T) {
 
 func TestGetProcessIdInDocker(t *testing.T) {
 	now := time.Now()
-	err := initContainerInfo()
+	err := ContainerInfo.Update()
 	if err != nil {
 		t.Error(err)
 	}
@@ -100,6 +100,11 @@ func TestGetProcessIdInDocker(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
-	t.Logf("%d", d.Nanoseconds())
-	t.Logf("%+v", pids)
+	t.Logf("%d ms", d.Milliseconds())
+	for k, v := range pids {
+		t.Logf("======> container %s has %d processes", k, len(v))
+		for _, pidInfo := range v {
+			t.Logf("  %+v", pidInfo)
+		}
+	}
 }