pxmon/internal/agent/top.go

package agent

import (
	"bufio"
	"context"
	"errors"
	"io/fs"
	"os"
	"os/user"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// ProcessStat is one process summary for /api/v1/top.
type ProcessStat struct {
	PID        int     `json:"pid"`
	User       string  `json:"user"`
	Command    string  `json:"command"`
	CPUPercent float64 `json:"cpu_percent"`
	RSSBytes   uint64  `json:"rss_bytes"`
	VSZBytes   uint64  `json:"vsz_bytes"`
	IOReadTot  uint64  `json:"io_read_bytes"`
	IOWriteTot uint64  `json:"io_write_bytes"`
}

// TopResponse is the payload returned by /api/v1/top.
type TopResponse struct {
	Timestamp   time.Time     `json:"timestamp"`
	SampleMs    int64         `json:"sample_ms"`
	TotalProcs  int           `json:"total_procs"`
	TopByCPU    []ProcessStat `json:"top_by_cpu"`
	TopByMemory []ProcessStat `json:"top_by_memory"`
	TopByIO     []ProcessStat `json:"top_by_io"`
}

// CollectTopProcesses takes two samples of /proc/[pid]/stat separated by
// sampleWindow to compute CPU%, then returns the top N processes by CPU, by
// RSS, and by IO total. Linux-only; other OSes return an empty response.
func CollectTopProcesses(sampleWindow time.Duration, limit int) (TopResponse, error) {
	if runtime.GOOS != "linux" {
		return TopResponse{Timestamp: time.Now().UTC()}, nil
	}
	if sampleWindow <= 0 {
		sampleWindow = 250 * time.Millisecond
	}
	if sampleWindow > 2*time.Second {
		sampleWindow = 2 * time.Second
	}
	if limit <= 0 || limit > 200 {
		limit = 20
	}

	hz := clockTicksPerSecond()
	pageSize := uint64(os.Getpagesize())

	first, err := snapshotProcesses()
	if err != nil {
		return TopResponse{}, err
	}
	time.Sleep(sampleWindow)
	second, err := snapshotProcesses()
	if err != nil {
		return TopResponse{}, err
	}

	usernameCache := newUsernameCache()
	elapsedTicks := float64(sampleWindow.Seconds()) * hz
	if elapsedTicks <= 0 {
		elapsedTicks = 1
	}

	merged := make([]ProcessStat, 0, len(second))
	for pid, s2 := range second {
		s1, ok := first[pid]
		cpu := 0.0
		if ok {
			dTicks := float64((s2.utime + s2.stime) - (s1.utime + s1.stime))
			if dTicks > 0 {
				cpu = (dTicks / elapsedTicks) * 100.0
			}
		}
		if cpu < 0 {
			cpu = 0
		}

		merged = append(merged, ProcessStat{
			PID:        pid,
			User:       usernameCache.lookup(s2.uid),
			Command:    s2.command,
			CPUPercent: round2(cpu),
			RSSBytes:   s2.rssPages * pageSize,
			VSZBytes:   s2.vsize,
			IOReadTot:  s2.ioRead,
			IOWriteTot: s2.ioWrite,
		})
	}

	byCPU := topN(merged, limit, func(a, b ProcessStat) bool { return a.CPUPercent > b.CPUPercent })
	byMem := topN(merged, limit, func(a, b ProcessStat) bool { return a.RSSBytes > b.RSSBytes })
	byIO := topN(merged, limit, func(a, b ProcessStat) bool {
		return (a.IOReadTot + a.IOWriteTot) > (b.IOReadTot + b.IOWriteTot)
	})

	return TopResponse{
		Timestamp:   time.Now().UTC(),
		SampleMs:    sampleWindow.Milliseconds(),
		TotalProcs:  len(merged),
		TopByCPU:    byCPU,
		TopByMemory: byMem,
		TopByIO:     byIO,
	}, nil
}

type procSample struct {
	pid      int
	command  string
	utime    uint64
	stime    uint64
	vsize    uint64
	rssPages uint64
	uid      int
	ioRead   uint64
	ioWrite  uint64
}

func snapshotProcesses() (map[int]procSample, error) {
	entries, err := os.ReadDir("/proc")
	if err != nil {
		return nil, err
	}
	out := make(map[int]procSample, 256)
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		pid, err := strconv.Atoi(e.Name())
		if err != nil || pid <= 0 {
			continue
		}
		s, ok := readProcSample(pid)
		if !ok {
			continue
		}
		out[pid] = s
	}
	return out, nil
}

func readProcSample(pid int) (procSample, bool) {
	pidStr := strconv.Itoa(pid)
	statData, err := os.ReadFile("/proc/" + pidStr + "/stat")
	if err != nil {
		return procSample{}, false
	}

	// comm is in parens and may contain spaces; parse after the last ')'.
	line := string(statData)
	rp := strings.LastIndexByte(line, ')')
	if rp <= 0 {
		return procSample{}, false
	}
	lp := strings.IndexByte(line, '(')
	if lp < 0 || lp >= rp {
		return procSample{}, false
	}
	comm := line[lp+1 : rp]
	rest := strings.Fields(line[rp+2:])
	// After comm and the state char, indices inside `rest`:
	// 0: state ... but we split from after ') '. fields[0]=state, fields[1]=ppid,
	// fields[2]=pgrp, ..., fields[11]=utime, fields[12]=stime, fields[19]=vsize, fields[20]=rss.
	if len(rest) < 22 {
		return procSample{}, false
	}
	utime, _ := strconv.ParseUint(rest[11], 10, 64)
	stime, _ := strconv.ParseUint(rest[12], 10, 64)
	vsize, _ := strconv.ParseUint(rest[20], 10, 64)
	rssPages, _ := strconv.ParseUint(rest[21], 10, 64)

	uid := readProcUID(pidStr)
	ioRead, ioWrite := readProcIO(pidStr)

	return procSample{
		pid:      pid,
		command:  comm,
		utime:    utime,
		stime:    stime,
		vsize:    vsize,
		rssPages: rssPages,
		uid:      uid,
		ioRead:   ioRead,
		ioWrite:  ioWrite,
	}, true
}

func readProcUID(pidStr string) int {
	data, err := os.ReadFile("/proc/" + pidStr + "/status")
	if err != nil {
		return -1
	}
	for _, line := range strings.Split(string(data), "\n") {
		if strings.HasPrefix(line, "Uid:") {
			fields := strings.Fields(line)
			if len(fields) >= 2 {
				if v, err := strconv.Atoi(fields[1]); err == nil {
					return v
				}
			}
			return -1
		}
	}
	return -1
}

func readProcIO(pidStr string) (uint64, uint64) {
	data, err := os.ReadFile("/proc/" + pidStr + "/io")
	if err != nil {
		return 0, 0
	}
	var rb, wb uint64
	for _, line := range strings.Split(string(data), "\n") {
		parts := strings.SplitN(line, ":", 2)
		if len(parts) != 2 {
			continue
		}
		key := strings.TrimSpace(parts[0])
		val := strings.TrimSpace(parts[1])
		v, _ := strconv.ParseUint(val, 10, 64)
		switch key {
		case "read_bytes":
			rb = v
		case "write_bytes":
			wb = v
		}
	}
	return rb, wb
}

var cachedClockTicks float64
var cachedClockOnce sync.Once

func clockTicksPerSecond() float64 {
	cachedClockOnce.Do(func() {
		// SC_CLK_TCK is almost always 100 on Linux; read from /proc/self/stat
		// + uptime as a sanity check would be nicer but we accept the default.
		cachedClockTicks = 100.0
	})
	return cachedClockTicks
}

type usernameCache struct {
	cache map[int]string
}

func newUsernameCache() *usernameCache {
	return &usernameCache{cache: make(map[int]string, 16)}
}

func (c *usernameCache) lookup(uid int) string {
	if uid < 0 {
		return "-"
	}
	if v, ok := c.cache[uid]; ok {
		return v
	}
	u, err := user.LookupId(strconv.Itoa(uid))
	if err != nil || u == nil {
		name := strconv.Itoa(uid)
		c.cache[uid] = name
		return name
	}
	c.cache[uid] = u.Username
	return u.Username
}

func topN(in []ProcessStat, n int, less func(a, b ProcessStat) bool) []ProcessStat {
	cp := make([]ProcessStat, len(in))
	copy(cp, in)
	sort.Slice(cp, func(i, j int) bool { return less(cp[i], cp[j]) })
	if len(cp) > n {
		cp = cp[:n]
	}
	return cp
}

// DirStat is one directory entry in a du-style listing.
type DirStat struct {
	Path  string `json:"path"`
	Name  string `json:"name"`
	Bytes uint64 `json:"bytes"`
	Files uint64 `json:"files"`
}

// DUResponse is returned by /api/v1/du.
type DUResponse struct {
	Timestamp time.Time `json:"timestamp"`
	Root      string    `json:"root"`
	Entries   []DirStat `json:"entries"`
	Truncated bool      `json:"truncated"`
}

// CollectDirSizes lists immediate children of root (only directories) and
// sums file sizes under each child recursively, respecting ctx deadline.
// Returned entries are sorted by Bytes descending. Truncated is true if the
// walker stopped early due to timeout.
func CollectDirSizes(ctx context.Context, root string, limit int) (DUResponse, error) {
	root = strings.TrimSpace(root)
	if root == "" {
		root = "/"
	}
	absRoot, err := filepath.Abs(root)
	if err != nil {
		return DUResponse{}, err
	}
	info, err := os.Stat(absRoot)
	if err != nil {
		return DUResponse{}, err
	}
	if !info.IsDir() {
		return DUResponse{}, errors.New("root is not a directory")
	}
	if limit <= 0 || limit > 100 {
		limit = 15
	}

	entries, err := os.ReadDir(absRoot)
	if err != nil {
		return DUResponse{}, err
	}

	truncated := false
	out := make([]DirStat, 0, len(entries))
	for _, e := range entries {
		if ctx.Err() != nil {
			truncated = true
			break
		}
		if !e.IsDir() {
			continue
		}
		name := e.Name()
		// Skip virtual filesystems when rooted at /.
		if absRoot == "/" && isSkippedSystemDir(name) {
			continue
		}
		full := filepath.Join(absRoot, name)
		size, files, stopped := sumDirectory(ctx, full)
		if stopped {
			truncated = true
		}
		out = append(out, DirStat{
			Path:  full,
			Name:  name,
			Bytes: size,
			Files: files,
		})
		if stopped {
			break
		}
	}

	sort.Slice(out, func(i, j int) bool { return out[i].Bytes > out[j].Bytes })
	if len(out) > limit {
		out = out[:limit]
	}

	return DUResponse{
		Timestamp: time.Now().UTC(),
		Root:      absRoot,
		Entries:   out,
		Truncated: truncated,
	}, nil
}

func isSkippedSystemDir(name string) bool {
	switch name {
	case "proc", "sys", "dev", "run", "tmp":
		return true
	}
	return false
}

func sumDirectory(ctx context.Context, path string) (uint64, uint64, bool) {
	var total uint64
	var files uint64
	stopped := false

	walkFn := func(p string, d fs.DirEntry, err error) error {
		if err != nil {
			if d != nil && d.IsDir() {
				return filepath.SkipDir
			}
			return nil
		}
		if ctx.Err() != nil {
			stopped = true
			return filepath.SkipAll
		}
		if d.IsDir() {
			// Skip known pseudo filesystems we may cross into.
			name := d.Name()
			if p != path && (name == "proc" || name == "sys" || name == "dev") {
				return filepath.SkipDir
			}
			return nil
		}
		info, err := d.Info()
		if err != nil {
			return nil
		}
		total += uint64(info.Size())
		files++
		return nil
	}
	_ = filepath.WalkDir(path, walkFn)
	return total, files, stopped
}

// scanLinesToFields is a tiny helper for tests and future use.
func scanLinesToFields(data []byte) [][]string {
	var out [][]string
	sc := bufio.NewScanner(strings.NewReader(string(data)))
	for sc.Scan() {
		out = append(out, strings.Fields(sc.Text()))
	}
	return out
}