并行统计文件目录总大小

统计一下给定目录下的文件目录大小

知识点

• WaitGroup保证并行结束
• chan计算共享变量
• Mutex计算共享变量
• 有缓存的chan控制并行协程数
• GOMAXPROCS控制并行cpu数

方式1

package main

import (
    "fmt"
    "io/ioutil"
    "os"
    "path/filepath"
    "runtime"
    "sync"
)

var lock sync.Mutex

func main() {
    runtime.GOMAXPROCS(20)
    if len(os.Args) != 2 { 
        fmt.Println("du_multi root")
        os.Exit(1)
    }   
    root := os.Args[1]
    var fileSize int64
    wg:=&sync.WaitGroup{}
   	nt := make(chan struct{}, 1000)
    wg.Add(1)
    go getFileSize(root, nt, wg, &fileSize)
    wg.Wait()
    fmt.Printf("TotalSize %dG in %s\n", fileSize>>30, root)
}

func getFileSize(dirPath string, nt chan struct{}, wg *sync.WaitGroup, fileSize *int64) {
    defer wg.Done()
    fi,err :=os.Lstat(dirPath)    // 忽略文件夹的符号链接
    if err!=nil {
        fmt.Fprintf(os.Stderr, err.Error())
        return
    }   
    if fi.IsDir() {
        nt <- struct{}{}    // 防止调用句柄次数超出系统限制
        dir, err := ioutil.ReadDir(dirPath)
        <- nt
        if err != nil {
            fmt.Println("Error read dir", err, dirPath)
            os.Exit(1)
        }   
        for _, fi := range dir {
            wg.Add(1)
            go getFileSize(filepath.Join(dirPath, fi.Name()), nt, wg, fileSize)
        }   
    } else {
        lock.Lock()
        (*fileSize)+=fi.Size()
        lock.Unlock()
    }   
}

方式2

• 通道会占用不少内存

package main

import (
    "fmt"
    "io/ioutil"
    "os"
    "path/filepath"
    "runtime"
    "sync"
)

func main() {
    runtime.GOMAXPROCS(20)
    if len(os.Args) != 2 { 
        fmt.Println("du_multi root")
        os.Exit(1)
    }   
    root := os.Args[1]
    var fileSize int64
    wg:=&sync.WaitGroup{}
    fs := make(chan int64)
   	nt := make(chan struct{}, 1000)
    wg.Add(1)
    go getFileSize(root, nt, wg, fs)
    go func() {
    	wg.Wait()
    	close(fs)
	}()
    for s := range fs {
    	fileSize+=s
    }
    fmt.Printf("TotalSize %dG in %s\n", fileSize>>30, root)
}

func getFileSize(dirPath string, nt chan struct{}, wg *sync.WaitGroup, fs chan int64) {
    defer wg.Done()
    fi,err :=os.Lstat(dirPath)    // 忽略文件夹的符号链接
    if err!=nil {
        fmt.Fprintf(os.Stderr, err.Error())
        return
    }   
    if fi.IsDir() {
        nt <- struct{}{}    // 防止调用句柄次数超出系统限制
        dir, err := ioutil.ReadDir(dirPath)
        <- nt
        if err != nil {
            fmt.Println("Error read dir", err, dirPath)
            os.Exit(1)
        }   
        for _, fi := range dir {
            wg.Add(1)
            go getFileSize(filepath.Join(dirPath, fi.Name()), nt, wg, fs)
        }   
    } else {
        fs <- fi.Size()
    }   
}

参考

【golang】并发遍历指定目录下的所有文件大小
golang服务的文件句柄超出系统限制(too many open files)