walk.go

Documentation: github.com/Shopify/go-storage

     1  package storage
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  // WalkFn is a function type which is passed to Walk.
     9  type WalkFn func(path string) error
    10  
    11  // Walker is an interface which defines the Walk method.
    12  type Walker interface {
    13  	// Walk traverses a path listing by prefix, calling fn with each path.
    14  	Walk(ctx context.Context, path string, fn WalkFn) error
    15  }
    16  
    17  // List runs the Walker on the given path and returns the list of visited paths.
    18  func List(ctx context.Context, w Walker, path string) ([]string, error) {
    19  	var out []string
    20  	if err := w.Walk(ctx, path, func(path string) error {
    21  		out = append(out, path)
    22  
    23  		return nil
    24  	}); err != nil {
    25  		return nil, err
    26  	}
    27  
    28  	return out, nil
    29  }
    30  
    31  // WalkN creates n workers which accept paths from the Walker.  If a WalkFn
    32  // returns non-nil error we wait for other running WalkFns to finish before
    33  // returning.
    34  func WalkN(ctx context.Context, w Walker, path string, n int, fn WalkFn) error {
    35  	errCh := make(chan error, n)
    36  	ch := make(chan string)
    37  	wg := sync.WaitGroup{}
    38  
    39  	for i := 0; i < n; i++ {
    40  		wg.Add(1)
    41  		go func() {
    42  			for f := range ch {
    43  				if err := fn(f); err != nil {
    44  					errCh <- err
    45  
    46  					break
    47  				}
    48  			}
    49  			wg.Done()
    50  		}()
    51  	}
    52  
    53  	err := w.Walk(ctx, path, func(path string) error {
    54  		select {
    55  		case err := <-errCh:
    56  			return err
    57  		case ch <- path:
    58  			return nil
    59  		}
    60  	})
    61  
    62  	close(ch)
    63  	wg.Wait()
    64  	close(errCh)
    65  
    66  	return err
    67  }
    68
View as plain text