restic/internal/walk/walk.go

package walk

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"restic"
	"sync"

	"restic/debug"
)

// TreeJob is a job sent from the tree walker.
type TreeJob struct {
	Path  string
	Error error

	Node *restic.Node
	Tree *restic.Tree
}

// TreeWalker traverses a tree in the repository depth-first and sends a job
// for each item (file or dir) that it encounters.
type TreeWalker struct {
	ch  chan<- loadTreeJob
	out chan<- TreeJob
}

// NewTreeWalker uses ch to load trees from the repository and sends jobs to
// out.
func NewTreeWalker(ch chan<- loadTreeJob, out chan<- TreeJob) *TreeWalker {
	return &TreeWalker{ch: ch, out: out}
}

// Walk starts walking the tree given by id. When the channel done is closed,
// processing stops.
func (tw *TreeWalker) Walk(ctx context.Context, path string, id restic.ID) {
	debug.Log("starting on tree %v for %v", id.Str(), path)
	defer debug.Log("done walking tree %v for %v", id.Str(), path)

	resCh := make(chan loadTreeResult, 1)
	tw.ch <- loadTreeJob{
		id:  id,
		res: resCh,
	}

	res := <-resCh
	if res.err != nil {
		select {
		case tw.out <- TreeJob{Path: path, Error: res.err}:
		case <-ctx.Done():
			return
		}
		return
	}

	tw.walk(ctx, path, res.tree)

	select {
	case tw.out <- TreeJob{Path: path, Tree: res.tree}:
	case <-ctx.Done():
		return
	}
}

func (tw *TreeWalker) walk(ctx context.Context, path string, tree *restic.Tree) {
	debug.Log("start on %q", path)
	defer debug.Log("done for %q", path)

	debug.Log("tree %#v", tree)

	// load all subtrees in parallel
	results := make([]<-chan loadTreeResult, len(tree.Nodes))
	for i, node := range tree.Nodes {
		if node.Type == "dir" {
			resCh := make(chan loadTreeResult, 1)
			tw.ch <- loadTreeJob{
				id:  *node.Subtree,
				res: resCh,
			}

			results[i] = resCh
		}
	}

	for i, node := range tree.Nodes {
		p := filepath.Join(path, node.Name)
		var job TreeJob

		if node.Type == "dir" {
			if results[i] == nil {
				panic("result chan should not be nil")
			}

			res := <-results[i]
			if res.err == nil {
				tw.walk(ctx, p, res.tree)
			} else {
				fmt.Fprintf(os.Stderr, "error loading tree: %v\n", res.err)
			}

			job = TreeJob{Path: p, Tree: res.tree, Error: res.err}
		} else {
			job = TreeJob{Path: p, Node: node}
		}

		select {
		case tw.out <- job:
		case <-ctx.Done():
			return
		}
	}
}

type loadTreeResult struct {
	tree *restic.Tree
	err  error
}

type loadTreeJob struct {
	id  restic.ID
	res chan<- loadTreeResult
}

type treeLoader func(restic.ID) (*restic.Tree, error)

func loadTreeWorker(ctx context.Context, wg *sync.WaitGroup, in <-chan loadTreeJob, load treeLoader) {
	debug.Log("start")
	defer debug.Log("exit")
	defer wg.Done()

	for {
		select {
		case <-ctx.Done():
			debug.Log("done channel closed")
			return
		case job, ok := <-in:
			if !ok {
				debug.Log("input channel closed, exiting")
				return
			}

			debug.Log("received job to load tree %v", job.id.Str())
			tree, err := load(job.id)

			debug.Log("tree %v loaded, error %v", job.id.Str(), err)

			select {
			case job.res <- loadTreeResult{tree, err}:
				debug.Log("job result sent")
			case <-ctx.Done():
				debug.Log("done channel closed before result could be sent")
				return
			}
		}
	}
}

// TreeLoader loads tree objects.
type TreeLoader interface {
	LoadTree(context.Context, restic.ID) (*restic.Tree, error)
}

const loadTreeWorkers = 10

// Tree walks the tree specified by id recursively and sends a job for each
// file and directory it finds. When the channel done is closed, processing
// stops.
func Tree(ctx context.Context, repo TreeLoader, id restic.ID, jobCh chan<- TreeJob) {
	debug.Log("start on %v, start workers", id.Str())

	load := func(id restic.ID) (*restic.Tree, error) {
		tree, err := repo.LoadTree(ctx, id)
		if err != nil {
			return nil, err
		}
		return tree, nil
	}

	ch := make(chan loadTreeJob)

	var wg sync.WaitGroup
	for i := 0; i < loadTreeWorkers; i++ {
		wg.Add(1)
		go loadTreeWorker(ctx, &wg, ch, load)
	}

	tw := NewTreeWalker(ch, jobCh)
	tw.Walk(ctx, "", id)
	close(jobCh)

	close(ch)
	wg.Wait()

	debug.Log("done")
}