hamlab/internal/space/walk.go

// Package space measures what is using a disk: a parallel walk of the filesystems a disk
// backs, kept as a tree of directories and the files worth seeing, the way WizTree does it.
//
// The walk never crosses a filesystem boundary by itself — /proc, /sys, another disk's
// mount — but every mountpoint the disk backs is a root of the same scan, and the roots
// are grafted into one tree at their paths, so a ZFS pool of many datasets, or a disk with
// / and /home on separate partitions, reads as one picture.
//
// Files below KeepMin are not kept one by one: a directory of a million small files is
// one "smaller files" entry with their count and total, which is all a treemap can show
// of them anyway, and what keeps a scan of a few million files within a few hundred MB.
//
// A walk can start from the last one. Given the previous tree, a directory is read again
// only when something says it changed: on ZFS, `zfs diff` against a snapshot older than
// the last scan names every path that changed, and directories it does not name are kept
// as they were without so much as a stat; elsewhere a directory whose mtime is what it was
// keeps its list and only its kept files are stat'ed again, since a file can grow without
// its directory noticing. A directory the last walk did not finish (it was cancelled) is
// read again whatever else is true, which is what makes a cancelled scan resumable.
package space

import (
	"context"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

// KeepMin is the smallest file kept as its own node.
const KeepMin = 128 << 10

// Node is a directory or a file. Sizes are bytes allocated on disk (what fills the disk);
// Apparent is the length. For a directory both include everything beneath it.
type Node struct {
	Name      string  `json:"name"`
	Dir       bool    `json:"dir"`
	Size      int64   `json:"size"`
	Apparent  int64   `json:"apparent"`
	MTime     int64   `json:"mtime"` // unix seconds
	MTimeNs   int64   `json:"-"`     // directories: the same to the nanosecond, so a change in the same second is not missed
	Files     uint32  `json:"files"` // beneath a directory (kept and small alike)
	Dirs      uint32  `json:"dirs"`
	SmallN    uint32  `json:"small_n,omitempty"`    // files here below KeepMin, not listed
	SmallSize int64   `json:"small_size,omitempty"` // their allocated bytes
	Mount     bool    `json:"mount,omitempty"`      // another filesystem starts here (not this disk's, or a root of its own)
	Denied    bool    `json:"denied,omitempty"`     // could not be read
	Partial   bool    `json:"partial,omitempty"`    // the walk was cancelled before this directory was read
	Children  []*Node `json:"children,omitempty"`   // directories: sorted by Size, largest first
}

// Progress is how far a running walk has got.
type Progress struct {
	Files   int64  `json:"files"`
	Dirs    int64  `json:"dirs"`   // directories read
	Reused  int64  `json:"reused"` // directories kept from the last scan without reading
	Bytes   int64  `json:"bytes"`
	Denied  int64  `json:"denied"`
	Current string `json:"current"`
}

// Mode is how a root is walked when there is a previous tree.
type Mode int

const (
	ModeFull  Mode = iota // read everything
	ModeMTime             // read a directory again only if its mtime changed; stat its kept files
	ModeDiff              // read only the directories a change list names
)

// Options is what a walk starts from.
type Options struct {
	Prev *Node // the last scan's tree, or nil
	// Modes says how each root is walked (by root path); a root not listed is ModeFull.
	Modes map[string]Mode
	// Dirty is every directory a change list named (ModeDiff roots): those are read again.
	// Touched is every ancestor of a dirty directory: those are kept but looked into.
	Dirty, Touched map[string]bool
}

type walker struct {
	ctx      context.Context
	opts     Options
	mu       sync.Mutex
	queue    []job
	cond     *sync.Cond
	pending  int // dirs queued or being read
	files    atomic.Int64
	dirs     atomic.Int64
	reused   atomic.Int64
	bytes    atomic.Int64
	denied   atomic.Int64
	partial  atomic.Int64
	curMu    sync.Mutex
	current  string
	stopping bool
}

type job struct {
	path string
	node *Node
	prev *Node  // the same directory in the last scan, if it had one
	dev  uint64 // the filesystem this root is on; a child on another is a mountpoint, not descended
	mode Mode
}

// Result is what a walk returns besides the tree.
type Result struct {
	Root     *Node
	Partial  int64 // directories not read because the walk was cancelled
	Reused   int64
	Read     int64
	Complete bool
}

// Walk measures roots (mountpoints, absolute) and returns them grafted into one tree whose
// root is "/". progress, if set, is called about four times a second. A cancelled walk
// still returns the tree, with the directories it did not get to marked Partial.
func Walk(ctx context.Context, roots []string, opts Options, progress func(Progress)) (Result, error) {
	w := &walker{ctx: ctx, opts: opts}
	w.cond = sync.NewCond(&w.mu)
	root := &Node{Name: "/", Dir: true}
	sort.Strings(roots)
	trees := map[string]*Node{}
	for _, r := range roots {
		r = filepath.Clean(r)
		if !filepath.IsAbs(r) {
			continue
		}
		var st syscall.Stat_t
		if err := syscall.Lstat(r, &st); err != nil || st.Mode&syscall.S_IFMT != syscall.S_IFDIR {
			continue
		}
		n := &Node{Name: filepath.Base(r), Dir: true, MTime: int64(st.Mtim.Sec), MTimeNs: st.Mtim.Nano()}
		if r == "/" {
			n = root
		}
		trees[r] = n
		var prev *Node
		if opts.Prev != nil {
			prev = Find(opts.Prev, r)
		}
		mode := opts.Modes[r]
		if prev == nil {
			mode = ModeFull
		}
		w.push(job{path: r, node: n, prev: prev, dev: uint64(st.Dev), mode: mode})
	}
	if len(trees) == 0 {
		return Result{Root: root}, os.ErrNotExist
	}
	workers := 32
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() { defer wg.Done(); w.run() }()
	}
	done := make(chan struct{})
	go func() { wg.Wait(); close(done) }()
	tick := time.NewTicker(250 * time.Millisecond)
	defer tick.Stop()
	cancelled := false
	for waiting := true; waiting; {
		select {
		case <-done:
			waiting = false
		case <-tick.C:
			if progress != nil {
				progress(w.snapshot())
			}
		case <-ctx.Done():
			cancelled = true
			w.mu.Lock()
			w.stopping = true
			for _, j := range w.queue { // never read: say so, so the next walk reads them
				j.node.Partial = true
				w.partial.Add(1)
			}
			w.queue = nil
			w.cond.Broadcast()
			w.mu.Unlock()
			<-done
			waiting = false
		}
	}
	if progress != nil {
		progress(w.snapshot())
	}
	// graft each root at its path; deepest last so parents exist first
	for _, r := range roots {
		n, ok := trees[filepath.Clean(r)]
		if !ok || n == root {
			continue
		}
		graft(root, filepath.Clean(r), n)
	}
	total(root)
	res := Result{Root: root, Partial: w.partial.Load(), Reused: w.reused.Load(), Read: w.dirs.Load(), Complete: !cancelled}
	if cancelled {
		return res, ctx.Err()
	}
	return res, nil
}

func (w *walker) snapshot() Progress {
	w.curMu.Lock()
	cur := w.current
	w.curMu.Unlock()
	return Progress{Files: w.files.Load(), Dirs: w.dirs.Load(), Reused: w.reused.Load(), Bytes: w.bytes.Load(), Denied: w.denied.Load(), Current: cur}
}

func (w *walker) push(j job) {
	w.mu.Lock()
	if w.stopping {
		j.node.Partial = true
		w.partial.Add(1)
		w.mu.Unlock()
		return
	}
	w.queue = append(w.queue, j)
	w.pending++
	w.cond.Signal()
	w.mu.Unlock()
}

// run takes directories off the queue until there are none left and none being read.
func (w *walker) run() {
	for {
		w.mu.Lock()
		for len(w.queue) == 0 && w.pending > 0 && !w.stopping {
			w.cond.Wait()
		}
		if w.stopping || (len(w.queue) == 0 && w.pending == 0) {
			w.cond.Broadcast()
			w.mu.Unlock()
			return
		}
		j := w.queue[len(w.queue)-1] // depth first: keeps the queue short and the disk head local
		w.queue = w.queue[:len(w.queue)-1]
		w.mu.Unlock()
		w.visit(j)
		w.mu.Lock()
		w.pending--
		if w.pending == 0 {
			w.cond.Broadcast()
		}
		w.mu.Unlock()
	}
}

func join(dir, name string) string {
	if dir == "/" {
		return "/" + name
	}
	return dir + "/" + name
}

// visit decides whether a directory must be read or can be taken from the last scan.
func (w *walker) visit(j job) {
	if j.prev == nil || j.prev.Partial || j.prev.Denied || j.mode == ModeFull {
		w.readDir(j)
		return
	}
	switch j.mode {
	case ModeDiff:
		if w.opts.Dirty[j.path] {
			w.readDir(j)
			return
		}
		if !w.opts.Touched[j.path] {
			w.reuse(j, false) // nothing under here changed: keep the whole subtree, no stat at all
			return
		}
		w.reuse(j, true) // something deeper changed: keep this list, look into the children
	case ModeMTime:
		var st syscall.Stat_t
		if err := syscall.Lstat(j.path, &st); err != nil {
			w.readDir(j)
			return
		}
		if st.Mtim.Nano() != j.prev.MTimeNs {
			w.readDir(j) // an entry was added, removed or renamed
			return
		}
		w.reuse(j, true) // same entries; files may have grown, and subdirectories are their own question
	}
}

// reuse keeps a directory as the last scan saw it. With descend, its subdirectories are
// visited (each decides for itself) and, in mtime mode, its kept files are stat'ed again.
func (w *walker) reuse(j job, descend bool) {
	w.reused.Add(1)
	n, p := j.node, j.prev
	n.MTime, n.MTimeNs, n.SmallN, n.SmallSize, n.Denied, n.Mount = p.MTime, p.MTimeNs, p.SmallN, p.SmallSize, p.Denied, p.Mount
	if !descend {
		n.Children = p.Children
		n.Apparent = p.Apparent - childApparent(p) // total() adds the children back
		w.files.Add(int64(p.Files))
		w.bytes.Add(p.Size)
		w.reused.Add(int64(p.Dirs))
		return
	}
	w.files.Add(int64(p.SmallN))
	w.bytes.Add(p.SmallSize)
	n.Apparent = p.Apparent - childApparent(p) // the small files' lengths; total() adds the children back
	var st syscall.Stat_t
	for _, c := range p.Children {
		if c.Dir {
			d := &Node{Name: c.Name, Dir: true, MTime: c.MTime, MTimeNs: c.MTimeNs, Mount: c.Mount}
			n.Children = append(n.Children, d)
			if c.Mount {
				continue
			}
			w.push(job{path: join(j.path, c.Name), node: d, prev: c, dev: j.dev, mode: j.mode})
			continue
		}
		if j.mode == ModeMTime {
			if err := syscall.Lstat(join(j.path, c.Name), &st); err != nil {
				continue // gone since (the directory's mtime would have said so, but be safe)
			}
			c = &Node{Name: c.Name, Size: int64(st.Blocks) * 512, Apparent: int64(st.Size), MTime: int64(st.Mtim.Sec)}
		}
		n.Children = append(n.Children, c)
		w.files.Add(1)
		w.bytes.Add(c.Size)
	}
}

func childApparent(p *Node) int64 {
	var a int64
	for _, c := range p.Children {
		a += c.Apparent
	}
	return a
}

func (w *walker) readDir(j job) {
	w.curMu.Lock()
	w.current = j.path
	w.curMu.Unlock()
	w.dirs.Add(1)
	f, err := os.Open(j.path)
	if err != nil {
		j.node.Denied = true
		w.denied.Add(1)
		return
	}
	names, err := f.Readdirnames(-1)
	f.Close()
	if err != nil {
		j.node.Denied = true
		w.denied.Add(1)
	}
	var prevDirs map[string]*Node
	if j.prev != nil {
		prevDirs = map[string]*Node{}
		for _, c := range j.prev.Children {
			if c.Dir {
				prevDirs[c.Name] = c
			}
		}
	}
	var st syscall.Stat_t
	for _, name := range names {
		p := join(j.path, name)
		if err := syscall.Lstat(p, &st); err != nil {
			continue
		}
		switch st.Mode & syscall.S_IFMT {
		case syscall.S_IFDIR:
			c := &Node{Name: name, Dir: true, MTime: int64(st.Mtim.Sec), MTimeNs: st.Mtim.Nano()}
			j.node.Children = append(j.node.Children, c)
			if uint64(st.Dev) != j.dev {
				c.Mount = true // something else is mounted here; if it is this disk's too, it is a root of its own
				continue
			}
			w.push(job{path: p, node: c, prev: prevDirs[name], dev: j.dev, mode: j.mode})
		case syscall.S_IFREG, syscall.S_IFLNK:
			alloc := int64(st.Blocks) * 512
			w.files.Add(1)
			w.bytes.Add(alloc)
			if alloc >= KeepMin || int64(st.Size) >= KeepMin {
				j.node.Children = append(j.node.Children, &Node{Name: name, Size: alloc, Apparent: int64(st.Size), MTime: int64(st.Mtim.Sec)})
			} else {
				j.node.SmallN++
				j.node.SmallSize += alloc
				j.node.Apparent += int64(st.Size) // added to the directory's own apparent total below
			}
		}
	}
}

// graft puts a root's tree at path under root, making any directories in between.
func graft(root *Node, path string, tree *Node) {
	parts := strings.Split(strings.Trim(path, "/"), "/")
	cur := root
	for i, part := range parts {
		var next *Node
		for _, c := range cur.Children {
			if c.Name == part && c.Dir {
				next = c
				break
			}
		}
		if i == len(parts)-1 {
			if next == nil {
				cur.Children = append(cur.Children, tree)
			} else {
				// the parent's walk saw this directory as a mountpoint; the walk of the mount fills it
				next.Children, next.SmallN, next.SmallSize, next.Apparent, next.Denied, next.Partial = tree.Children, tree.SmallN, tree.SmallSize, tree.Apparent, tree.Denied, tree.Partial
				next.Mount = true
			}
			return
		}
		if next == nil {
			next = &Node{Name: part, Dir: true}
			cur.Children = append(cur.Children, next)
		}
		cur = next
	}
}

// total sums every directory from its children and sorts them largest first.
func total(n *Node) {
	if !n.Dir {
		return
	}
	n.Size = n.SmallSize
	n.Files = n.SmallN
	n.Dirs = 0
	apparent := n.Apparent // the small files' lengths
	for _, c := range n.Children {
		total(c)
		n.Size += c.Size
		apparent += c.Apparent
		if c.Dir {
			n.Dirs += 1 + c.Dirs
			n.Files += c.Files
		} else {
			n.Files++
		}
	}
	n.Apparent = apparent
	sort.Slice(n.Children, func(i, j int) bool { return n.Children[i].Size > n.Children[j].Size })
}

// Find follows path (absolute, "/" for the root) down the tree.
func Find(root *Node, path string) *Node {
	cur := root
	for _, part := range strings.Split(strings.Trim(filepath.Clean(path), "/"), "/") {
		if part == "" {
			continue
		}
		var next *Node
		for _, c := range cur.Children {
			if c.Name == part {
				next = c
				break
			}
		}
		if next == nil {
			return nil
		}
		cur = next
	}
	return cur
}

// File is one of the largest files, with its full path.
type File struct {
	Path     string `json:"path"`
	Size     int64  `json:"size"`
	Apparent int64  `json:"apparent"`
	MTime    int64  `json:"mtime"`
}

// Largest is the n biggest files in the tree, largest first.
func Largest(root *Node, n int) []File {
	var out []File
	minIdx := -1
	var walk func(n *Node, path string)
	walk = func(node *Node, path string) {
		for _, c := range node.Children {
			p := join(path, c.Name)
			if c.Dir {
				if c.Size > 0 && (len(out) < n || c.Size > out[minIdx].Size) {
					walk(c, p)
				}
				continue
			}
			if len(out) < n {
				out = append(out, File{Path: p, Size: c.Size, Apparent: c.Apparent, MTime: c.MTime})
				if minIdx < 0 || c.Size < out[minIdx].Size {
					minIdx = len(out) - 1
				}
			} else if c.Size > out[minIdx].Size {
				out[minIdx] = File{Path: p, Size: c.Size, Apparent: c.Apparent, MTime: c.MTime}
				minIdx = 0
				for i := range out {
					if out[i].Size < out[minIdx].Size {
						minIdx = i
					}
				}
			}
		}
	}
	walk(root, "/")
	sort.Slice(out, func(i, j int) bool { return out[i].Size > out[j].Size })
	return out
}