data: replace Tree with TreeNodeIterator

The TreeNodeIterator decodes nodes while iterating over a tree blob.
This should reduce peak memory usage as now only the serialized tree
blob and a single node have to alive at the same time. Using the
iterator has implications for the error handling however. Now it is
necessary that all loops that iterate through a tree check for errors
before using the node returned by the iterator.

The other change is that it is no longer possible to iterate over a tree
multiple times. Instead it must be loaded a second time. This only
affects the tree rewriting code.

internal/data/find.go

@@ -22,16 +22,19 @@ func FindUsedBlobs(ctx context.Context, repo restic.Loader, treeIDs restic.IDs,
 		blobs.Insert(h)
 		lock.Unlock()
 		return blobReferenced
-	}, func(_ restic.ID, err error, tree *Tree) error {
+	}, func(_ restic.ID, err error, nodes TreeNodeIterator) error {
 		if err != nil {
 			return err
 		}
 
-		for _, node := range tree.Nodes {
+		for item := range nodes {
+			if item.Error != nil {
+				return item.Error
+			}
 			lock.Lock()
-			switch node.Type {
+			switch item.Node.Type {
 			case NodeTypeFile:
-				for _, blob := range node.Content {
+				for _, blob := range item.Node.Content {
 					blobs.Insert(restic.BlobHandle{ID: blob, Type: restic.DataBlob})
 				}
 			}

internal/data/testing.go

@@ -5,12 +5,15 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
+	"slices"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/restic/chunker"
 	"github.com/restic/restic/internal/restic"
 	"github.com/restic/restic/internal/test"
+	rtest "github.com/restic/restic/internal/test"
 )
 
 // fakeFile returns a reader which yields deterministic pseudo-random data.
@@ -72,9 +75,8 @@ func (fs *fakeFileSystem) saveTree(ctx context.Context, uploader restic.BlobSave
 	rnd := rand.NewSource(seed)
 	numNodes := int(rnd.Int63() % maxNodes)
 
-	var tree Tree
+	var nodes []*Node
 	for i := 0; i < numNodes; i++ {
-
 		// randomly select the type of the node, either tree (p = 1/4) or file (p = 3/4).
 		if depth > 1 && rnd.Int63()%4 == 0 {
 			treeSeed := rnd.Int63() % maxSeed
@@ -87,7 +89,7 @@ func (fs *fakeFileSystem) saveTree(ctx context.Context, uploader restic.BlobSave
 				Subtree: &id,
 			}
 
-			tree.Nodes = append(tree.Nodes, node)
+			nodes = append(nodes, node)
 			continue
 		}
 
@@ -102,14 +104,24 @@ func (fs *fakeFileSystem) saveTree(ctx context.Context, uploader restic.BlobSave
 		}
 
 		node.Content = fs.saveFile(ctx, uploader, fakeFile(fileSeed, fileSize))
-		tree.Nodes = append(tree.Nodes, node)
+		nodes = append(nodes, node)
 	}
 
-	tree.Sort()
-	id, err := SaveTree(ctx, uploader, &tree)
-	if err != nil {
-		fs.t.Fatalf("SaveTree returned error: %v", err)
+	return TestSaveNodes(fs.t, ctx, uploader, nodes)
+}
+
+//nolint:revive // as this is a test helper, t should go first
+func TestSaveNodes(t testing.TB, ctx context.Context, uploader restic.BlobSaver, nodes []*Node) restic.ID {
+	slices.SortFunc(nodes, func(a, b *Node) int {
+		return strings.Compare(a.Name, b.Name)
+	})
+	treeWriter := NewTreeWriter(uploader)
+	for _, node := range nodes {
+		err := treeWriter.AddNode(node)
+		rtest.OK(t, err)
 	}
+	id, err := treeWriter.Finalize(ctx)
+	rtest.OK(t, err)
 	return id
 }
 

internal/data/tree.go

@@ -5,124 +5,164 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+
+	"io"
+	"iter"
 	"path"
-	"sort"
 	"strings"
 
 	"github.com/restic/restic/internal/errors"
 	"github.com/restic/restic/internal/restic"
-
-	"github.com/restic/restic/internal/debug"
 )
 
+// For documentation purposes only:
+// // Tree is an ordered list of nodes.
+// type Tree struct {
+//         Nodes []*Node `json:"nodes"`
+// }
+
 var ErrTreeNotOrdered = errors.New("nodes are not ordered or duplicate")
 
-// Tree is an ordered list of nodes.
-type Tree struct {
-	Nodes []*Node `json:"nodes"`
+type treeIterator struct {
+	dec     json.Decoder
+	started bool
 }
 
-// NewTree creates a new tree object with the given initial capacity.
-func NewTree(capacity int) *Tree {
-	return &Tree{
-		Nodes: make([]*Node, 0, capacity),
-	}
+type NodeOrError struct {
+	Node  *Node
+	Error error
 }
 
-func (t *Tree) String() string {
-	return fmt.Sprintf("Tree<%d nodes>", len(t.Nodes))
-}
+type TreeNodeIterator = iter.Seq[NodeOrError]
 
-// Equals returns true if t and other have exactly the same nodes.
-func (t *Tree) Equals(other *Tree) bool {
-	if len(t.Nodes) != len(other.Nodes) {
-		debug.Log("tree.Equals(): trees have different number of nodes")
-		return false
+func NewTreeNodeIterator(rd io.Reader) (TreeNodeIterator, error) {
+	t := &treeIterator{
+		dec: *json.NewDecoder(rd),
 	}
 
-	for i := 0; i < len(t.Nodes); i++ {
-		if !t.Nodes[i].Equals(*other.Nodes[i]) {
-			debug.Log("tree.Equals(): node %d is different:", i)
-			debug.Log("  %#v", t.Nodes[i])
-			debug.Log("  %#v", other.Nodes[i])
-			return false
+	err := t.init()
+	if err != nil {
+		return nil, err
+	}
+
+	return func(yield func(NodeOrError) bool) {
+		if t.started {
+			panic("tree iterator is single use only")
 		}
-	}
-
-	return true
+		t.started = true
+		for {
+			n, err := t.next()
+			if err != nil && errors.Is(err, io.EOF) {
+				return
+			}
+			if !yield(NodeOrError{Node: n, Error: err}) {
+				return
+			}
+			// errors are final
+			if err != nil {
+				return
+			}
+		}
+	}, nil
 }
 
-// Insert adds a new node at the correct place in the tree.
-func (t *Tree) Insert(node *Node) error {
-	pos, found := t.find(node.Name)
-	if found != nil {
-		return errors.Errorf("node %q already present", node.Name)
+func (t *treeIterator) init() error {
+	// `{"nodes":[` `]}`
+
+	if err := t.assertToken(json.Delim('{')); err != nil {
+		return err
+	}
+	if err := t.assertToken("nodes"); err != nil {
+		return err
+	}
+	if err := t.assertToken(json.Delim('[')); err != nil {
+		return err
 	}
-
-	// https://github.com/golang/go/wiki/SliceTricks
-	t.Nodes = append(t.Nodes, nil)
-	copy(t.Nodes[pos+1:], t.Nodes[pos:])
-	t.Nodes[pos] = node
-
 	return nil
 }
 
-func (t *Tree) find(name string) (int, *Node) {
-	pos := sort.Search(len(t.Nodes), func(i int) bool {
-		return t.Nodes[i].Name >= name
-	})
-
-	if pos < len(t.Nodes) && t.Nodes[pos].Name == name {
-		return pos, t.Nodes[pos]
-	}
-
-	return pos, nil
-}
-
-// Find returns a node with the given name, or nil if none could be found.
-func (t *Tree) Find(name string) *Node {
-	if t == nil {
-		return nil
-	}
-
-	_, node := t.find(name)
-	return node
-}
-
-// Sort sorts the nodes by name.
-func (t *Tree) Sort() {
-	list := Nodes(t.Nodes)
-	sort.Sort(list)
-	t.Nodes = list
-}
-
-// Subtrees returns a slice of all subtree IDs of the tree.
-func (t *Tree) Subtrees() (trees restic.IDs) {
-	for _, node := range t.Nodes {
-		if node.Type == NodeTypeDir && node.Subtree != nil {
-			trees = append(trees, *node.Subtree)
+func (t *treeIterator) next() (*Node, error) {
+	if t.dec.More() {
+		var n Node
+		err := t.dec.Decode(&n)
+		if err != nil {
+			return nil, err
 		}
+		return &n, nil
 	}
 
-	return trees
+	if err := t.assertToken(json.Delim(']')); err != nil {
+		return nil, err
+	}
+	if err := t.assertToken(json.Delim('}')); err != nil {
+		return nil, err
+	}
+	return nil, io.EOF
 }
 
-// LoadTree loads a tree from the repository.
-func LoadTree(ctx context.Context, r restic.BlobLoader, id restic.ID) (*Tree, error) {
-	debug.Log("load tree %v", id)
+func (t *treeIterator) assertToken(token json.Token) error {
+	to, err := t.dec.Token()
+	if err != nil {
+		return err
+	}
+	if to != token {
+		return errors.Errorf("error decoding tree: expected %v, got %v", token, to)
+	}
+	return nil
+}
 
-	buf, err := r.LoadBlob(ctx, restic.TreeBlob, id, nil)
+func LoadTree(ctx context.Context, loader restic.BlobLoader, content restic.ID) (TreeNodeIterator, error) {
+	rd, err := loader.LoadBlob(ctx, restic.TreeBlob, content, nil)
 	if err != nil {
 		return nil, err
 	}
+	return NewTreeNodeIterator(bytes.NewReader(rd))
+}
 
-	t := &Tree{}
-	err = json.Unmarshal(buf, t)
-	if err != nil {
-		return nil, err
+type TreeFinder struct {
+	next    func() (NodeOrError, bool)
+	stop    func()
+	current *Node
+}
+
+func NewTreeFinder(tree TreeNodeIterator) *TreeFinder {
+	if tree == nil {
+		return &TreeFinder{stop: func() {}}
+	}
+	next, stop := iter.Pull(tree)
+	return &TreeFinder{next: next, stop: stop}
+}
+
+// Find finds the node with the given name. If the node is not found, it returns nil.
+// If Find was called before, the new name must be strictly greater than the last name.
+func (t *TreeFinder) Find(name string) (*Node, error) {
+	if t.next == nil {
+		return nil, nil
+	}
+	// loop until `t.current.Name` is >= name
+	for t.current == nil || t.current.Name < name {
+		current, ok := t.next()
+		if current.Error != nil {
+			return nil, current.Error
+		}
+		if !ok {
+			return nil, nil
+		}
+		t.current = current.Node
 	}
 
-	return t, nil
+	if t.current.Name == name {
+		// forget the current node to free memory as early as possible
+		current := t.current
+		t.current = nil
+		return current, nil
+	}
+	// we have already passed the name
+	return nil, nil
+}
+
+func (t *TreeFinder) Close() {
+	t.stop()
 }
 
 type TreeWriter struct {
@@ -148,10 +188,13 @@ func (t *TreeWriter) Finalize(ctx context.Context) (restic.ID, error) {
 	return id, err
 }
 
-func SaveTree(ctx context.Context, saver restic.BlobSaver, t *Tree) (restic.ID, error) {
+func SaveTree(ctx context.Context, saver restic.BlobSaver, nodes TreeNodeIterator) (restic.ID, error) {
 	treeWriter := NewTreeWriter(saver)
-	for _, node := range t.Nodes {
-		err := treeWriter.AddNode(node)
+	for item := range nodes {
+		if item.Error != nil {
+			return restic.ID{}, item.Error
+		}
+		err := treeWriter.AddNode(item.Node)
 		if err != nil {
 			return restic.ID{}, err
 		}
@@ -214,7 +257,12 @@ func FindTreeDirectory(ctx context.Context, repo restic.BlobLoader, id *restic.I
 		if err != nil {
 			return nil, fmt.Errorf("path %s: %w", subfolder, err)
 		}
-		node := tree.Find(name)
+		finder := NewTreeFinder(tree)
+		node, err := finder.Find(name)
+		finder.Close()
+		if err != nil {
+			return nil, fmt.Errorf("path %s: %w", subfolder, err)
+		}
 		if node == nil {
 			return nil, fmt.Errorf("path %s: not found", subfolder)
 		}

internal/data/tree_stream.go

@@ -23,12 +23,36 @@ type trackedID struct {
 	rootIdx int
 }
 
+// subtreesCollector wraps a TreeNodeIterator and returns a new iterator that collects the subtrees.
+func subtreesCollector(tree TreeNodeIterator) (TreeNodeIterator, func() restic.IDs) {
+	subtrees := restic.IDs{}
+	isComplete := false
+
+	return func(yield func(NodeOrError) bool) {
+			for item := range tree {
+				if !yield(item) {
+					return
+				}
+				// be defensive and check for nil subtree as this code is also used by the checker
+				if item.Node != nil && item.Node.Type == NodeTypeDir && item.Node.Subtree != nil {
+					subtrees = append(subtrees, *item.Node.Subtree)
+				}
+			}
+			isComplete = true
+		}, func() restic.IDs {
+			if !isComplete {
+				panic("tree was not read completely")
+			}
+			return subtrees
+		}
+}
+
 // loadTreeWorker loads trees from repo and sends them to out.
 func loadTreeWorker(
 	ctx context.Context,
 	repo restic.Loader,
 	in <-chan trackedID,
-	process func(id restic.ID, error error, tree *Tree) error,
+	process func(id restic.ID, error error, nodes TreeNodeIterator) error,
 	out chan<- trackedTreeItem,
 ) error {
 
@@ -39,14 +63,21 @@ func loadTreeWorker(
 		}
 		debug.Log("load tree %v (%v) returned err: %v", tree, treeID, err)
 
+		//  wrap iterator to collect subtrees while `process` iterates over `tree`
+		var collectSubtrees func() restic.IDs
+		if tree != nil {
+			tree, collectSubtrees = subtreesCollector(tree)
+		}
+
 		err = process(treeID.ID, err, tree)
 		if err != nil {
 			return err
 		}
 
+		// assume that the number of subtrees is within reasonable limits, such that the memory usage is not a problem
 		var subtrees restic.IDs
-		if tree != nil {
-			subtrees = tree.Subtrees()
+		if collectSubtrees != nil {
+			subtrees = collectSubtrees()
 		}
 
 		job := trackedTreeItem{ID: treeID.ID, Subtrees: subtrees, rootIdx: treeID.rootIdx}
@@ -159,7 +190,7 @@ func StreamTrees(
 	trees restic.IDs,
 	p *progress.Counter,
 	skip func(tree restic.ID) bool,
-	process func(id restic.ID, error error, tree *Tree) error,
+	process func(id restic.ID, error error, nodes TreeNodeIterator) error,
 ) error {
 	loaderChan := make(chan trackedID)
 	hugeTreeChan := make(chan trackedID, 10)

internal/data/tree_test.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"os"
 	"path/filepath"
+	"slices"
 	"strconv"
 	"testing"
 
@@ -105,37 +106,45 @@ func TestNodeComparison(t *testing.T) {
 func TestEmptyLoadTree(t *testing.T) {
 	repo := repository.TestRepository(t)
 
-	tree := data.NewTree(0)
+	nodes := []*data.Node{}
 	var id restic.ID
 	rtest.OK(t, repo.WithBlobUploader(context.TODO(), func(ctx context.Context, uploader restic.BlobSaverWithAsync) error {
-		var err error
 		// save tree
-		id, err = data.SaveTree(ctx, uploader, tree)
-		return err
+		id = data.TestSaveNodes(t, ctx, uploader, nodes)
+		return nil
 	}))
 
 	// load tree again
-	tree2, err := data.LoadTree(context.TODO(), repo, id)
+	it, err := data.LoadTree(context.TODO(), repo, id)
 	rtest.OK(t, err)
+	nodes2 := []*data.Node{}
+	for item := range it {
+		rtest.OK(t, item.Error)
+		nodes2 = append(nodes2, item.Node)
+	}
 
-	rtest.Assert(t, tree.Equals(tree2),
-		"trees are not equal: want %v, got %v",
-		tree, tree2)
+	rtest.Assert(t, slices.Equal(nodes, nodes2),
+		"tree nodes are not equal: want %v, got %v",
+		nodes, nodes2)
+}
+
+// Basic type for comparing the serialization of the tree
+type Tree struct {
+	Nodes []*data.Node `json:"nodes"`
 }
 
 func TestTreeEqualSerialization(t *testing.T) {
 	files := []string{"node.go", "tree.go", "tree_test.go"}
 	for i := 1; i <= len(files); i++ {
-		tree := data.NewTree(i)
+		tree := Tree{Nodes: make([]*data.Node, 0, i)}
 		builder := data.NewTreeJSONBuilder()
 
 		for _, fn := range files[:i] {
 			node := nodeForFile(t, fn)
 
-			rtest.OK(t, tree.Insert(node))
+			tree.Nodes = append(tree.Nodes, node)
 			rtest.OK(t, builder.AddNode(node))
 
-			rtest.Assert(t, tree.Insert(node) != nil, "no error on duplicate node")
 			rtest.Assert(t, builder.AddNode(node) != nil, "no error on duplicate node")
 			rtest.Assert(t, errors.Is(builder.AddNode(node), data.ErrTreeNotOrdered), "wrong error returned")
 		}
@@ -144,11 +153,11 @@ func TestTreeEqualSerialization(t *testing.T) {
 		treeBytes = append(treeBytes, '\n')
 		rtest.OK(t, err)
 
-		stiBytes, err := builder.Finalize()
+		buf, err := builder.Finalize()
 		rtest.OK(t, err)
 
 		// compare serialization of an individual node and the SaveTreeIterator
-		rtest.Equals(t, treeBytes, stiBytes)
+		rtest.Equals(t, treeBytes, buf)
 	}
 }
 
@@ -165,11 +174,12 @@ func BenchmarkBuildTree(b *testing.B) {
 	b.ReportAllocs()
 
 	for i := 0; i < b.N; i++ {
-		t := data.NewTree(size)
-
+		t := data.NewTreeJSONBuilder()
 		for i := range nodes {
-			_ = t.Insert(&nodes[i])
+			rtest.OK(b, t.AddNode(&nodes[i]))
 		}
+		_, err := t.Finalize()
+		rtest.OK(b, err)
 	}
 }
 
@@ -186,8 +196,11 @@ func testLoadTree(t *testing.T, version uint) {
 	repo, _, _ := repository.TestRepositoryWithVersion(t, version)
 	sn := archiver.TestSnapshot(t, repo, rtest.BenchArchiveDirectory, nil)
 
-	_, err := data.LoadTree(context.TODO(), repo, *sn.Tree)
+	nodes, err := data.LoadTree(context.TODO(), repo, *sn.Tree)
 	rtest.OK(t, err)
+	for item := range nodes {
+		rtest.OK(t, item.Error)
+	}
 }
 
 func BenchmarkLoadTree(t *testing.B) {