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.
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
	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
	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"`
	Bytes   int64  `json:"bytes"`
	Denied  int64  `json:"denied"`
	Current string `json:"current"`
}

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

type job struct {
	path string
	node *Node
	dev  uint64 // the filesystem this root is on; a child on another is a mountpoint, not descended
}

// 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.
func Walk(ctx context.Context, roots []string, progress func(Progress)) (*Node, error) {
	w := &walker{ctx: ctx}
	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: st.Mtim.Sec}
		if r == "/" {
			n = root
		}
		trees[r] = n
		w.push(job{path: r, node: n, dev: uint64(st.Dev)})
	}
	if len(trees) == 0 {
		return root, os.ErrNotExist
	}
	workers := 16
	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()
	for waiting := true; waiting; {
		select {
		case <-done:
			waiting = false
		case <-tick.C:
			if progress != nil {
				progress(w.snapshot())
			}
		case <-ctx.Done():
			w.mu.Lock()
			w.stopping = true
			w.queue = nil
			w.cond.Broadcast()
			w.mu.Unlock()
			<-done
			return root, ctx.Err()
		}
	}
	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)
	return root, 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(), Bytes: w.bytes.Load(), Denied: w.denied.Load(), Current: cur}
}

func (w *walker) push(j job) {
	w.mu.Lock()
	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.readDir(j)
		w.mu.Lock()
		w.pending--
		if w.pending == 0 {
			w.cond.Broadcast()
		}
		w.mu.Unlock()
	}
}

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 st syscall.Stat_t
	for _, name := range names {
		p := j.path + "/" + name
		if j.path == "/" {
			p = "/" + 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: st.Mtim.Sec}
			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, dev: j.dev})
		case syscall.S_IFREG, syscall.S_IFLNK:
			alloc := st.Blocks * 512
			w.files.Add(1)
			w.bytes.Add(alloc)
			if alloc >= KeepMin || st.Size >= KeepMin {
				j.node.Children = append(j.node.Children, &Node{Name: name, Size: alloc, Apparent: st.Size, MTime: st.Mtim.Sec})
			} else {
				j.node.SmallN++
				j.node.SmallSize += alloc
				j.node.Apparent += 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 = tree.Children, tree.SmallN, tree.SmallSize, tree.Apparent, tree.Denied
				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 := path + "/" + c.Name
			if path == "/" {
				p = "/" + 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
}