github.com/mckael/restic@v0.8.3/internal/repository/parallel.go (about)

     1  package repository
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/restic/restic/internal/debug"
     7  	"github.com/restic/restic/internal/restic"
     8  	"golang.org/x/sync/errgroup"
     9  )
    10  
    11  // ParallelWorkFunc gets one file ID to work on. If an error is returned,
    12  // processing stops. When the contect is cancelled the function should return.
    13  type ParallelWorkFunc func(ctx context.Context, id string) error
    14  
    15  // ParallelIDWorkFunc gets one restic.ID to work on. If an error is returned,
    16  // processing stops. When the context is cancelled the function should return.
    17  type ParallelIDWorkFunc func(ctx context.Context, id restic.ID) error
    18  
    19  // FilesInParallel runs n workers of f in parallel, on the IDs that
    20  // repo.List(t) yields. If f returns an error, the process is aborted and the
    21  // first error is returned.
    22  func FilesInParallel(ctx context.Context, repo restic.Lister, t restic.FileType, n int, f ParallelWorkFunc) error {
    23  	g, ctx := errgroup.WithContext(ctx)
    24  
    25  	ch := make(chan string, n)
    26  	g.Go(func() error {
    27  		defer close(ch)
    28  		return repo.List(ctx, t, func(fi restic.FileInfo) error {
    29  			select {
    30  			case <-ctx.Done():
    31  			case ch <- fi.Name:
    32  			}
    33  			return nil
    34  		})
    35  	})
    36  
    37  	for i := 0; i < n; i++ {
    38  		g.Go(func() error {
    39  			for name := range ch {
    40  				err := f(ctx, name)
    41  				if err != nil {
    42  					return err
    43  				}
    44  			}
    45  			return nil
    46  		})
    47  	}
    48  
    49  	return g.Wait()
    50  }
    51  
    52  // ParallelWorkFuncParseID converts a function that takes a restic.ID to a
    53  // function that takes a string. Filenames that do not parse as a restic.ID
    54  // are ignored.
    55  func ParallelWorkFuncParseID(f ParallelIDWorkFunc) ParallelWorkFunc {
    56  	return func(ctx context.Context, s string) error {
    57  		id, err := restic.ParseID(s)
    58  		if err != nil {
    59  			debug.Log("invalid ID %q: %v", id, err)
    60  			return nil
    61  		}
    62  
    63  		return f(ctx, id)
    64  	}
    65  }