github.com/fawick/restic@v0.1.1-0.20171126184616-c02923fbfc79/internal/repository/parallel.go (about)

     1  package repository
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/restic/restic/internal/debug"
     8  	"github.com/restic/restic/internal/restic"
     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) yield. 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 uint, f ParallelWorkFunc) error {
    23  	wg := &sync.WaitGroup{}
    24  	ch := repo.List(ctx, t)
    25  	errors := make(chan error, n)
    26  
    27  	for i := 0; uint(i) < n; i++ {
    28  		wg.Add(1)
    29  		go func() {
    30  			defer wg.Done()
    31  
    32  			for {
    33  				select {
    34  				case id, ok := <-ch:
    35  					if !ok {
    36  						return
    37  					}
    38  
    39  					err := f(ctx, id)
    40  					if err != nil {
    41  						errors <- err
    42  						return
    43  					}
    44  				case <-ctx.Done():
    45  					return
    46  				}
    47  			}
    48  		}()
    49  	}
    50  
    51  	wg.Wait()
    52  
    53  	select {
    54  	case err := <-errors:
    55  		return err
    56  	default:
    57  		break
    58  	}
    59  
    60  	return nil
    61  }
    62  
    63  // ParallelWorkFuncParseID converts a function that takes a restic.ID to a
    64  // function that takes a string. Filenames that do not parse as a restic.ID
    65  // are ignored.
    66  func ParallelWorkFuncParseID(f ParallelIDWorkFunc) ParallelWorkFunc {
    67  	return func(ctx context.Context, s string) error {
    68  		id, err := restic.ParseID(s)
    69  		if err != nil {
    70  			debug.Log("invalid ID %q: %v", id, err)
    71  			return nil
    72  		}
    73  
    74  		return f(ctx, id)
    75  	}
    76  }