hamlab/internal/space/manager.go

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 <root>…`, 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 <mountpoint>...")
		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
}