Merge branch 'lukeshu/rebuild-nodes-take6'

11 files changed, 82 insertions, 49 deletions

diff --git a/cmd/btrfs-rec/main.go b/cmd/btrfs-rec/main.go
index d4165bf..dc00dab 100644
--- a/cmd/btrfs-rec/main.go
+++ b/cmd/btrfs-rec/main.go
@@ -103,7 +103,9 @@ func main() {
 				ctx := cmd.Context()
 				logger := textui.NewLogger(os.Stderr, logLevelFlag.Level)
 				ctx = dlog.WithLogger(ctx, logger)
-				ctx = dlog.WithField(ctx, "mem", new(textui.LiveMemUse))
+				if logLevelFlag.Level >= dlog.LogLevelDebug {
+					ctx = dlog.WithField(ctx, "mem", new(textui.LiveMemUse))
+				}
 				dlog.SetFallbackLogger(logger.WithField("btrfs-progs.THIS_IS_A_BUG", true))
 
 				grp := dgroup.NewGroup(ctx, dgroup.GroupConfig{
diff --git a/go.mod b/go.mod
index 3cd0cac..1ef3515 100644
--- a/go.mod
+++ b/go.mod
@@ -7,7 +7,7 @@ module git.lukeshu.com/btrfs-progs-ng
 go 1.20
 
 require (
-	git.lukeshu.com/go/lowmemjson v0.3.4
+	git.lukeshu.com/go/lowmemjson v0.3.8
 	git.lukeshu.com/go/typedsync v0.1.0
 	github.com/datawire/dlib v1.3.0
 	github.com/datawire/ocibuild v0.0.3-0.20220423003204-fc6a4e9f90dc
diff --git a/go.sum b/go.sum
index 716bc99..79eadc9 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-git.lukeshu.com/go/lowmemjson v0.3.4 h1:lGvZhZ6qclJOulDoSNUHRsdzZIGolHesg392wrBXskY=
-git.lukeshu.com/go/lowmemjson v0.3.4/go.mod h1:cP+ybDhmhZYlTNZjqMMhEjp0kmGDwzkygw/3fXcME0U=
+git.lukeshu.com/go/lowmemjson v0.3.8 h1:wFm5jg8iC1owd1veYMryMgbSG8m4aeXb+wNEJnIyUiI=
+git.lukeshu.com/go/lowmemjson v0.3.8/go.mod h1:cP+ybDhmhZYlTNZjqMMhEjp0kmGDwzkygw/3fXcME0U=
 git.lukeshu.com/go/typedsync v0.1.0 h1:BYv123nWCymA3zZpokP6nDdtNQ6p7Q51hSWGno/U3Dc=
 git.lukeshu.com/go/typedsync v0.1.0/go.mod h1:EAn7NcfoGeGMv3DWxKQnifcT/rYPAIEqp9Rsz//oYqY=
 github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
index 9ec2849..dbbc6eb 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/forrest.go
@@ -22,6 +22,7 @@ import (
 
 type Callbacks interface {
 	AddedItem(ctx context.Context, tree btrfsprim.ObjID, key btrfsprim.Key)
+	AddedRoot(ctx context.Context, tree btrfsprim.ObjID, root btrfsvol.LogicalAddr)
 	LookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool)
 	LookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool)
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
index eab3eb2..39d8871 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/btrees/tree.go
@@ -304,6 +304,7 @@ func (tree *RebuiltTree) AddRoot(ctx context.Context, rootNode btrfsvol.LogicalA
 	if (tree.ID == btrfsprim.ROOT_TREE_OBJECTID || tree.ID == btrfsprim.UUID_TREE_OBJECTID) && stats.AddedItems > 0 {
 		tree.forrest.flushNegativeCache(ctx)
 	}
+	tree.forrest.cb.AddedRoot(ctx, tree.ID, rootNode)
 }
 
 // main public API /////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
index 9e3b144..04df2b6 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/keyio/keyio.go
@@ -7,6 +7,7 @@ package keyio
 import (
 	"context"
 	"fmt"
+	"sync"
 
 	"github.com/datawire/dlib/dlog"
 
@@ -48,6 +49,7 @@ type Handle struct {
 	Names map[ItemPtr][]byte      // DIR_INDEX
 	Sizes map[ItemPtr]SizeAndErr  // EXTENT_CSUM and EXTENT_DATA
 
+	mu    sync.Mutex
 	cache containers.ARCache[btrfsvol.LogicalAddr, *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]]
 }
 
@@ -153,6 +155,8 @@ func (o *Handle) readNode(ctx context.Context, laddr btrfsvol.LogicalAddr) *disk
 }
 
 func (o *Handle) ReadItem(ctx context.Context, ptr ItemPtr) btrfsitem.Item {
+	o.mu.Lock()
+	defer o.mu.Unlock()
 	if o.graph.Nodes[ptr.Node].Level != 0 {
 		panic(fmt.Errorf("should not happen: keyio.Handle.ReadItem called for non-leaf node@%v", ptr.Node))
 	}
@@ -164,5 +168,5 @@ func (o *Handle) ReadItem(ctx context.Context, ptr ItemPtr) btrfsitem.Item {
 		panic(fmt.Errorf("should not happen: keyio.Handle.ReadItem called for out-of-bounds item index: index=%v len=%v",
 			ptr.Idx, len(items)))
 	}
-	return items[ptr.Idx].Body
+	return items[ptr.Idx].Body.CloneItem()
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index dc78c2e..bc36485 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -56,6 +56,7 @@ type rebuilder struct {
 		Key    containers.Optional[btrfsprim.Key]
 	}
 	treeQueue          containers.Set[btrfsprim.ObjID]
+	retryItemQueue     map[btrfsprim.ObjID]containers.Set[keyAndTree]
 	addedItemQueue     containers.Set[keyAndTree]
 	settledItemQueue   containers.Set[keyAndTree]
 	augmentQueue       map[btrfsprim.ObjID]*treeAugmentQueue
@@ -98,6 +99,7 @@ func (o *rebuilder) Rebuild(ctx context.Context) error {
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.step", "rebuild")
 
 	// Initialize
+	o.retryItemQueue = make(map[btrfsprim.ObjID]containers.Set[keyAndTree])
 	o.addedItemQueue = make(containers.Set[keyAndTree])
 	o.settledItemQueue = make(containers.Set[keyAndTree])
 	o.augmentQueue = make(map[btrfsprim.ObjID]*treeAugmentQueue)
@@ -133,7 +135,7 @@ func (o *rebuilder) Rebuild(ctx context.Context) error {
 		}
 		runtime.GC()
 
-		// Apply augments (drain o.augmentQueue, fill o.addedItemQueue).
+		// Apply augments (drain o.augmentQueue (and maybe o.retryItemQueue), fill o.addedItemQueue).
 		if err := o.processAugmentQueue(ctx); err != nil {
 			return err
 		}
@@ -286,6 +288,7 @@ func (o *rebuilder) processSettledItemQueue(ctx context.Context) error {
 				Key:  item.Key,
 				Body: item.Body,
 			})
+			item.Body.Free()
 			if item.ItemType == btrfsitem.ROOT_ITEM_KEY {
 				o.treeQueue.Insert(item.ObjectID)
 			}
@@ -300,7 +303,7 @@ func (o *rebuilder) processSettledItemQueue(ctx context.Context) error {
 	return grp.Wait()
 }
 
-// processAugmentQueue drains o.augmentQueue, filling o.addedItemQueue.
+// processAugmentQueue drains o.augmentQueue (and maybe o.retryItemQueue), filling o.addedItemQueue.
 func (o *rebuilder) processAugmentQueue(ctx context.Context) error {
 	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.substep", "apply-augments")
 
@@ -342,9 +345,12 @@ func (o *rebuilder) processAugmentQueue(ctx context.Context) error {
 	return nil
 }
 
-func (o *rebuilder) enqueueRetry() {
+func (o *rebuilder) enqueueRetry(ifTreeID btrfsprim.ObjID) {
 	if o.curKey.Key.OK {
-		o.settledItemQueue.Insert(keyAndTree{
+		if o.retryItemQueue[ifTreeID] == nil {
+			o.retryItemQueue[ifTreeID] = make(containers.Set[keyAndTree])
+		}
+		o.retryItemQueue[ifTreeID].Insert(keyAndTree{
 			TreeID: o.curKey.TreeID,
 			Key:    o.curKey.Key.Val,
 		})
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_treecb.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_treecb.go
index eb0204c..492436b 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_treecb.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_treecb.go
@@ -10,6 +10,7 @@ import (
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 )
 
 // AddedItem implements btrees.Callbacks.
@@ -20,6 +21,13 @@ func (o *rebuilder) AddedItem(ctx context.Context, tree btrfsprim.ObjID, key btr
 	})
 }
 
+// AddedRoot implements btrees.Callbacks.
+func (o *rebuilder) AddedRoot(ctx context.Context, tree btrfsprim.ObjID, root btrfsvol.LogicalAddr) {
+	if retries := o.retryItemQueue[tree]; retries != nil {
+		o.addedItemQueue.InsertFrom(retries)
+	}
+}
+
 // LookupRoot implements btrees.Callbacks.
 func (o *rebuilder) LookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool) {
 	wantKey := WantWithTree{
@@ -33,10 +41,12 @@ func (o *rebuilder) LookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offse
 	ctx = withWant(ctx, logFieldTreeWant, "tree Root", wantKey)
 	foundKey, ok := o._want(ctx, wantKey)
 	if !ok {
-		o.enqueueRetry()
+		o.enqueueRetry(btrfsprim.ROOT_TREE_OBJECTID)
 		return 0, btrfsitem.Root{}, false
 	}
-	switch itemBody := o.rebuilt.Tree(ctx, wantKey.TreeID).ReadItem(ctx, foundKey).(type) {
+	itemBody := o.rebuilt.Tree(ctx, wantKey.TreeID).ReadItem(ctx, foundKey)
+	defer itemBody.Free()
+	switch itemBody := itemBody.(type) {
 	case *btrfsitem.Root:
 		return btrfsprim.Generation(foundKey.Offset), *itemBody, true
 	case *btrfsitem.Error:
@@ -57,10 +67,12 @@ func (o *rebuilder) LookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id btr
 	}
 	ctx = withWant(ctx, logFieldTreeWant, "resolve parent UUID", wantKey)
 	if !o._wantOff(ctx, wantKey) {
-		o.enqueueRetry()
+		o.enqueueRetry(btrfsprim.UUID_TREE_OBJECTID)
 		return 0, false
 	}
-	switch itemBody := o.rebuilt.Tree(ctx, wantKey.TreeID).ReadItem(ctx, wantKey.Key.Key()).(type) {
+	itemBody := o.rebuilt.Tree(ctx, wantKey.TreeID).ReadItem(ctx, wantKey.Key.Key())
+	defer itemBody.Free()
+	switch itemBody := itemBody.(type) {
 	case *btrfsitem.UUIDMap:
 		return itemBody.ObjID, true
 	case *btrfsitem.Error:
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wantcb.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wantcb.go
index c57b2bb..adf3cff 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wantcb.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wantcb.go
@@ -40,7 +40,7 @@ func (o *rebuilder) want(ctx context.Context, reason string, treeID btrfsprim.Ob
 
 func (o *rebuilder) _want(ctx context.Context, wantKey WantWithTree) (key btrfsprim.Key, ok bool) {
 	if o.rebuilt.Tree(ctx, wantKey.TreeID) == nil {
-		o.enqueueRetry()
+		o.enqueueRetry(wantKey.TreeID)
 		return btrfsprim.Key{}, false
 	}
 
@@ -90,7 +90,7 @@ func (o *rebuilder) wantOff(ctx context.Context, reason string, treeID btrfsprim
 
 func (o *rebuilder) _wantOff(ctx context.Context, wantKey WantWithTree) (ok bool) {
 	if o.rebuilt.Tree(ctx, wantKey.TreeID) == nil {
-		o.enqueueRetry()
+		o.enqueueRetry(wantKey.TreeID)
 		return false
 	}
 
@@ -131,7 +131,7 @@ func (o *rebuilder) wantDirIndex(ctx context.Context, reason string, treeID btrf
 	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
 
 	if o.rebuilt.Tree(ctx, treeID) == nil {
-		o.enqueueRetry()
+		o.enqueueRetry(treeID)
 		return
 	}
 
@@ -256,7 +256,7 @@ func (o *rebuilder) _wantRange(
 	wantKey.Key.OffsetType = offsetRange
 
 	if o.rebuilt.Tree(ctx, treeID) == nil {
-		o.enqueueRetry()
+		o.enqueueRetry(treeID)
 		return
 	}
 
@@ -365,7 +365,7 @@ func (o *rebuilder) wantCSum(ctx context.Context, reason string, inodeTree, inod
 	}
 	inodeCtx := withWant(ctx, logFieldItemWant, reason, inodeWant)
 	if !o._wantOff(inodeCtx, inodeWant) {
-		o.enqueueRetry()
+		o.enqueueRetry(inodeTree)
 		return
 	}
 	inodePtr, ok := o.rebuilt.Tree(inodeCtx, inodeTree).Items(inodeCtx).Load(inodeWant.Key.Key())
diff --git a/lib/textui/log.go b/lib/textui/log.go
index b4dcea4..2a6fdd4 100644
--- a/lib/textui/log.go
+++ b/lib/textui/log.go
@@ -237,28 +237,30 @@ func (l *logger) log(lvl dlog.LogLevel, writeMsg func(io.Writer)) {
 	}
 
 	// caller //////////////////////////////////////////////////////////////
-	const (
-		thisModule             = "git.lukeshu.com/btrfs-progs-ng"
-		thisPackage            = "git.lukeshu.com/btrfs-progs-ng/lib/textui"
-		maximumCallerDepth int = 25
-		minimumCallerDepth int = 3 // runtime.Callers + .log + .Log
-	)
-	var pcs [maximumCallerDepth]uintptr
-	depth := runtime.Callers(minimumCallerDepth, pcs[:])
-	frames := runtime.CallersFrames(pcs[:depth])
-	for f, again := frames.Next(); again; f, again = frames.Next() {
-		if !strings.HasPrefix(f.Function, thisModule+"/") {
-			continue
-		}
-		if strings.HasPrefix(f.Function, thisPackage+".") {
-			continue
-		}
-		if nextField == len(fieldKeys) {
-			logBuf.WriteString(" :")
+	if lvl >= dlog.LogLevelDebug {
+		const (
+			thisModule             = "git.lukeshu.com/btrfs-progs-ng"
+			thisPackage            = "git.lukeshu.com/btrfs-progs-ng/lib/textui"
+			maximumCallerDepth int = 25
+			minimumCallerDepth int = 3 // runtime.Callers + .log + .Log
+		)
+		var pcs [maximumCallerDepth]uintptr
+		depth := runtime.Callers(minimumCallerDepth, pcs[:])
+		frames := runtime.CallersFrames(pcs[:depth])
+		for f, again := frames.Next(); again; f, again = frames.Next() {
+			if !strings.HasPrefix(f.Function, thisModule+"/") {
+				continue
+			}
+			if strings.HasPrefix(f.Function, thisPackage+".") {
+				continue
+			}
+			if nextField == len(fieldKeys) {
+				logBuf.WriteString(" :")
+			}
+			file := f.File[strings.LastIndex(f.File, thisModDir+"/")+len(thisModDir+"/"):]
+			fmt.Fprintf(logBuf, " (from %s:%d)", file, f.Line)
+			break
 		}
-		file := f.File[strings.LastIndex(f.File, thisModDir+"/")+len(thisModDir+"/"):]
-		fmt.Fprintf(logBuf, " (from %s:%d)", file, f.Line)
-		break
 	}
 
 	// boilerplate /////////////////////////////////////////////////////////
diff --git a/lib/textui/log_test.go b/lib/textui/log_test.go
index bcd9c39..44958cb 100644
--- a/lib/textui/log_test.go
+++ b/lib/textui/log_test.go
@@ -15,8 +15,13 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
-func logLineRegexp(inner string) string {
-	return `[0-9]{2}:[0-9]{2}:[0-9]{2}\.[0-9]{4} ` + inner + ` \(from lib/textui/log_test\.go:[0-9]+\)\n`
+func logLineRegexp(includePos bool, inner string) string {
+	ret := `[0-9]{2}:[0-9]{2}:[0-9]{2}\.[0-9]{4} ` + inner
+	if includePos {
+		ret += ` : \(from lib/textui/log_test\.go:[0-9]+\)`
+	}
+	ret += `\n`
+	return ret
 }
 
 func TestLogFormat(t *testing.T) {
@@ -25,7 +30,7 @@ func TestLogFormat(t *testing.T) {
 	ctx := dlog.WithLogger(context.Background(), textui.NewLogger(&out, dlog.LogLevelTrace))
 	dlog.Debugf(ctx, "foo %d", 12345)
 	assert.Regexp(t,
-		`^`+logLineRegexp(`DBG : foo 12,345 :`)+`$`,
+		`^`+logLineRegexp(true, `DBG : foo 12,345`)+`$`,
 		out.String())
 }
 
@@ -45,12 +50,12 @@ func TestLogLevel(t *testing.T) {
 	dlog.Error(ctx, "Error")
 	assert.Regexp(t,
 		`^`+
-			logLineRegexp(`ERR : Error :`)+
-			logLineRegexp(`WRN : Warn :`)+
-			logLineRegexp(`INF : Info :`)+
-			logLineRegexp(`INF : Info :`)+
-			logLineRegexp(`WRN : Warn :`)+
-			logLineRegexp(`ERR : Error :`)+
+			logLineRegexp(false, `ERR : Error`)+
+			logLineRegexp(false, `WRN : Warn`)+
+			logLineRegexp(false, `INF : Info`)+
+			logLineRegexp(false, `INF : Info`)+
+			logLineRegexp(false, `WRN : Warn`)+
+			logLineRegexp(false, `ERR : Error`)+
 			`$`,
 		out.String())
 }
@@ -62,6 +67,6 @@ func TestLogField(t *testing.T) {
 	ctx = dlog.WithField(ctx, "foo", 12345)
 	dlog.Info(ctx, "msg")
 	assert.Regexp(t,
-		`^`+logLineRegexp(`INF : msg : foo=12,345`)+`$`,
+		`^`+logLineRegexp(false, `INF : msg : foo=12,345`)+`$`,
 		out.String())
 }