From c892c0bab950aef9547cdbf20b01a2b158875b29 Mon Sep 17 00:00:00 2001
From: greatroar <61184462+greatroar@users.noreply.github.com>
Date: Sun, 26 Sep 2021 17:18:42 +0200
Subject: [PATCH] internal/restic: Don't allocate in Tree.Insert
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

name         old time/op    new time/op    delta
BuildTree-8    34.6µs ± 4%     7.0µs ± 3%  -79.68%  (p=0.000 n=18+19)

name         old alloc/op   new alloc/op   delta
BuildTree-8    34.0kB ± 0%     0.9kB ± 0%  -97.37%  (p=0.000 n=20+20)

name         old allocs/op  new allocs/op  delta
BuildTree-8       108 ± 0%         1 ± 0%  -99.07%  (p=0.000 n=20+20)
---
 cmd/restic/cmd_recover.go       |  2 +-
 internal/archiver/archiver.go   |  5 +++--
 internal/archiver/tree_saver.go |  3 ++-
 internal/restic/tree.go         | 10 +++++-----
 internal/restic/tree_test.go    | 24 +++++++++++++++++++++++-
 internal/walker/walker_test.go  |  2 +-
 6 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/cmd/restic/cmd_recover.go b/cmd/restic/cmd_recover.go
index 44d93d233..860c45a57 100644
--- a/cmd/restic/cmd_recover.go
+++ b/cmd/restic/cmd_recover.go
@@ -107,7 +107,7 @@ func runRecover(gopts GlobalOptions) error {
 		return nil
 	}
 
-	tree := restic.NewTree()
+	tree := restic.NewTree(len(roots))
 	for id := range roots {
 		var subtreeID = id
 		node := restic.Node{
diff --git a/internal/archiver/archiver.go b/internal/archiver/archiver.go
index 0fb6e1ae7..e22c097ec 100644
--- a/internal/archiver/archiver.go
+++ b/internal/archiver/archiver.go
@@ -550,12 +550,13 @@ func (arch *Archiver) statDir(dir string) (os.FileInfo, error) {
 func (arch *Archiver) SaveTree(ctx context.Context, snPath string, atree *Tree, previous *restic.Tree) (*restic.Tree, error) {
 	debug.Log("%v (%v nodes), parent %v", snPath, len(atree.Nodes), previous)
 
-	tree := restic.NewTree()
+	nodeNames := atree.NodeNames()
+	tree := restic.NewTree(len(nodeNames))
 
 	futureNodes := make(map[string]FutureNode)
 
 	// iterate over the nodes of atree in lexicographic (=deterministic) order
-	for _, name := range atree.NodeNames() {
+	for _, name := range nodeNames {
 		subatree := atree.Nodes[name]
 
 		// test if context has been cancelled
diff --git a/internal/archiver/tree_saver.go b/internal/archiver/tree_saver.go
index c41392643..867bad6aa 100644
--- a/internal/archiver/tree_saver.go
+++ b/internal/archiver/tree_saver.go
@@ -103,7 +103,8 @@ type saveTreeResponse struct {
 func (s *TreeSaver) save(ctx context.Context, snPath string, node *restic.Node, nodes []FutureNode) (*restic.Node, ItemStats, error) {
 	var stats ItemStats
 
-	tree := restic.NewTree()
+	tree := restic.NewTree(len(nodes))
+
 	for _, fn := range nodes {
 		fn.wait(ctx)
 
diff --git a/internal/restic/tree.go b/internal/restic/tree.go
index 81650105a..09ea44dfa 100644
--- a/internal/restic/tree.go
+++ b/internal/restic/tree.go
@@ -14,10 +14,10 @@ type Tree struct {
 	Nodes []*Node `json:"nodes"`
 }
 
-// NewTree creates a new tree object.
-func NewTree() *Tree {
+// NewTree creates a new tree object with the given initial capacity.
+func NewTree(capacity int) *Tree {
 	return &Tree{
-		Nodes: []*Node{},
+		Nodes: make([]*Node, 0, capacity),
 	}
 }
 
@@ -51,8 +51,8 @@ func (t *Tree) Insert(node *Node) error {
 		return errors.Errorf("node %q already present", node.Name)
 	}
 
-	// https://code.google.com/p/go-wiki/wiki/SliceTricks
-	t.Nodes = append(t.Nodes, &Node{})
+	// 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
 
diff --git a/internal/restic/tree_test.go b/internal/restic/tree_test.go
index 2bcda6760..598a0ad4b 100644
--- a/internal/restic/tree_test.go
+++ b/internal/restic/tree_test.go
@@ -6,6 +6,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strconv"
 	"testing"
 
 	"github.com/restic/restic/internal/repository"
@@ -98,7 +99,7 @@ func TestLoadTree(t *testing.T) {
 	defer cleanup()
 
 	// save tree
-	tree := restic.NewTree()
+	tree := restic.NewTree(0)
 	id, err := repo.SaveTree(context.TODO(), tree)
 	rtest.OK(t, err)
 
@@ -113,3 +114,24 @@ func TestLoadTree(t *testing.T) {
 		"trees are not equal: want %v, got %v",
 		tree, tree2)
 }
+
+func BenchmarkBuildTree(b *testing.B) {
+	const size = 100 // Directories of this size are not uncommon.
+
+	nodes := make([]restic.Node, size)
+	for i := range nodes {
+		// Archiver.SaveTree inputs in sorted order, so do that here too.
+		nodes[i].Name = strconv.Itoa(i)
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		t := restic.NewTree(size)
+
+		for i := range nodes {
+			_ = t.Insert(&nodes[i])
+		}
+	}
+}
diff --git a/internal/walker/walker_test.go b/internal/walker/walker_test.go
index a4e820940..7a939b1e2 100644
--- a/internal/walker/walker_test.go
+++ b/internal/walker/walker_test.go
@@ -23,7 +23,7 @@ func BuildTreeMap(tree TestTree) (m TreeMap, root restic.ID) {
 }
 
 func buildTreeMap(tree TestTree, m TreeMap) restic.ID {
-	res := restic.NewTree()
+	res := restic.NewTree(0)
 
 	for name, item := range tree {
 		switch elem := item.(type) {