feat(snapshots): big upload performance improvements by parallelizing directory traversal (#1752)

* feat(general): added internal/workshare package

This introduces work sharing utility useful when walking trees of
things (such as filesystem), which allows N threads/goroutines to be
used.

Whenever a routine is visiting its children, it can share some of that
work with another idle goroutine in the pool (when available). If
no other goroutine is idle, we are already at capacity and the caller
simply does the work in their own goroutine.

The API introduced here is not the most beautiful, but allows us to
avoid allocations in most cases, which is critical for high-performance
data processing.

* feat(snapshots): speed up uploads by parallelizing directory traversal

Previously directories were walked strictly sequentially which means
we could never be uploading data from multiple directories in parallel,
even if they had just a few files each.

This change switches to using the new `workshare` utility which improves
parallelism. It also reduces memory allocations, goroutine creations
and overall memory usage when taking large snapshots, while increasing
CPU utilization.

Tests on realistic directory structures show huge speed-ups during cold
snapshots (without any metadata caching:)

Photo library - 160GB, files:41717 dirs:1350

    Before: 3m11s
    After: 1m50s
    Total time reduction: 43%

Working code directory - 30.7 GB files:194560 dirs:42455

    Before: 55s
    After: 25s
    Total time reduction: 55%

* do not report multiple cancelation errors during parallel uploads

* do not report multiple cancelation errors during parallel uploads

* pr feedback, clarified usage, added comments

* fixed flaky test

internal/workshare/workshare_pool.go

@@ -0,0 +1,80 @@
+// Package workshare implements work sharing worker pool.
+package workshare
+
+import (
+	"sync"
+	"sync/atomic"
+)
+
+// ProcessFunc processes the provided request, which is typically a pointer to a structure.
+// The result will typically be written within the request structure. To avoid allocations,
+// the function should not be a local closure but a named function (either global or a method).
+type ProcessFunc func(c *Pool, request interface{})
+
+type workItem struct {
+	process ProcessFunc     // function to call
+	request interface{}     // parameter to the function
+	wg      *sync.WaitGroup // signal the provided WaitGroup after work is done
+}
+
+// Pool manages a pool of generic workers that can process workItem.
+type Pool struct {
+	activeWorkers int32
+
+	semaphore chan struct{}
+
+	work   chan workItem
+	closed chan struct{}
+
+	wg sync.WaitGroup
+}
+
+// ActiveWorkers returns the number of active workers.
+func (w *Pool) ActiveWorkers() int {
+	return int(atomic.LoadInt32(&w.activeWorkers))
+}
+
+// NewPool creates a worker pool that launches a given number of goroutines that can invoke shared work.
+func NewPool(numWorkers int) *Pool {
+	if numWorkers < 0 {
+		numWorkers = 0
+	}
+
+	w := &Pool{
+		// channel must be unbuffered so that it has exactly as many slots as there are goroutines capable of reading from it
+		// this way by pushing to the channel we can be sure that a pre-spun goroutine will pick it up soon.
+		work:      make(chan workItem),
+		closed:    make(chan struct{}),
+		semaphore: make(chan struct{}, numWorkers),
+	}
+
+	for i := 0; i < numWorkers; i++ {
+		w.wg.Add(1)
+
+		go func() {
+			defer w.wg.Done()
+
+			for {
+				select {
+				case it := <-w.work:
+					atomic.AddInt32(&w.activeWorkers, 1)
+					it.process(w, it.request)
+					atomic.AddInt32(&w.activeWorkers, -1)
+					<-w.semaphore
+					it.wg.Done()
+
+				case <-w.closed:
+					return
+				}
+			}
+		}()
+	}
+
+	return w
+}
+
+// Close closes the worker pool.
+func (w *Pool) Close() {
+	close(w.closed)
+	w.wg.Wait()
+}

internal/workshare/workshare_test.go

@@ -0,0 +1,128 @@
+package workshare_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/kopia/kopia/internal/workshare"
+)
+
+type treeNode struct {
+	value    int
+	children []*treeNode
+}
+
+func buildTree(level int) *treeNode {
+	n := &treeNode{
+		value: 1,
+	}
+
+	if level <= 0 {
+		return n
+	}
+
+	for i := 0; i < level; i++ {
+		n.children = append(n.children, buildTree(level-1))
+	}
+
+	return n
+}
+
+type computeTreeSumRequest struct {
+	input *treeNode
+
+	result int
+	err    error
+}
+
+func dispatchComputeTreeSumRequest(w *workshare.Pool, input interface{}) {
+	req := input.(*computeTreeSumRequest)
+
+	if w.ActiveWorkers() == 0 {
+		panic("unexpected worker count")
+	}
+
+	res, err := computeTreeSum(w, req.input)
+	if err != nil {
+		req.err = err
+		return
+	}
+
+	req.result = res
+}
+
+func computeTreeSum(workPool *workshare.Pool, n *treeNode) (int, error) {
+	total := n.value
+
+	var cs workshare.AsyncGroup
+
+	for _, child := range n.children {
+		if cs.CanShareWork(workPool) {
+			// run the request on another goroutine, the results will be available
+			cs.RunAsync(workPool, dispatchComputeTreeSumRequest, &computeTreeSumRequest{
+				input: child,
+			})
+		} else {
+			chtot, err := computeTreeSum(workPool, child)
+			if err != nil {
+				return 0, err
+			}
+
+			total += chtot
+		}
+	}
+
+	for _, req := range cs.Wait() {
+		twr := req.(*computeTreeSumRequest)
+
+		if twr.err != nil {
+			return 0, twr.err
+		}
+
+		total += twr.result
+	}
+
+	return total, nil
+}
+
+func TestComputeTreeSum10(t *testing.T) {
+	testComputeTreeSum(t, 10)
+}
+
+func TestComputeTreeSum1(t *testing.T) {
+	testComputeTreeSum(t, 1)
+}
+
+func TestComputeTreeSum0(t *testing.T) {
+	testComputeTreeSum(t, 0)
+}
+
+func TestComputeTreeSumNegative(t *testing.T) {
+	testComputeTreeSum(t, -1)
+}
+
+// nolint:thelper
+func testComputeTreeSum(t *testing.T, numWorkers int) {
+	w := workshare.NewPool(numWorkers)
+	defer w.Close()
+
+	n := buildTree(6)
+
+	sum, err := computeTreeSum(w, n)
+	require.NoError(t, err)
+	require.Equal(t, 1957, sum)
+}
+
+var treeToWalk = buildTree(6)
+
+func BenchmarkComputeTreeSum(b *testing.B) {
+	w := workshare.NewPool(10)
+	defer w.Close()
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		computeTreeSum(w, treeToWalk)
+	}
+}

internal/workshare/workshare_waitgroup.go

@@ -0,0 +1,101 @@
+// Package workshare implements work sharing worker pool.
+//
+// It is commonly used to traverse tree-like structures:
+//
+// type processWorkRequest struct {
+//   node *someNode
+//   err error
+// }
+//
+// func processWork(p *workshare.Pool, req interface{}) {
+//   req := req.(*processWorkRequest)
+//   req.err = visitNode(p, req.node)
+// }
+//
+// func visitNode(p *workshare.Pool, n *someNode) error {
+//   var wg workshare.AsyncGroup
+//
+//   for _, child := range n.children {
+//     if wg.CanShareWork(p) {
+//       // run asynchronously, collect result using wg.Wait() below.
+//       RunAsync(p, processWork, processWorkRequest{child})
+//     } else {
+// 	      if err := visitNode(p, n); err != nil {
+//          return err
+//        }
+//     }
+//   }
+//
+//   // wait for results from all shared work and handle them.
+//   for _, req := range wg.Wait() {
+//     if err := req.(*processWorkRequest).err; err != nil {
+//       return err
+//     }
+//   }
+//
+//   return nil
+// }
+//
+// wp = workshare.NewPool(10)
+// defer wp.Close()
+// visitNode(wp, root)
+package workshare
+
+import (
+	"sync"
+)
+
+// AsyncGroup launches and awaits asynchronous work through a WorkerPool.
+// It provides API designed to minimize allocations while being reasonably easy to use.
+// AsyncGroup is very lightweight and NOT safe for concurrent use, all interactions must
+// be from the same goroutine.
+type AsyncGroup struct {
+	wg       *sync.WaitGroup
+	requests []interface{}
+}
+
+// Wait waits for scheduled asynchronous work to complete and returns all asynchronously processed inputs.
+func (g *AsyncGroup) Wait() []interface{} {
+	if g.wg == nil {
+		return nil
+	}
+
+	g.wg.Wait()
+
+	return g.requests
+}
+
+// RunAsync starts the asynchronous work to process the provided request,
+// the user must call Wait() after all RunAsync() have been scheduled.
+func (g *AsyncGroup) RunAsync(w *Pool, process ProcessFunc, request interface{}) {
+	if g.wg == nil {
+		g.wg = &sync.WaitGroup{}
+	}
+
+	g.wg.Add(1)
+
+	g.requests = append(g.requests, request)
+
+	w.work <- workItem{
+		process: process,
+		request: request,
+		wg:      g.wg,
+	}
+}
+
+// CanShareWork determines if the provided worker pool has capacity to share work.
+// If the function returns true, the use MUST call RunAsync() exactly once. This pattern avoids
+// allocations required to create asynchronous input if the worker pool is full.
+func (g *AsyncGroup) CanShareWork(w *Pool) bool {
+	select {
+	case w.semaphore <- struct{}{}:
+		// we successfully added token to the channel, because we have exactly the same number
+		// of workers as the capacity of the channel, one worker will wake up to process
+		// item from w.work, which will be added by RunAsync().
+		return true
+
+	default:
+		// all workers are busy.
+		return false
+	}
+}

snapshot/snapshotfs/upload.go

@@ -16,13 +16,13 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
-	"golang.org/x/sync/errgroup"
 
 	"github.com/kopia/kopia/fs"
 	"github.com/kopia/kopia/fs/ignorefs"
 	"github.com/kopia/kopia/internal/clock"
 	"github.com/kopia/kopia/internal/iocopy"
 	"github.com/kopia/kopia/internal/timetrack"
+	"github.com/kopia/kopia/internal/workshare"
 	"github.com/kopia/kopia/repo"
 	"github.com/kopia/kopia/repo/logging"
 	"github.com/kopia/kopia/repo/object"
@@ -104,6 +104,8 @@ type Uploader struct {
 
 	// disable snapshot size estimation
 	disableEstimation bool
+
+	workerPool *workshare.Pool
 }
 
 // IsCanceled returns true if the upload is canceled.
@@ -497,68 +499,33 @@ func (u *Uploader) uploadDirWithCheckpointing(ctx context.Context, rootDir fs.Di
 	return uploadDirInternal(ctx, u, rootDir, policyTree, previousDirs, localDirPathOrEmpty, ".", &dmb, &cp)
 }
 
-func (u *Uploader) foreachEntryUnlessCanceled(ctx context.Context, parallel int, relativePath string, entries fs.Entries, cb func(ctx context.Context, entry fs.Entry, entryRelativePath string) error) error {
-	if parallel > len(entries) {
-		// don't launch more goroutines than needed
-		parallel = len(entries)
-	}
+type uploadWorkItem struct {
+	err error
+}
 
-	if parallel == 0 {
-		return nil
-	}
+func (u *Uploader) foreachEntryUnlessCanceled(ctx context.Context, wg *workshare.AsyncGroup, relativePath string, entries fs.Entries, cb func(ctx context.Context, entry fs.Entry, entryRelativePath string) error) error {
+	for _, entry := range entries {
+		entry := entry
 
-	if parallel == 1 {
-		for _, entry := range entries {
-			if u.IsCanceled() {
-				return errCanceled
-			}
+		if u.IsCanceled() {
+			return errCanceled
+		}
 
-			entryRelativePath := path.Join(relativePath, entry.Name())
+		entryRelativePath := path.Join(relativePath, entry.Name())
+
+		if wg.CanShareWork(u.workerPool) {
+			wg.RunAsync(u.workerPool, func(c *workshare.Pool, input interface{}) {
+				wi, _ := input.(*uploadWorkItem)
+				wi.err = cb(ctx, entry, entryRelativePath)
+			}, &uploadWorkItem{})
+		} else {
 			if err := cb(ctx, entry, entryRelativePath); err != nil {
 				return err
 			}
 		}
-
-		return nil
 	}
 
-	ch := make(chan fs.Entry)
-	eg, ctx := errgroup.WithContext(ctx)
-
-	// one goroutine to pump entries into channel until ctx is closed.
-	eg.Go(func() error {
-		defer close(ch)
-
-		for _, e := range entries {
-			select {
-			case ch <- e: // sent to channel
-			case <-ctx.Done(): // context closed
-				return nil
-			}
-		}
-		return nil
-	})
-
-	// launch N workers in parallel
-	for i := 0; i < parallel; i++ {
-		eg.Go(func() error {
-			for entry := range ch {
-				if u.IsCanceled() {
-					return errCanceled
-				}
-
-				entryRelativePath := path.Join(relativePath, entry.Name())
-				if err := cb(ctx, entry, entryRelativePath); err != nil {
-					return err
-				}
-			}
-
-			return nil
-		})
-	}
-
-	// nolint:wrapcheck
-	return eg.Wait()
+	return nil
 }
 
 func rootCauseError(err error) error {
@@ -698,14 +665,34 @@ func (u *Uploader) processChildren(
 	policyTree *policy.Tree,
 	previousEntries []fs.Entries,
 ) error {
-	if err := u.processSubdirectories(ctx, parentDirCheckpointRegistry, parentDirBuilder, localDirPathOrEmpty, relativePath, entries, policyTree, previousEntries); err != nil {
+	var wg workshare.AsyncGroup
+
+	// ignore errCancel because a more serious error may be reported in wg.Wait()
+	// we'll check for cancelation later.
+
+	if err := u.processSubdirectories(ctx, parentDirCheckpointRegistry, parentDirBuilder, localDirPathOrEmpty, relativePath, entries, policyTree, previousEntries, &wg); err != nil && !errors.Is(err, errCanceled) {
 		return errors.Wrap(err, "processing subdirectories")
 	}
 
-	if err := u.processNonDirectories(ctx, parentDirCheckpointRegistry, parentDirBuilder, relativePath, entries, policyTree, previousEntries); err != nil {
+	if err := u.processNonDirectories(ctx, parentDirCheckpointRegistry, parentDirBuilder, relativePath, entries, policyTree, previousEntries, &wg); err != nil && !errors.Is(err, errCanceled) {
 		return errors.Wrap(err, "processing non-directories")
 	}
 
+	for _, wi := range wg.Wait() {
+		wi, ok := wi.(*uploadWorkItem)
+		if !ok {
+			return errors.Errorf("unexpected work item type %T", wi)
+		}
+
+		if wi.err != nil {
+			return wi.err
+		}
+	}
+
+	if u.IsCanceled() {
+		return errCanceled
+	}
+
 	return nil
 }
 
@@ -717,12 +704,9 @@ func (u *Uploader) processSubdirectories(
 	entries fs.Entries,
 	policyTree *policy.Tree,
 	previousEntries []fs.Entries,
+	wg *workshare.AsyncGroup,
 ) error {
-	// for now don't process subdirectories in parallel, we need a mechanism to
-	// prevent explosion of parallelism
-	const parallelism = 1
-
-	return u.foreachEntryUnlessCanceled(ctx, parallelism, relativePath, entries, func(ctx context.Context, entry fs.Entry, entryRelativePath string) error {
+	return u.foreachEntryUnlessCanceled(ctx, wg, relativePath, entries, func(ctx context.Context, entry fs.Entry, entryRelativePath string) error {
 		dir, ok := entry.(fs.Directory)
 		if !ok {
 			// skip non-directories
@@ -844,7 +828,16 @@ func (u *Uploader) effectiveParallelUploads() int {
 }
 
 // nolint:funlen
-func (u *Uploader) processNonDirectories(ctx context.Context, parentCheckpointRegistry *checkpointRegistry, parentDirBuilder *dirManifestBuilder, dirRelativePath string, entries fs.Entries, policyTree *policy.Tree, prevEntries []fs.Entries) error {
+func (u *Uploader) processNonDirectories(
+	ctx context.Context,
+	parentCheckpointRegistry *checkpointRegistry,
+	parentDirBuilder *dirManifestBuilder,
+	dirRelativePath string,
+	entries fs.Entries,
+	policyTree *policy.Tree,
+	prevEntries []fs.Entries,
+	wg *workshare.AsyncGroup,
+) error {
 	workerCount := u.effectiveParallelUploads()
 
 	var asyncWritesPerFile int
@@ -856,11 +849,9 @@ func (u *Uploader) processNonDirectories(ctx context.Context, parentCheckpointRe
 				asyncWritesPerFile = 0
 			}
 		}
-
-		workerCount = len(entries)
 	}
 
-	return u.foreachEntryUnlessCanceled(ctx, workerCount, dirRelativePath, entries, func(ctx context.Context, entry fs.Entry, entryRelativePath string) error {
+	return u.foreachEntryUnlessCanceled(ctx, wg, dirRelativePath, entries, func(ctx context.Context, entry fs.Entry, entryRelativePath string) error {
 		// note this function runs in parallel and updates 'u.stats', which must be done using atomic operations.
 		if _, ok := entry.(fs.Directory); ok {
 			// skip directories
@@ -1213,6 +1204,11 @@ func (u *Uploader) writeDirManifest(ctx context.Context, dirRelativePath string,
 }
 
 func (u *Uploader) reportErrorAndMaybeCancel(err error, isIgnored bool, dmb *dirManifestBuilder, entryRelativePath string) {
+	if u.IsCanceled() && errors.Is(err, errCanceled) {
+		// alrady canceled, do not report another.
+		return
+	}
+
 	if isIgnored {
 		atomic.AddInt32(&u.stats.IgnoredErrorCount, 1)
 	} else {
@@ -1280,6 +1276,9 @@ func (u *Uploader) Upload(
 
 	defer u.Progress.UploadFinished()
 
+	u.workerPool = workshare.NewPool(u.effectiveParallelUploads() - 1)
+	defer u.workerPool.Close()
+
 	u.stats = &snapshot.Stats{}
 	u.totalWrittenBytes = 0