package space import ( "bufio" "context" "encoding/json" "errors" "fmt" "io" "log" "os" "os/exec" "os/signal" "strings" "sync" "syscall" "time" "hamlab/internal/priv" ) // ---- the helper: `hamlab space-walk …`, run by the daemon (as root when it can) ---- // // The walk runs in its own process so that root can read every directory, so that the // memory of the walk is given back when it ends, and so that a runaway scan can be killed // without touching the console. Progress goes to stderr as JSON lines; the tree to stdout. func Helper(args []string, stdout, stderr io.Writer) int { if len(args) == 0 { fmt.Fprintln(stderr, "usage: hamlab space-walk ...") return 2 } ctx, cancel := context.WithCancel(context.Background()) defer cancel() sig := make(chan os.Signal, 1) signal.Notify(sig, syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP) go func() { <-sig; cancel() }() go func() { _, _ = io.Copy(io.Discard, os.Stdin); cancel() }() // the daemon went away enc := json.NewEncoder(stderr) root, err := Walk(ctx, args, func(p Progress) { _ = enc.Encode(p) }) if err != nil && !errors.Is(err, context.Canceled) { fmt.Fprintln(stderr, "error:", err) return 1 } if ctx.Err() != nil { return 3 } if err := Encode(stdout, root); err != nil { fmt.Fprintln(stderr, "error:", err) return 1 } return 0 } // ---- the daemon side: one scan per disk, the last result kept ---- type Scan struct { Disk string `json:"disk"` Roots []string `json:"roots"` Started time.Time `json:"started"` Finished time.Time `json:"finished,omitempty"` Running bool `json:"running"` Error string `json:"error,omitempty"` Progress Progress `json:"progress"` AsRoot bool `json:"as_root"` // once finished: Files uint32 `json:"files"` Dirs uint32 `json:"dirs"` Size int64 `json:"size"` Denied int64 `json:"denied"` root *Node largest []File cancel context.CancelFunc } type Manager struct { self string // this binary, for the helper mu sync.Mutex scans map[string]*Scan } func New(self string) *Manager { return &Manager{self: self, scans: map[string]*Scan{}} } // Status is the scan of disk (a copy without the tree), or nil. func (m *Manager) Status(disk string) *Scan { m.mu.Lock() defer m.mu.Unlock() s := m.scans[disk] if s == nil { return nil } c := *s c.root, c.largest, c.cancel = nil, nil, nil return &c } func (m *Manager) All() map[string]*Scan { m.mu.Lock() defer m.mu.Unlock() out := map[string]*Scan{} for k, s := range m.scans { c := *s c.root, c.largest, c.cancel = nil, nil, nil out[k] = &c } return out } // Start scans the roots of disk in the background. One scan per disk at a time. func (m *Manager) Start(disk string, roots []string) (*Scan, error) { if len(roots) == 0 { return nil, errors.New("nothing is mounted from this disk") } m.mu.Lock() if s := m.scans[disk]; s != nil && s.Running { m.mu.Unlock() return nil, errors.New("a scan of this disk is already running") } ctx, cancel := context.WithCancel(context.Background()) s := &Scan{Disk: disk, Roots: roots, Started: time.Now(), Running: true, AsRoot: priv.Available(), cancel: cancel} m.scans[disk] = s m.mu.Unlock() go m.run(ctx, s) return m.Status(disk), nil } func (m *Manager) run(ctx context.Context, s *Scan) { args := append([]string{m.self, "space-walk"}, s.Roots...) if s.AsRoot { args = priv.Wrap(args...) } cmd := exec.Command(args[0], args[1:]...) stdin, _ := cmd.StdinPipe() // held open: closing it tells the helper the daemon is gone stdout, _ := cmd.StdoutPipe() stderr, _ := cmd.StderrPipe() finish := func(err error) { m.mu.Lock() defer m.mu.Unlock() s.Running = false s.Finished = time.Now() if err != nil { s.Error = err.Error() } } if err := cmd.Start(); err != nil { finish(err) return } go func() { <-ctx.Done(); _ = cmd.Process.Signal(syscall.SIGTERM); stdin.Close() }() var lastErr string go func() { sc := bufio.NewScanner(stderr) sc.Buffer(make([]byte, 64<<10), 1<<20) for sc.Scan() { line := sc.Text() var p Progress if strings.HasPrefix(line, "{") && json.Unmarshal([]byte(line), &p) == nil { m.mu.Lock() s.Progress = p m.mu.Unlock() continue } if line != "" { lastErr = line } } }() root, decErr := Decode(stdout) waitErr := cmd.Wait() if ctx.Err() != nil { finish(errors.New("cancelled")) return } if waitErr != nil { if lastErr != "" { finish(errors.New(strings.TrimPrefix(lastErr, "error: "))) } else { finish(waitErr) } return } if decErr != nil { finish(decErr) return } m.mu.Lock() s.root = root s.Files, s.Dirs, s.Size, s.Denied = root.Files, root.Dirs, root.Size, s.Progress.Denied s.Running = false s.Finished = time.Now() m.mu.Unlock() log.Printf("space: %s scanned: %d files, %d dirs, %d bytes in %s%s", s.Disk, root.Files, root.Dirs, root.Size, s.Finished.Sub(s.Started).Round(time.Millisecond), map[bool]string{true: " (as root)", false: ""}[s.AsRoot]) } func (m *Manager) Cancel(disk string) bool { m.mu.Lock() s := m.scans[disk] m.mu.Unlock() if s == nil || !s.Running || s.cancel == nil { return false } s.cancel() return true } // Tree is the subtree at path, cut to depth levels of at most limit children each, leaving // out children smaller than min bytes; what is left out is summed into MoreN/MoreSize. A // treemap asks deep with a min the size of the smallest tile it could draw; a folder list // asks one level with no min. type Tree struct { Node Path string `json:"path"` Children []*Tree `json:"children,omitempty"` MoreN int `json:"more_n,omitempty"` MoreSize int64 `json:"more_size,omitempty"` Total int64 `json:"total"` // the whole scan's size, for percentages } func (m *Manager) Tree(disk, path string, depth, limit int, min int64) (*Tree, error) { m.mu.Lock() s := m.scans[disk] var root *Node if s != nil { root = s.root } m.mu.Unlock() if root == nil { return nil, errors.New("no scan of this disk yet") } n := Find(root, path) if n == nil { return nil, errors.New("no such path in the scan") } if path == "" { path = "/" } t := cut(n, path, depth, limit, min) t.Total = root.Size return t, nil } func cut(n *Node, path string, depth, limit int, min int64) *Tree { t := &Tree{Node: *n, Path: path} t.Node.Children = nil if depth <= 0 || !n.Dir { return t } for i, c := range n.Children { if i >= limit || (c.Size < min && i > 0) { // children are sorted largest first: the first is always kept t.MoreN++ t.MoreSize += c.Size continue } p := path + "/" + c.Name if path == "/" { p = "/" + c.Name } t.Children = append(t.Children, cut(c, p, depth-1, limit, min)) } return t } // Largest is the n biggest files of the last scan of disk. func (m *Manager) Largest(disk string, n int) ([]File, error) { m.mu.Lock() s := m.scans[disk] var root *Node var have []File if s != nil { root, have = s.root, s.largest } m.mu.Unlock() if root == nil { return nil, errors.New("no scan of this disk yet") } if len(have) < n { have = Largest(root, n) // the tree is never changed once made, so no lock while walking it m.mu.Lock() if s.root == root { s.largest = have } m.mu.Unlock() } if len(have) > n { return have[:n], nil } return have, nil }