pwalkdir.go

Documentation: github.com/opencontainers/selinux/pkg/pwalkdir

     1  //go:build go1.16
     2  // +build go1.16
     3  
     4  package pwalkdir
     5  
     6  import (
     7  	"fmt"
     8  	"io/fs"
     9  	"path/filepath"
    10  	"runtime"
    11  	"sync"
    12  )
    13  
    14  // Walk is a wrapper for filepath.WalkDir which can call multiple walkFn
    15  // in parallel, allowing to handle each item concurrently. A maximum of
    16  // twice the runtime.NumCPU() walkFn will be called at any one time.
    17  // If you want to change the maximum, use WalkN instead.
    18  //
    19  // The order of calls is non-deterministic.
    20  //
    21  // Note that this implementation only supports primitive error handling:
    22  //
    23  // - no errors are ever passed to walkFn;
    24  //
    25  // - once a walkFn returns any error, all further processing stops
    26  // and the error is returned to the caller of Walk;
    27  //
    28  // - filepath.SkipDir is not supported;
    29  //
    30  // - if more than one walkFn instance will return an error, only one
    31  // of such errors will be propagated and returned by Walk, others
    32  // will be silently discarded.
    33  func Walk(root string, walkFn fs.WalkDirFunc) error {
    34  	return WalkN(root, walkFn, runtime.NumCPU()*2)
    35  }
    36  
    37  // WalkN is a wrapper for filepath.WalkDir which can call multiple walkFn
    38  // in parallel, allowing to handle each item concurrently. A maximum of
    39  // num walkFn will be called at any one time.
    40  //
    41  // Please see Walk documentation for caveats of using this function.
    42  func WalkN(root string, walkFn fs.WalkDirFunc, num int) error {
    43  	// make sure limit is sensible
    44  	if num < 1 {
    45  		return fmt.Errorf("walk(%q): num must be > 0", root)
    46  	}
    47  
    48  	files := make(chan *walkArgs, 2*num)
    49  	errCh := make(chan error, 1) // Get the first error, ignore others.
    50  
    51  	// Start walking a tree asap.
    52  	var (
    53  		err error
    54  		wg  sync.WaitGroup
    55  
    56  		rootLen   = len(root)
    57  		rootEntry *walkArgs
    58  	)
    59  	wg.Add(1)
    60  	go func() {
    61  		err = filepath.WalkDir(root, func(p string, entry fs.DirEntry, err error) error {
    62  			if err != nil {
    63  				close(files)
    64  				return err
    65  			}
    66  			if len(p) == rootLen {
    67  				// Root entry is processed separately below.
    68  				rootEntry = &walkArgs{path: p, entry: entry}
    69  				return nil
    70  			}
    71  			// Add a file to the queue unless a callback sent an error.
    72  			select {
    73  			case e := <-errCh:
    74  				close(files)
    75  				return e
    76  			default:
    77  				files <- &walkArgs{path: p, entry: entry}
    78  				return nil
    79  			}
    80  		})
    81  		if err == nil {
    82  			close(files)
    83  		}
    84  		wg.Done()
    85  	}()
    86  
    87  	wg.Add(num)
    88  	for i := 0; i < num; i++ {
    89  		go func() {
    90  			for file := range files {
    91  				if e := walkFn(file.path, file.entry, nil); e != nil {
    92  					select {
    93  					case errCh <- e: // sent ok
    94  					default: // buffer full
    95  					}
    96  				}
    97  			}
    98  			wg.Done()
    99  		}()
   100  	}
   101  
   102  	wg.Wait()
   103  
   104  	if err == nil {
   105  		err = walkFn(rootEntry.path, rootEntry.entry, nil)
   106  	}
   107  
   108  	return err
   109  }
   110  
   111  // walkArgs holds the arguments that were passed to the Walk or WalkN
   112  // functions.
   113  type walkArgs struct {
   114  	entry fs.DirEntry
   115  	path  string
   116  }
   117
View as plain text