use task childs instead of pgids to track build process

proto_package.go

@@ -225,7 +225,6 @@ func (p *ProtoPackage) build(ctx context.Context) (time.Duration, error) {
 	go pollMemoryUsage(pgid, 1*time.Second, done, &peakMem)
 
 	err = cmd.Wait()
-
 	close(done)
 
 	Rusage, ok := cmd.ProcessState.SysUsage().(*syscall.Rusage)

utils.go

@@ -768,6 +768,7 @@ func downloadSRCINFO(pkg, tag string) (*srcinfo.Srcinfo, error) {
 	return nSrcInfo, nil
 }
 
+// getProcessMemory reads both RSS and Swap memory from /proc/<pid>/status
 func getProcessMemory(pid int) (int, int, error) {
 	statusPath := fmt.Sprintf("/proc/%d/status", pid)
 	file, err := os.Open(statusPath)
@@ -803,58 +804,55 @@ func getProcessMemory(pid int) (int, int, error) {
 	return rss, swap, nil
 }
 
-func getProcessTreeMemory(pgid int) (int, int, error) {
+func getChildProcesses(pid int) ([]int, error) {
-	procDir, err := os.Open("/proc")
+	var children []int
+	taskPath := fmt.Sprintf("/proc/%d/task", pid)
+
+	taskDirs, err := os.ReadDir(taskPath)
 	if err != nil {
-		return 0, 0, err
+		return nil, err
-	}
-	defer func(procDir *os.File) {
-		_ = procDir.Close()
-	}(procDir)
-
-	totalRSS, totalSwap := 0, 0
-	entries, err := procDir.Readdir(-1)
-	if err != nil {
-		return 0, 0, fmt.Errorf("failed to read /proc: %v", err)
 	}
 
-	for _, entry := range entries {
+	for _, task := range taskDirs {
-		if !entry.IsDir() {
+		childFile := fmt.Sprintf("/proc/%d/task/%s/children", pid, task.Name())
-			continue
+		file, err := os.Open(childFile)
-		}
-
-		pid, err := strconv.Atoi(entry.Name())
-		if err != nil {
-			continue
-		}
-
-		statPath := fmt.Sprintf("/proc/%d/stat", pid)
-		file, err := os.Open(statPath)
 		if err != nil {
 			continue
 		}
 
 		scanner := bufio.NewScanner(file)
-		if scanner.Scan() {
+		for scanner.Scan() {
-			fields := strings.Fields(scanner.Text())
+			childPIDs := strings.Fields(scanner.Text())
-			if len(fields) >= 5 {
+			for _, childStr := range childPIDs {
-				processPGID, err := strconv.Atoi(fields[4])
+				childPID, err := strconv.Atoi(childStr)
-				if err != nil {
-					return 0, 0, fmt.Errorf("failed to parse process PG ID: %v", err)
-				}
-				if processPGID == pgid {
-					rss, swap, err := getProcessMemory(pid)
 				if err == nil {
-						totalRSS += rss
+					children = append(children, childPID)
-						totalSwap += swap
+					subChildren, _ := getChildProcesses(childPID)
-					} else {
+					children = append(children, subChildren...)
-						return 0, 0, fmt.Errorf("failed to get process memory: %v", err)
-					}
 				}
 			}
 		}
 		_ = file.Close()
 	}
+	return children, nil
+}
+
+func getProcessTreeMemory(rootPID int) (int, int, error) {
+	pids := []int{rootPID}
+
+	childPIDs, err := getChildProcesses(rootPID)
+	if err == nil {
+		pids = append(pids, childPIDs...)
+	}
+
+	totalRSS, totalSwap := 0, 0
+	for _, pid := range pids {
+		rss, swap, err := getProcessMemory(pid)
+		if err == nil {
+			totalRSS += rss
+			totalSwap += swap
+		}
+	}
 
 	return totalRSS, totalSwap, nil
 }