hamlab/internal/space/manager.go

package space

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"sort"
	"strconv"
	"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. The daemon writes a header line to the helper's stdin
// (how each root is to be walked, and what changed) followed by the last scan's tree when
// there is one, then holds stdin open: EOF means the daemon has gone or wants the walk
// stopped. Progress goes to stderr as JSON lines; a result line and the tree to stdout.

type header struct {
	Prev    bool            `json:"prev"`
	Modes   map[string]Mode `json:"modes,omitempty"`
	Dirty   []string        `json:"dirty,omitempty"`
	Touched []string        `json:"touched,omitempty"`
}

type resultLine struct {
	Complete bool  `json:"complete"`
	Partial  int64 `json:"partial"`
	Reused   int64 `json:"reused"`
	Read     int64 `json:"read"`
}

func Helper(args []string, stdin io.Reader, 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() }()
	in := bufio.NewReaderSize(stdin, 1<<20)
	line, err := in.ReadBytes('\n')
	if err != nil {
		fmt.Fprintln(stderr, "error: no header:", err)
		return 2
	}
	var h header
	if err := json.Unmarshal(line, &h); err != nil {
		fmt.Fprintln(stderr, "error: bad header:", err)
		return 2
	}
	opts := Options{Modes: h.Modes, Dirty: set(h.Dirty), Touched: set(h.Touched)}
	if h.Prev {
		if opts.Prev, err = Decode(in); err != nil {
			fmt.Fprintln(stderr, "error: last tree:", err)
			return 2
		}
	}
	go func() { _, _ = io.Copy(io.Discard, in); cancel() }() // the daemon went away, or asked us to stop
	enc := json.NewEncoder(stderr)
	res, err := Walk(ctx, args, opts, func(p Progress) { _ = enc.Encode(p) })
	if err != nil && !errors.Is(err, context.Canceled) {
		fmt.Fprintln(stderr, "error:", err)
		return 1
	}
	rl, _ := json.Marshal(resultLine{Complete: res.Complete, Partial: res.Partial, Reused: res.Reused, Read: res.Read})
	if _, err := stdout.Write(append(rl, '\n')); err != nil {
		return 1
	}
	if err := Encode(stdout, res.Root); err != nil {
		fmt.Fprintln(stderr, "error:", err)
		return 1
	}
	if !res.Complete {
		return 3
	}
	return 0
}

func set(list []string) map[string]bool {
	if len(list) == 0 {
		return nil
	}
	m := make(map[string]bool, len(list))
	for _, s := range list {
		m[s] = true
	}
	return m
}

// ---- the daemon side: one scan per target, the last result kept, on disk too ----

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"`
	Loading  bool      `json:"loading,omitempty"` // read back from disk after a restart, not yet decoded
	Error    string    `json:"error,omitempty"`
	Progress Progress  `json:"progress"`
	AsRoot   bool      `json:"as_root"`
	// How this scan was done: "full", or "refresh" / "resume" from the one before, and by what
	// means per root ("zfs diff" or "mtime").
	Method  string            `json:"method,omitempty"`
	ByRoot  map[string]string `json:"by_root,omitempty"`
	Partial bool              `json:"partial,omitempty"` // cancelled before it finished; Scan again resumes it
	// once finished:
	Files       uint32 `json:"files"`
	Dirs        uint32 `json:"dirs"`
	Size        int64  `json:"size"`
	Denied      int64  `json:"denied"`
	Read        int64  `json:"read"`         // directories read this time
	Reused      int64  `json:"reused"`       // directories kept from the scan before
	PartialDirs int64  `json:"partial_dirs"` // directories not read (cancelled)
	// Bases is, per root that is a ZFS dataset, the newest snapshot that existed when the
	// scan started: the next scan asks `zfs diff` what changed since it.
	Bases map[string]string `json:"bases,omitempty"`

	root    *Node
	largest []File
	cancel  context.CancelFunc
}

type Manager struct {
	self  string // this binary, for the helper
	dir   string // where results are kept between restarts
	mu    sync.Mutex
	scans map[string]*Scan
}

// New makes a manager that keeps results under dir and reads back the ones there.
func New(self, dir string) *Manager {
	m := &Manager{self: self, dir: dir, scans: map[string]*Scan{}}
	m.loadAll()
	return m
}

func (m *Manager) fileFor(id, ext string) string {
	safe := strings.NewReplacer(":", "_", "/", "_", "\\", "_").Replace(id)
	return filepath.Join(m.dir, safe+ext)
}

// loadAll reads the saved scans' metadata now and their trees in the background, one at a
// time, so a restart does not double the memory of the last scans all at once.
func (m *Manager) loadAll() {
	metas, _ := filepath.Glob(filepath.Join(m.dir, "*.json"))
	var pending []*Scan
	for _, mp := range metas {
		b, err := os.ReadFile(mp)
		if err != nil {
			continue
		}
		s := &Scan{}
		if json.Unmarshal(b, s) != nil || s.Disk == "" {
			continue
		}
		s.Running, s.cancel = false, nil
		s.Loading = true
		m.scans[s.Disk] = s
		pending = append(pending, s)
	}
	if len(pending) == 0 {
		return
	}
	go func() {
		for _, s := range pending {
			root, err := m.readTree(s.Disk)
			m.mu.Lock()
			if err != nil {
				log.Printf("space: saved scan of %s unreadable (%v); dropped", s.Disk, err)
				if m.scans[s.Disk] == s {
					delete(m.scans, s.Disk)
				}
			} else if m.scans[s.Disk] == s {
				s.root, s.Loading = root, false
			}
			m.mu.Unlock()
		}
	}()
}

func (m *Manager) readTree(id string) (*Node, error) {
	f, err := os.Open(m.fileFor(id, ".scan"))
	if err != nil {
		return nil, err
	}
	defer f.Close()
	br := bufio.NewReaderSize(f, 1<<20)
	if _, err := br.ReadBytes('\n'); err != nil { // the result line
		return nil, err
	}
	return Decode(br)
}

func (m *Manager) saveMeta(s *Scan) {
	if m.dir == "" {
		return
	}
	_ = os.MkdirAll(m.dir, 0o700)
	b, err := json.MarshalIndent(s, "", "  ")
	if err == nil {
		err = writeAtomic(m.fileFor(s.Disk, ".json"), b)
	}
	if err != nil {
		log.Printf("space: save %s: %v", s.Disk, err)
	}
}

func writeAtomic(path string, b []byte) error {
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, b, 0o600); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

// Status is the scan of id (a copy without the tree), or nil.
func (m *Manager) Status(id string) *Scan {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.copyOf(m.scans[id])
}

func (m *Manager) copyOf(s *Scan) *Scan {
	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 {
		out[k] = m.copyOf(s)
	}
	return out
}

// Start scans the roots of id in the background, from the last scan of it when there is
// one. datasets names the ZFS dataset behind each root that has one.
func (m *Manager) Start(id string, roots []string, datasets map[string]string) (*Scan, error) {
	if len(roots) == 0 {
		return nil, errors.New("nothing is mounted from this disk")
	}
	m.mu.Lock()
	prev := m.scans[id]
	if prev != nil && prev.Running {
		m.mu.Unlock()
		return nil, errors.New("a scan of this disk is already running")
	}
	if prev != nil && (prev.root == nil || prev.Loading) {
		prev = nil // nothing to start from
	}
	ctx, cancel := context.WithCancel(context.Background())
	s := &Scan{Disk: id, Roots: roots, Started: time.Now(), Running: true, AsRoot: priv.Available(), cancel: cancel, Bases: map[string]string{}}
	if prev != nil {
		// the new scan takes the old one's place in the list while it runs, but keeps its
		// tree, so the pages still answer until the new one is in
		s.root, s.largest = prev.root, prev.largest
		s.Files, s.Dirs, s.Size = prev.Files, prev.Dirs, prev.Size
	}
	m.scans[id] = s
	m.mu.Unlock()
	go m.run(ctx, s, prev, datasets)
	return m.Status(id), nil
}

// plan decides how each root is walked this time, given the last scan.
func (m *Manager) plan(ctx context.Context, s, prev *Scan, datasets map[string]string) (header, Options) {
	h := header{Modes: map[string]Mode{}}
	opts := Options{Modes: h.Modes}
	if prev == nil || prev.root == nil {
		s.Method = "full"
		return h, opts
	}
	h.Prev = true
	opts.Prev = prev.root
	s.ByRoot = map[string]string{}
	dirty, touched := map[string]bool{}, map[string]bool{}
	// every dataset is asked at once what changed; the answers come back in root order
	type diffed struct {
		changed []change
		err     error
		took    time.Duration
	}
	answers := map[string]chan diffed{}
	for _, r := range s.Roots {
		if base := prev.Bases[r]; base != "" && datasets[r] != "" {
			ch := make(chan diffed, 1)
			answers[r] = ch
			go func(r, base string) {
				t0 := time.Now()
				c, err := zfsDiff(ctx, base, r)
				ch <- diffed{changed: c, err: err, took: time.Since(t0)}
			}(r, base)
		}
	}
	for _, r := range s.Roots {
		mode := ModeMTime
		how := "mtime"
		if ch, ok := answers[r]; ok {
			m.mu.Lock()
			s.Progress.Current = "asking zfs what changed in " + r
			m.mu.Unlock()
			d := <-ch
			if d.err == nil {
				mode, how = ModeDiff, fmt.Sprintf("zfs diff since %s: %d changes in %s", strings.TrimPrefix(prev.Bases[r], datasets[r]), len(d.changed), d.took.Round(time.Millisecond))
				for _, c := range d.changed {
					dirty[c.parent] = true
					if c.dir {
						dirty[c.path] = true
					}
				}
			} else {
				how = "mtime (zfs diff: " + d.err.Error() + ")"
			}
		}
		if Find(prev.root, r) == nil {
			mode, how = ModeFull, "full (new root)"
		}
		h.Modes[r] = mode
		s.ByRoot[r] = how
	}
	m.mu.Lock()
	s.Progress.Current = ""
	m.mu.Unlock()
	// what the last walk never got to is read whatever the change list says
	for _, p := range partialPaths(prev.root, "/") {
		dirty[p] = true
	}
	for p := range dirty {
		for a := filepath.Dir(p); ; a = filepath.Dir(a) {
			touched[a] = true
			if a == "/" {
				break
			}
		}
	}
	for p := range dirty {
		h.Dirty = append(h.Dirty, p)
	}
	for p := range touched {
		h.Touched = append(h.Touched, p)
	}
	sort.Strings(h.Dirty)
	sort.Strings(h.Touched)
	opts.Dirty, opts.Touched = dirty, touched
	if prev.Partial {
		s.Method = "resume"
	} else {
		s.Method = "refresh"
	}
	return h, opts
}

func partialPaths(n *Node, path string) []string {
	var out []string
	var walk func(n *Node, path string)
	walk = func(n *Node, path string) {
		if n.Partial {
			out = append(out, path)
		}
		for _, c := range n.Children {
			if c.Dir {
				walk(c, join(path, c.Name))
			}
		}
	}
	walk(n, path)
	return out
}

func (m *Manager) run(ctx context.Context, s *Scan, prev *Scan, datasets map[string]string) {
	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()
		}
		m.saveMeta(s)
	}
	// the newest snapshot of each dataset, so the next scan can ask what changed since
	for _, r := range s.Roots {
		if ds := datasets[r]; ds != "" {
			if snap := newestSnapshot(ctx, ds); snap != "" {
				s.Bases[r] = snap
			}
		}
	}
	h, _ := m.plan(ctx, s, prev, datasets)
	if ctx.Err() != nil { // cancelled while zfs was being asked
		finish(errors.New("cancelled"))
		return
	}
	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 to stop
	stdout, _ := cmd.StdoutPipe()
	stderr, _ := cmd.StderrPipe()
	if err := cmd.Start(); err != nil {
		finish(err)
		return
	}
	go func() {
		hb, _ := json.Marshal(h)
		_, err := stdin.Write(append(hb, '\n'))
		if err == nil && h.Prev {
			err = Encode(stdin, prev.root)
		}
		if err != nil {
			log.Printf("space: %s: sending the last tree: %v", s.Disk, err)
		}
		<-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
			}
		}
	}()
	// the result streams to disk as it is decoded, so a restart finds it
	_ = os.MkdirAll(m.dir, 0o700)
	tmp := m.fileFor(s.Disk, ".scan.tmp")
	out, ferr := os.OpenFile(tmp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
	var src io.Reader = stdout
	if ferr == nil {
		src = io.TeeReader(stdout, out)
	}
	br := bufio.NewReaderSize(src, 1<<20)
	var rl resultLine
	line, rerr := br.ReadBytes('\n')
	if rerr == nil {
		rerr = json.Unmarshal(line, &rl)
	}
	var root *Node
	var decErr error
	if rerr == nil {
		root, decErr = Decode(br)
	} else {
		decErr = rerr
		_, _ = io.Copy(io.Discard, br)
	}
	if ferr == nil {
		out.Close()
	}
	waitErr := cmd.Wait()
	cancelled := ctx.Err() != nil
	if root == nil {
		_ = os.Remove(tmp)
		switch {
		case cancelled:
			finish(errors.New("cancelled"))
		case waitErr != nil && lastErr != "":
			finish(errors.New(strings.TrimPrefix(lastErr, "error: ")))
		case waitErr != nil:
			finish(waitErr)
		default:
			finish(decErr)
		}
		return
	}
	if ferr == nil {
		_ = os.Rename(tmp, m.fileFor(s.Disk, ".scan"))
	}
	m.mu.Lock()
	s.root, s.largest = root, nil
	s.Files, s.Dirs, s.Size, s.Denied = root.Files, root.Dirs, root.Size, s.Progress.Denied
	s.Read, s.Reused, s.PartialDirs = rl.Read, rl.Reused, rl.Partial
	s.Partial = !rl.Complete
	s.Running = false
	s.Finished = time.Now()
	if s.Partial {
		s.Error = "cancelled"
	}
	m.saveMeta(s)
	m.mu.Unlock()
	what := "scanned"
	if s.Partial {
		what = "partly scanned (cancelled)"
	}
	log.Printf("space: %s %s: %d files, %d dirs, %d bytes in %s (%s; %d dirs read, %d kept%s)", s.Disk, what, root.Files, root.Dirs, root.Size, s.Finished.Sub(s.Started).Round(time.Millisecond), s.Method, rl.Read, rl.Reused, map[bool]string{true: ", as root", false: ""}[s.AsRoot])
}

func (m *Manager) Cancel(id string) bool {
	m.mu.Lock()
	s := m.scans[id]
	m.mu.Unlock()
	if s == nil || !s.Running || s.cancel == nil {
		return false
	}
	s.cancel()
	return true
}

// Forget drops the scan of id, in memory and on disk.
func (m *Manager) Forget(id string) {
	m.mu.Lock()
	if s := m.scans[id]; s != nil && s.cancel != nil {
		s.cancel()
	}
	delete(m.scans, id)
	m.mu.Unlock()
	_ = os.Remove(m.fileFor(id, ".scan"))
	_ = os.Remove(m.fileFor(id, ".json"))
}

// ---- ZFS: what changed since the last scan ----

type change struct {
	path, parent string
	dir          bool
}

// newestSnapshot is the most recent snapshot of a dataset, "" when it has none.
func newestSnapshot(ctx context.Context, dataset string) string {
	ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
	defer cancel()
	out, err := exec.CommandContext(ctx, "zfs", "list", "-H", "-t", "snapshot", "-o", "name", "-s", "creation", "-d", "1", dataset).Output()
	if err != nil {
		return ""
	}
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")
	last := lines[len(lines)-1]
	if !strings.Contains(last, "@") {
		return ""
	}
	return last
}

// zfsDiff lists what changed in a dataset since a snapshot: every changed, added, removed
// or renamed path, with the directory whose listing that touches.
func zfsDiff(ctx context.Context, snapshot, mount string) ([]change, error) {
	// a diff slower than this is not worth waiting for: the walk by mtime is bounded
	ctx, cancel := context.WithTimeout(ctx, 4*time.Minute)
	defer cancel()
	args := priv.Wrap("zfs", "diff", "-H", "-F", snapshot)
	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
	// on timeout the whole process group goes, not just sudo: zfs itself would otherwise run
	// on holding the pipe, and Wait would sit there with it
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Cancel = func() error { return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) }
	cmd.WaitDelay = 5 * time.Second
	var stderr strings.Builder
	cmd.Stderr = &stderr
	out, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}
	var changes []change
	sc := bufio.NewScanner(out)
	sc.Buffer(make([]byte, 64<<10), 4<<20)
	for sc.Scan() {
		f := strings.Split(sc.Text(), "\t")
		if len(f) < 3 {
			continue
		}
		isDir := f[1] == "/"
		for _, p := range f[2:] { // a rename has the old and the new path
			p = unescape(p)
			if !strings.HasPrefix(p, mount) {
				continue
			}
			changes = append(changes, change{path: p, parent: filepath.Dir(p), dir: isDir})
		}
	}
	if err := cmd.Wait(); err != nil {
		msg := strings.TrimSpace(stderr.String())
		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
			msg = "gave up after 4 min"
		}
		if msg == "" {
			msg = err.Error()
		}
		if i := strings.IndexByte(msg, '\n'); i > 0 {
			msg = msg[:i]
		}
		return nil, errors.New(msg)
	}
	return changes, nil
}

// unescape undoes zfs diff's \ooo escapes for spaces and odd bytes in names.
func unescape(s string) string {
	if !strings.Contains(s, "\\") {
		return s
	}
	var b strings.Builder
	for i := 0; i < len(s); i++ {
		if s[i] == '\\' && i+3 < len(s) {
			if v, err := strconv.ParseUint(s[i+1:i+4], 8, 8); err == nil {
				b.WriteByte(byte(v))
				i += 3
				continue
			}
		}
		b.WriteByte(s[i])
	}
	return b.String()
}

// 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(id, path string, depth, limit int, min int64) (*Tree, error) {
	m.mu.Lock()
	s := m.scans[id]
	var root *Node
	if s != nil {
		root = s.root
	}
	m.mu.Unlock()
	if root == nil {
		if s != nil && s.Loading {
			return nil, errors.New("the last scan is still being read back")
		}
		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 !n.Dir {
		return t
	}
	if depth <= 0 { // as deep as was asked: what is inside is one lump, not nothing
		if len(n.Children) > 0 {
			t.MoreN, t.MoreSize = len(n.Children), n.Size-n.SmallSize
		}
		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
		}
		t.Children = append(t.Children, cut(c, join(path, c.Name), depth-1, limit, min))
	}
	return t
}

// Largest is the n biggest files of the last scan of id.
func (m *Manager) Largest(id string, n int) ([]File, error) {
	m.mu.Lock()
	s := m.scans[id]
	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
}