Rework visualize-nodes

2 files changed, 186 insertions, 70 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
index 9a2bd89..caccfb9 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
@@ -65,6 +65,46 @@ func (m uuidMap) ParentTree(tree btrfsprim.ObjID) (btrfsprim.ObjID, bool) {
 	return parentObjID, true
 }
 
+type fullAncestorLister struct {
+	uuidMap       uuidMap
+	treeAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]
+
+	memos map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]
+}
+
+func newFullAncestorLister(uuidMap uuidMap, treeAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]) fullAncestorLister {
+	return fullAncestorLister{
+		uuidMap:       uuidMap,
+		treeAncestors: treeAncestors,
+		memos:         make(map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]),
+	}
+}
+
+func (fa fullAncestorLister) GetFullAncestors(child btrfsprim.ObjID) containers.Set[btrfsprim.ObjID] {
+	if memoized := fa.memos[child]; memoized != nil {
+		return memoized
+	}
+	ret := make(containers.Set[btrfsprim.ObjID])
+	fa.memos[child] = ret
+
+	// Try to use '.uuidMap' instead of '.treeAncestors' if possible.
+	knownAncestors := make(containers.Set[btrfsprim.ObjID])
+	if parent, ok := fa.uuidMap.ParentTree(child); ok {
+		if parent == 0 {
+			return ret
+		}
+		knownAncestors.Insert(parent)
+	} else {
+		knownAncestors.InsertFrom(fa.treeAncestors[child])
+	}
+
+	for _, ancestor := range maps.SortedKeys(knownAncestors) {
+		ret.Insert(ancestor)
+		ret.InsertFrom(fa.GetFullAncestors(ancestor))
+	}
+	return ret
+}
+
 func (m uuidMap) considerAncestors(ctx context.Context, treeAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]) {
 	if missing := m.missingRootItems(); len(missing) == 0 {
 		return
@@ -72,34 +112,8 @@ func (m uuidMap) considerAncestors(ctx context.Context, treeAncestors map[btrfsp
 		dlog.Infof(ctx, "Rebuilding %d root items...", len(missing))
 	}
 
-	var fullAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]
-	var getFullAncestors func(child btrfsprim.ObjID) containers.Set[btrfsprim.ObjID]
-	getFullAncestors = func(child btrfsprim.ObjID) containers.Set[btrfsprim.ObjID] {
-		if memoized := fullAncestors[child]; memoized != nil {
-			return memoized
-		}
-		ret := make(containers.Set[btrfsprim.ObjID])
-		fullAncestors[child] = ret
-
-		// Try to use 'm' instead of 'treeAncestors' if possible.
-		knownAncestors := make(containers.Set[btrfsprim.ObjID])
-		if parent, ok := m.ParentTree(child); ok {
-			if parent == 0 {
-				return ret
-			}
-			knownAncestors.Insert(parent)
-		} else {
-			knownAncestors.InsertFrom(treeAncestors[child])
-		}
+	fa := newFullAncestorLister(m, treeAncestors)
 
-		for _, ancestor := range maps.SortedKeys(knownAncestors) {
-			ret.Insert(ancestor)
-			ret.InsertFrom(getFullAncestors(ancestor))
-		}
-		return ret
-	}
-
-	fullAncestors = make(map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID])
 	for _, missingRoot := range maps.SortedKeys(m.missingRootItems()) {
 		if _, ok := m.TreeParent[missingRoot]; ok {
 			// This one is incomplete because it doesn't have a UUID, not
@@ -107,9 +121,9 @@ func (m uuidMap) considerAncestors(ctx context.Context, treeAncestors map[btrfsp
 			continue
 		}
 		potentialParents := make(containers.Set[btrfsprim.ObjID])
-		potentialParents.InsertFrom(getFullAncestors(missingRoot))
-		for ancestor := range getFullAncestors(missingRoot) {
-			potentialParents.DeleteFrom(getFullAncestors(ancestor))
+		potentialParents.InsertFrom(fa.GetFullAncestors(missingRoot))
+		for ancestor := range fa.GetFullAncestors(missingRoot) {
+			potentialParents.DeleteFrom(fa.GetFullAncestors(ancestor))
 		}
 		if len(potentialParents) == 1 {
 			parent := potentialParents.TakeOne()
@@ -166,6 +180,13 @@ func buildUUIDMap(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sc
 		SeenTrees: make(containers.Set[btrfsprim.ObjID]),
 	}
 
+	// These 4 trees are mentioned directly in the superblock, so
+	// they are always seen.
+	ret.SeenTrees.Insert(btrfsprim.ROOT_TREE_OBJECTID)
+	ret.SeenTrees.Insert(btrfsprim.CHUNK_TREE_OBJECTID)
+	ret.SeenTrees.Insert(btrfsprim.TREE_LOG_OBJECTID)
+	ret.SeenTrees.Insert(btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
+
 	progress()
 	for _, devResults := range scanResults {
 		for laddr := range devResults.FoundNodes {
@@ -189,6 +210,7 @@ func buildUUIDMap(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sc
 					if err := maybeSet("ParentUUID", ret.TreeParent, item.Key.ObjectID, itemBody.ParentUUID); err != nil {
 						return uuidMap{}, err
 					}
+					ret.SeenTrees.Insert(item.Key.ObjectID)
 				case btrfsitem.UUIDMap:
 					uuid := btrfsitem.KeyToUUID(item.Key)
 					if err := maybeSet("ObjID2UUID", ret.ObjID2UUID, itemBody.ObjID, uuid); err != nil {
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/visualizenodes.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/visualizenodes.go
index 38742a0..c3784bb 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/visualizenodes.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/visualizenodes.go
@@ -5,14 +5,18 @@
 package rebuildnodes
 
 import (
+	"bufio"
 	"context"
 	"errors"
 	"fmt"
 	"html"
-	"io"
 	iofs "io/fs"
+	"os"
+	"path/filepath"
 	"strings"
 
+	"github.com/datawire/dlib/dlog"
+
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
@@ -24,7 +28,36 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
-func VisualizeNodes(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) error {
+func ptrTo[T any](x T) *T {
+	return &x
+}
+
+func getCliques(uuidMap uuidMap, treeAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID]) map[btrfsprim.ObjID]*containers.Set[btrfsprim.ObjID] {
+	cliques := make(map[btrfsprim.ObjID]*containers.Set[btrfsprim.ObjID])
+	lister := newFullAncestorLister(uuidMap, treeAncestors)
+	for _, treeID := range maps.SortedKeys(uuidMap.SeenTrees) {
+		clique := ptrTo(make(containers.Set[btrfsprim.ObjID]))
+		clique.Insert(treeID)
+		clique.InsertFrom(lister.GetFullAncestors(treeID))
+		for _, id := range maps.SortedKeys(*clique) {
+			if existingClique, ok := cliques[id]; ok {
+				clique.InsertFrom(*existingClique)
+			}
+			cliques[id] = clique
+		}
+	}
+	return cliques
+}
+
+func getCliqueID(cliques map[btrfsprim.ObjID]*containers.Set[btrfsprim.ObjID], treeID btrfsprim.ObjID) btrfsprim.ObjID {
+	clique, ok := cliques[treeID]
+	if !ok {
+		panic(fmt.Errorf("tree ID %v was not in .SeenTrees", treeID))
+	}
+	return maps.SortedKeys(*clique)[0]
+}
+
+func VisualizeNodes(ctx context.Context, dir string, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) error {
 	uuidMap, err := buildUUIDMap(ctx, fs, nodeScanResults)
 	if err != nil {
 		return err
@@ -44,8 +77,12 @@ func VisualizeNodes(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanRe
 
 	////////////////////////////////////////////////////////////////////////////////////////////
 
-	nodes := make(map[btrfsprim.ObjID]containers.Set[string])
-	edges := make(containers.Set[string])
+	cliques := getCliques(uuidMap, treeAncestors)
+
+	dlog.Info(ctx, "Building graphviz graph...")
+
+	nodes := make(map[btrfsprim.ObjID]containers.Set[string]) // grouped by treeID
+	edges := make(map[btrfsprim.ObjID]containers.Set[string]) // grouped by cliqueID
 	visitedNodes := make(containers.Set[btrfsvol.LogicalAddr])
 	var isOrphan bool
 
@@ -54,42 +91,37 @@ func VisualizeNodes(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanRe
 
 		// Node
 		var treeID btrfsprim.ObjID
+		var cliqueID btrfsprim.ObjID
 		var nodeStr string
 		if err != nil && (errors.Is(err, btrfstree.ErrNotANode) || errors.As(err, new(*btrfstree.IOError))) {
 			treeID = 0
-			nodeStr = fmt.Sprintf(`n%d [shape=star label="%v"]`, addr, addr)
+			func() {
+				defer func() {
+					if r := recover(); r != nil {
+						panic(fmt.Errorf("path=%#v (%s): %v", path, path, r))
+					}
+				}()
+				cliqueID = getCliqueID(cliques, path.Node(-1).FromTree)
+			}()
+			nodeStr = fmt.Sprintf(`n%d [shape=star color=red label="%v"]`, addr, addr)
 		} else {
 			treeID = nodeRef.Data.Head.Owner
+			cliqueID = getCliqueID(cliques, treeID)
 			var buf strings.Builder
-			fmt.Fprintf(&buf, `n%d [shape=record label="%v\ngen=%v\nlvl=%v|`,
+			fmt.Fprintf(&buf, `n%d [shape=record label="%v\ngen=%v\nlvl=%v|{`,
 				addr,
 				nodeRef.Data.Head.Addr,
 				nodeRef.Data.Head.Generation,
 				nodeRef.Data.Head.Level)
-			if nodeRef.Data.Head.Level == 0 {
-				for i, item := range nodeRef.Data.BodyLeaf {
-					sep := "|"
-					if i == 0 {
-						sep = "{"
-					}
-					fmt.Fprintf(&buf, "%s<p%d>%d: (%d,%v,%d)",
-						sep, i, i,
-						item.Key.ObjectID,
-						item.Key.ItemType,
-						item.Key.Offset)
-				}
+			if nodeRef.Data.Head.NumItems == 0 {
+				buf.WriteString("(no items)")
 			} else {
-				for i, ptr := range nodeRef.Data.BodyInternal {
-					sep := "|"
+				for i := uint32(0); i < nodeRef.Data.Head.NumItems; i++ {
 					if i == 0 {
-						sep = "{"
+						fmt.Fprintf(&buf, "<p%[1]d>%[1]d", i)
+					} else {
+						fmt.Fprintf(&buf, "|<p%[1]d>%[1]d", i)
 					}
-					fmt.Fprintf(&buf, "%s<p%d>%d: (%d,%v,%d) gen=%v",
-						sep, i, i,
-						ptr.Key.ObjectID,
-						ptr.Key.ItemType,
-						ptr.Key.Offset,
-						ptr.Generation)
 				}
 			}
 			buf.WriteString(`}"]`)
@@ -112,11 +144,24 @@ func VisualizeNodes(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanRe
 		} else {
 			fmt.Fprintf(&edge, "n%d:p%d", path.Node(-2).ToNodeAddr, path.Node(-1).FromItemIdx)
 		}
-		fmt.Fprintf(&edge, " -> n%d", addr)
+		fmt.Fprintf(&edge, ` -> n%d [label="`, addr)
+		if path.Node(-1).FromItemIdx >= 0 {
+			fmt.Fprintf(&edge, "%d: key=(%d,%v,%d) gen=%v",
+				path.Node(-1).FromItemIdx,
+				path.Node(-1).ToKey.ObjectID,
+				path.Node(-1).ToKey.ItemType,
+				path.Node(-1).ToKey.Offset,
+				path.Node(-1).ToNodeGeneration)
+		}
 		if err != nil {
-			fmt.Fprintf(&edge, ` [color=red label="%s"]`, html.EscapeString(err.Error()))
+			fmt.Fprintf(&edge, `\n\n%s" color=red]`, html.EscapeString(err.Error()))
+		} else {
+			edge.WriteString(`"]`)
 		}
-		edges.Insert(edge.String())
+		if _, ok := edges[cliqueID]; !ok {
+			edges[cliqueID] = make(containers.Set[string])
+		}
+		edges[cliqueID].Insert(edge.String())
 
 		// Return
 		if visitedNodes.Has(addr) {
@@ -151,20 +196,69 @@ func VisualizeNodes(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanRe
 		)
 	}
 
+	dlog.Info(ctx, "... done building")
+
 	////////////////////////////////////////////////////////////////////////////////////////////
 
-	fmt.Fprintln(out, "digraph FS {")
-	for _, treeID := range maps.SortedKeys(nodes) {
-		fmt.Fprintf(out, "  subgraph cluster_%d {\n", treeID)
-		for _, node := range maps.SortedKeys(nodes[treeID]) {
-			fmt.Fprintf(out, "    %s;\n", node)
-		}
-		fmt.Fprintln(out, "  }")
+	dlog.Infof(ctx, "Writing graphviz output to %q...", dir)
+
+	cliqueIDs := maps.SortedKeys(edges)
+
+	if err := os.MkdirAll(dir, 0777); err != nil {
+		return err
 	}
-	for _, edge := range maps.SortedKeys(edges) {
-		fmt.Fprintf(out, "  %s;\n", edge)
+	for _, cliqueID := range cliqueIDs {
+		if err := func() (err error) {
+			maybeSetErr := func(_err error) {
+				if err == nil && _err != nil {
+					err = _err
+				}
+			}
+			fh, err := os.OpenFile(filepath.Join(dir, fmt.Sprintf("%d.dot", cliqueID)), os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
+			if err != nil {
+				return err
+			}
+			defer func() {
+				maybeSetErr(fh.Close())
+			}()
+			buf := bufio.NewWriter(fh)
+			defer func() {
+				maybeSetErr(buf.Flush())
+			}()
+
+			if _, err := fmt.Fprintf(buf, "digraph clique%d {\n", cliqueID); err != nil {
+				return err
+			}
+			clique := cliques[cliqueID]
+			for _, treeID := range maps.SortedKeys(*clique) {
+				if _, err := fmt.Fprintf(buf, "  subgraph cluster_tree%d {\n", treeID); err != nil {
+					return err
+				}
+				for _, node := range maps.SortedKeys(nodes[treeID]) {
+					if _, err := fmt.Fprintf(buf, "    %s;\n", node); err != nil {
+						return err
+					}
+				}
+				if _, err := fmt.Fprintln(buf, "  }"); err != nil {
+					return err
+				}
+			}
+			for _, edge := range maps.SortedKeys(edges[cliqueID]) {
+				if _, err := fmt.Fprintf(buf, "  %s;\n", edge); err != nil {
+					return err
+				}
+			}
+			if _, err := fmt.Fprintln(buf, "}"); err != nil {
+				return err
+			}
+
+			return nil
+		}(); err != nil {
+			return err
+		}
 	}
-	fmt.Fprintln(out, "}")
+
+	dlog.Info(ctx, "... done writing")
 
 	return nil
 }