talm

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package miniprocfs
6

7
import (
8
	"bytes"
9
	"errors"
10
	"io"
11
	"os"
12
	"strconv"
13

14
	"github.com/siderolabs/talos/pkg/machinery/api/machine"
15
)
16

17
const (
18
	procsPageSize = 256
19
	procsBufSize  = 16 * 1024
20
	userHz        = 100
21
)
22

23
// Processes wraps iterative walker over processes under /proc.
24
type Processes struct {
25
	fd       *os.File
26
	dirnames []string
27
	idx      int
28

29
	buf      []byte
30
	pagesize int
31

32
	RootPath string
33
}
34

35
// NewProcesses initializes process info iterator with path /proc.
36
func NewProcesses() (*Processes, error) {
37
	return NewProcessesWithPath("/proc")
38
}
39

40
// NewProcessesWithPath initializes process info iterator with non-default path.
41
func NewProcessesWithPath(rootPath string) (*Processes, error) {
42
	procs := &Processes{
43
		RootPath: rootPath,
44
		buf:      make([]byte, procsBufSize),
45
	}
46

47
	var err error
48

49
	procs.fd, err = os.Open(rootPath)
50
	if err != nil {
51
		return nil, err
52
	}
53

54
	procs.pagesize = os.Getpagesize()
55

56
	return procs, nil
57
}
58

59
// Close the iterator.
60
func (procs *Processes) Close() error {
61
	return procs.fd.Close()
62
}
63

64
// Next returns process info until the list of processes is exhausted.
65
//
66
// Next returns nil, nil when all processes were processed.
67
// Next skips processes which can't be analyzed.
68
func (procs *Processes) Next() (*machine.ProcessInfo, error) {
69
	for {
70
		if procs.idx >= len(procs.dirnames) {
71
			var err error
72

73
			procs.dirnames, err = procs.fd.Readdirnames(procsPageSize)
74
			if err == io.EOF {
75
				return nil, nil
76
			}
77

78
			if err != nil {
79
				return nil, err
80
			}
81

82
			procs.idx = 0
83
		}
84

85
		info, err := procs.readProc(procs.dirnames[procs.idx])
86
		procs.idx++
87

88
		// if err != nil, this process was killed before we were able to read /proc data
89
		if err == nil {
90
			return info, nil
91
		}
92
	}
93
}
94

95
//nolint:gocyclo
96
func (procs *Processes) readProc(pidString string) (*machine.ProcessInfo, error) {
97
	pid, err := strconv.ParseInt(pidString, 10, 32)
98
	if err != nil {
99
		return nil, err
100
	}
101

102
	path := procs.RootPath + "/" + pidString + "/"
103

104
	executable, err := os.Readlink(path + "exe")
105
	if err != nil {
106
		return nil, err
107
	}
108

109
	if err = procs.readFileIntoBuf(path + "comm"); err != nil {
110
		return nil, err
111
	}
112

113
	command := string(bytes.TrimSpace(procs.buf))
114

115
	if err = procs.readFileIntoBuf(path + "cmdline"); err != nil {
116
		return nil, err
117
	}
118

119
	args := string(bytes.ReplaceAll(bytes.TrimRight(procs.buf, "\x00"), []byte{0}, []byte{' '}))
120

121
	if err = procs.readFileIntoBuf(path + "stat"); err != nil {
122
		return nil, err
123
	}
124

125
	rbracket := bytes.LastIndexByte(procs.buf, ')')
126
	if rbracket == -1 {
127
		return nil, errors.New("unexpected format")
128
	}
129

130
	fields := bytes.Fields(procs.buf[rbracket+2:])
131

132
	state := string(fields[0])
133

134
	ppid, err := strconv.ParseInt(string(fields[1]), 10, 32)
135
	if err != nil {
136
		return nil, err
137
	}
138

139
	numThreads, err := strconv.ParseInt(string(fields[17]), 10, 32)
140
	if err != nil {
141
		return nil, err
142
	}
143

144
	uTime, err := strconv.ParseUint(string(fields[11]), 10, 64)
145
	if err != nil {
146
		return nil, err
147
	}
148

149
	sTime, err := strconv.ParseUint(string(fields[12]), 10, 64)
150
	if err != nil {
151
		return nil, err
152
	}
153

154
	vSize, err := strconv.ParseUint(string(fields[20]), 10, 64)
155
	if err != nil {
156
		return nil, err
157
	}
158

159
	rss, err := strconv.ParseUint(string(fields[21]), 10, 64)
160
	if err != nil {
161
		return nil, err
162
	}
163

164
	return &machine.ProcessInfo{
165
		Pid:            int32(pid),
166
		Ppid:           int32(ppid),
167
		State:          state,
168
		Threads:        int32(numThreads),
169
		CpuTime:        float64(uTime+sTime) / userHz,
170
		VirtualMemory:  vSize,
171
		ResidentMemory: rss * uint64(procs.pagesize),
172
		Command:        command,
173
		Executable:     executable,
174
		Args:           args,
175
	}, nil
176
}
177

178
func (procs *Processes) readFileIntoBuf(path string) error {
179
	f, err := os.Open(path)
180
	if err != nil {
181
		return err
182
	}
183

184
	defer f.Close() //nolint:errcheck
185

186
	procs.buf = procs.buf[:cap(procs.buf)]
187

188
	n, err := f.Read(procs.buf)
189
	if err != nil {
190
		return err
191
	}
192

193
	procs.buf = procs.buf[:n]
194

195
	return f.Close()
196
}
197