golang练手小项目系列(3)-并发读取文件夹信息
问题描述：
实现一个程序计算并打印输入的目录下所有文件的总数和总大小(以GB计算)。完成之后你将熟悉select、WaitGroup、ioutil的用法。
要点：
并发读取文件(夹)信息。
限制开启的goroutines的最大数量。
运行时每隔500ms打印当前已经统计的文件数和总大小（使用命令行参数指定此功能是否启用）。
拓展：
在执行中在有外部输入时退出程序。
实现：
import (
"flag"
"fmt"
"io/ioutil"
"os"
"path/filepath"
"sync"
"time"
)
var verbose = flag.Bool("v", false, "show verbose progress messages")
var sema = make(chan struct{}, 50)
var done = make(chan struct{})
func dirents(dir string) []os.FileInfo{
select {
case sema <- struct{}{}: //
case <- done:
return nil
}
defer func() {<- sema}()
entries, err := ioutil.ReadDir(dir)
if err != nil{
fmt.Fprintf(os.Stderr, "du1: %v\n", err)
return nil
}
return entries
}
func walkDir(dir string, n *sync.WaitGroup, fileSizes chan<- int64){
defer n.Done()
if cancelled(){
return
}
for _, entry := range dirents(dir){
if entry.IsDir(){
n.Add(1)
subdir := filepath.Join(dir, entry.Name())
go walkDir(subdir, n, fileSizes)
} else {
fileSizes <- entry.Size()
}
}
}
func printDiskUsage(nfiles, nbytes int64){
fmt.Printf("%d files %.1f GB\n", nfiles, float64(nbytes/1e9))
}
func cancelled() bool{
select {
case <- done:
return true
default:
return false
}
}
func main() {
flag.Parse()
roots := flag.Args()
var tick <-chan time.Time
if *verbose{
tick = time.Tick(500 * time.Millisecond)
}
if len(roots) == 0{
roots = []string{"."}
}
fileSizes := make(chan int64)
var nfiles, nbytes int64
var n sync.WaitGroup
for _, root := range roots{
n.Add(1)
go walkDir(root, &n, fileSizes)
}
go func() {
n.Wait()
close(fileSizes)
}()
go func() {
os.Stdin.Read(make([]byte, 1))
close(done)
}()
loop:
for {
select {
case <-done:
for range fileSizes{
//
}
case size, ok := <- fileSizes:
if !ok {
break loop
}
nfiles++
nbytes += size
case <- tick:
printDiskUsage(nfiles, nbytes)
}
}
printDiskUsage(nfiles, nbytes)
}
