wings/server/filesystem_walker.go

package server

import (
	"github.com/gammazero/workerpool"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"sync"
)

type FileWalker struct {
	*Filesystem
}

type PooledFileWalker struct {
	wg       sync.WaitGroup
	pool     *workerpool.WorkerPool
	callback filepath.WalkFunc

	Filesystem *Filesystem
}

// Returns a new walker instance.
func (fs *Filesystem) NewWalker() *FileWalker {
	return &FileWalker{fs}
}

// Creates a new pooled file walker that will concurrently walk over a given directory but limit itself
// to a worker pool as to not completely flood out the system or cause a process crash.
func newPooledWalker(fs *Filesystem) *PooledFileWalker {
	return &PooledFileWalker{
		Filesystem: fs,
		// Create a worker pool that is the same size as the number of processors available on the
		// system. Going much higher doesn't provide much of a performance boost, and is only more
		// likely to lead to resource overloading anyways.
		pool: workerpool.New(runtime.GOMAXPROCS(0)),
	}
}

// Process a given path by calling the callback function for all of the files and directories within
// the path, and then dropping into any directories that we come across.
func (w *PooledFileWalker) process(path string) error {
	p, err := w.Filesystem.SafePath(path)
	if err != nil {
		return err
	}

	files, err := ioutil.ReadDir(p)
	if err != nil {
		return err
	}

	// Loop over all of the files and directories in the given directory and call the provided
	// callback function. If we encounter a directory, push that directory onto the worker queue
	// to be processed.
	for _, f := range files {
		sp := filepath.Join(p, f.Name())
		i, err := os.Stat(sp)

		if err = w.callback(sp, i, err); err != nil {
			if err == filepath.SkipDir {
				return nil
			}

			return err
		}

		if i.IsDir() {
			w.push(sp)
		}
	}

	return nil
}

// Push a new path into the worker pool.
//
// @todo probably helps to handle errors.
func (w *PooledFileWalker) push(path string) {
	w.wg.Add(1)
	w.pool.Submit(func() {
		defer w.wg.Done()
		w.process(path)
	})
}

// Walks the given directory and executes the callback function for all of the files and directories
// that are encountered.
func (fs *Filesystem) Walk(dir string, callback filepath.WalkFunc) error {
	w := newPooledWalker(fs)
	w.callback = callback

	w.push(dir)

	w.wg.Wait()
	w.pool.StopWait()

	return nil
}