153 lines
3 KiB
Go
153 lines
3 KiB
Go
|
package budget
|
||
|
|
|
||
|
|
import (
|
||
|
|
"crypto/rand"
|
||
|
|
"encoding/hex"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"os"
|
||
|
|
"path/filepath"
|
||
|
|
"slices"
|
||
|
|
"strings"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type watcher struct {
|
||
|
|
config config
|
||
|
|
ctl unitCtl
|
||
|
|
procs procSource
|
||
|
|
cgroup string
|
||
|
|
owned map[string]bool
|
||
|
|
out io.Writer
|
||
|
|
}
|
||
|
|
|
||
|
|
func randomID() (string, error) {
|
||
|
|
var data [3]byte
|
||
|
|
_, err := rand.Read(data[:])
|
||
|
|
return hex.EncodeToString(data[:]), err
|
||
|
|
}
|
||
|
|
|
||
|
|
func scopeProps(c config, name, slice string, pids []uint32) ([]property, error) {
|
||
|
|
if slice == "" {
|
||
|
|
slice = "app.slice"
|
||
|
|
if name == "search" {
|
||
|
|
slice = "search.slice"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
values := []property{prop("Description", name+" workload"), prop("Slice", slice), prop("CollectMode", "inactive-or-failed"), prop("PIDs", pids)}
|
||
|
|
if name != "search" {
|
||
|
|
limits, err := c.properties(name)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
values = append(values, limits...)
|
||
|
|
}
|
||
|
|
return values, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *watcher) init() error {
|
||
|
|
units, err := w.ctl.Units()
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
w.owned = map[string]bool{}
|
||
|
|
for _, unit := range units {
|
||
|
|
if strings.HasPrefix(unit.Name, "budget-") && strings.HasSuffix(unit.Name, ".scope") {
|
||
|
|
w.owned[unit.Name] = true
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *watcher) reap() error {
|
||
|
|
for _, name := range keys(w.owned) {
|
||
|
|
state, err := w.ctl.State(name)
|
||
|
|
if busError(err, noUnit) {
|
||
|
|
delete(w.owned, name)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
switch state.Active {
|
||
|
|
case "inactive", "failed":
|
||
|
|
delete(w.owned, name)
|
||
|
|
case "active":
|
||
|
|
data, err := os.ReadFile(filepath.Join(w.cgroup, strings.TrimPrefix(state.Group, "/"), "cgroup.events"))
|
||
|
|
if os.IsNotExist(err) {
|
||
|
|
delete(w.owned, name)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if strings.Contains(string(data), "populated 0") {
|
||
|
|
if err := w.ctl.Stop(name); err != nil && !busError(err, noUnit) {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
delete(w.owned, name)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *watcher) step() error {
|
||
|
|
tasks, err := w.procs.Snapshot(w.config)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
pids := make([]uint32, 0, len(tasks))
|
||
|
|
for pid := range tasks {
|
||
|
|
pids = append(pids, pid)
|
||
|
|
}
|
||
|
|
slices.Sort(pids)
|
||
|
|
for _, pid := range pids {
|
||
|
|
value := tasks[pid]
|
||
|
|
if value.Profile == "" {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
group, err := w.procs.Group(pid)
|
||
|
|
if gone(err) {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if group != value.Group {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
members := family(tasks, pid)
|
||
|
|
values, err := scopeProps(w.config, value.Profile, value.Slice, members)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
id, err := randomID()
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
name := "budget-" + value.Profile[:1] + id + ".scope"
|
||
|
|
if err = w.ctl.Start(name, values); busError(err, noPID) {
|
||
|
|
continue
|
||
|
|
} else if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
w.owned[name] = true
|
||
|
|
if _, err := fmt.Fprintf(w.out, "%s: %s %v\n", value.Profile, name, members); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return w.reap()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *watcher) watch() error {
|
||
|
|
if err := w.init(); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
for {
|
||
|
|
if err := w.step(); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
time.Sleep(200 * time.Millisecond)
|
||
|
|
}
|
||
|
|
}
|