scanner.go

/*
   High performance filesystem scanner
   (c) Parul Jain, 2016
   [email protected]
   MIT License
*/

package main

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"sync"
	"sync/atomic"
)

func scan(path string) (fileCount uint64) {
	const MaxThreads = 5000
	//Create a throttling mechanism
	//Buffered channel
	tChan := make(chan int, MaxThreads)
	//Load up MaxThreads "tokens" into the channel
	for i := 0; i < MaxThreads; i++ {
		tChan <- 1
	}
	//Create a sync mechanism to wait for all threads to complete
	wg := &sync.WaitGroup{}
	//Create the func that will do the work
	//Declaring a var for the func is needed to use it recursively
	var doFolder func(string)
	doFolder = func(p string) {
		//Defer makes sure that cleanup is called no matter how we exit the func
		defer func() {
			//Return token
			tChan <- 1
			//Done
			wg.Done()
		}()
		items, err := ioutil.ReadDir(p)
		if err != nil {
			fmt.Println(err)
			return
		}
		for _, item := range items {
			if item.IsDir() {
				//Add to wait group so we can track when the scan is done
				wg.Add(1)
				//Wait for token to be available
				<-tChan
				//Start the scan for the folder on a new thread
				go doFolder(filepath.Join(p, item.Name()))
			} else {
				//Process the file
				atomic.AddUint64(&fileCount, 1)
			}
		}
	}
	//Add to wait group so we can track when the scan is done
	wg.Add(1)
	//Wait for token to be available
	<-tChan
	go doFolder(path)

	//Wait for all go routines to finish
	wg.Wait()
	return
}

func main() {
	fmt.Println(scan(`e:\`))
}