Golang的channel使用以及并发同步技巧

在学习《The Go Programming Language》第八章并发单元的时候还是遭遇了不少问题，和值得总结思考和记录的地方。

做一个类似于unix du命令的工具。但是阉割了一些功能，这里应该只实现-c(统计total大小) 和-h(以human比较容易辨识的显示出来)的功能。

首先我们需要构造一个 能够返回FileInfo信息数组的函数，我们把它取名为dirEntries:

func dirEntries(dir string) []os.FileInfo {
    entries, err := ioutil.ReadDir(dir)
    if err != nil {
        fmt.Fprintf(os.Stderr, "du: %v\n", err)
        return nil
    }
    return entries
}

传入一个路径字符串，然后使用ioutil.ReadDir解析这个路径下面的所有文件以及文件夹生成一个FileInfo的profile。

Fileinfo interface下面包含了：

type FileInfo interface {
    Name() string       // base name of the file
    Size() int64        // length in bytes for regular files; system-dependent for others
    Mode() FileMode     // file mode bits
    ModTime() time.Time // modification time
    IsDir() bool        // abbreviation for Mode().IsDir()
    Sys() interface{}   // underlying data source (can return nil)
}

多种方法，可以直接调用，其作用就是后面注释写的一样。

有了能够获取文件夹下面文件和文件夹的函数之后，我们需要一个调用方用来walk指定的目录：

// 入参是一个文件目录，一个INT64的只接收的单向channel
func walkDir(dir string, fileSizes chan<- int64) {
    for _, entry := range dirEntries(dir) {
        if entry.IsDir() {
            subdir := filepath.Join(dir, entry.Name())
            walkDir(subdir, fileSizes)
        } else {
            fileSizes <- entry.Size()
        }
    }
}

这里我们定义一个目录，然后需求传入一个单向接收channel用于在多goroutine中计算总共的文件大小。

使用range方法来遍历我们上面写的dirEntries的返回文件或文件夹，如果是文件夹则继续迭代。

如果不是则将文件大小存入放入fileSizes channel中。

搞定上面两个函数，我们来写主函数部分：

func main() {
    root := ""
    flag.StringVar(&root, "-p", ".", "input dir.")
    flag.Parse()

    fileSizes := make(chan int64)
    // 起一个goroutine去walk目录
    go func() {
        walkDir(root, fileSizes)
        // Walk完毕之后要关闭该channel下面使用range读取数据的时候才会有尽头
        close(fileSizes)
    }()

    var nfiles, nbytes int64
    for size := range fileSizes {
        nfiles++
        nbytes += size
    }
    fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
}

这里注意一点，因为起goroutine的walk函数，和下面同时在range遍历是在同步进行，如果下面range速度太快读到管道里面没有值了会阻塞住等待有数据继续进来之后读取，而不是会跳出。只有当close(fileSizes)这句执行到，显示关闭掉channel之后，才会跳出range循环并且这时已经读取完了所有的数据。这里有点像，close channel的时候给range发送了一个停止信号一样，感觉这个利用起来会比较有用？ 后续可能会再研究一下。

让我们继续来优化我们的程序，添加一个-v参数，打印出扫描文件的进度，当我们要扫描整个盘的时候，可能会花费大量的时间，我们需要知道进度如何了。

其实这个需求只需要很小的改动，让我们来重新改写一下main函数，用select多路复用来完成这个事情。

func main() {
    root := ""
    verbose := false
    tick := make(<-chan time.Time)
    var nfiles, nbytes int64

    flag.StringVar(&root, "p", ".", "input dir.")
    flag.BoolVar(&verbose, "v", false, "add verbose if you want")
    flag.Parse()

    if verbose {
        tick = time.Tick(500 * time.Millisecond)
    }

    fileSizes := make(chan int64)
    // 起一个goroutine去walk目录
    go func() {
        walkDir(root, fileSizes)
        // Walk完毕之后要关闭该channel下面使用range读取数据的时候才会有尽头
        close(fileSizes)
    }()

loop:
    for {
        select {
        case size, ok := <-fileSizes:
            if !ok {
                break loop
            }
            nfiles++
            nbytes += size
        case <-tick:
            fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
        }
    }
    fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
}

上面其实都差不多，这里我直接从loop那里开始说吧，遇到这个loop的时候我其实还蛮疑惑的，因为我在go语言保留关键字里面并没有看到他的身影，但是这里他的确是个关键字，和里面的break连用 里面break后面跟上的loop 可以直接跳出到最外层loop包裹的循环，而不是break默认的只跳出一层循环。明白了这个道理之后，这个就不难理解了，当我们还在遍历文件的时候，select 会持续读取文件大小赋值给size，并且返回true给ok。如果我们开启了verbose，每隔500毫秒tick会收到来自time.Tick的消息。我们都知道select会在都准备好的情况下随机pick一个执行，所以这里也或快或慢的被打印进度(前提是同时收到信号，但是实际上这个发生速度可能在nm级别，凭感受很难感觉到谁先)。当最后都执行完毕后filesSizes channel会被上面的携程函数close()，当close之后，在读取完剩余数据后，fileSizes会返回给ok nil。就可以跳出循环。

看到这里可能会觉得有点绕，所以要尽可能的多理解一下，当然我们可以让这个du程序更快。可以注意到我们并没有在walkdir里面开启goroutines进行并发处理。下面我将尝试开启goroutine处理它们，并且用channel给他们加个锁控制一下goroutine的数量，在此之前我们先来看看现在完成了的代码：

package main

import (
    "fmt"
    "io/ioutil"
    "os"
    "path/filepath"
    "flag"
    "time"
)

// 入参是一个文件目录，一个INT64的只接收的单向channel
func walkDir(dir string, fileSizes chan<- int64) {
    for _, entry := range dirEntries(dir) {
        if entry.IsDir() {
            subdir := filepath.Join(dir, entry.Name())
            walkDir(subdir, fileSizes)
        } else {
            fileSizes <- entry.Size()
        }
    }
}

func dirEntries(dir string) []os.FileInfo {
    entries, err := ioutil.ReadDir(dir)
    if err != nil {
        fmt.Fprintf(os.Stderr, "du: %v\n", err)
        return nil
    }
    return entries
}

func main() {
    t1 := time.Now()
    root := ""
    verbose := false
    tick := make(<-chan time.Time)
    var nfiles, nbytes int64

    flag.StringVar(&root, "p", ".", "input dir.")
    flag.BoolVar(&verbose, "v", false, "add verbose if you want")
    flag.Parse()

    if verbose {
        tick = time.Tick(500 * time.Millisecond)
    }

    fileSizes := make(chan int64)
    // 起一个goroutine去walk目录
    go func() {
        walkDir(root, fileSizes)
        // Walk完毕之后要关闭该channel下面使用range读取数据的时候才会有尽头
        close(fileSizes)
    }()

loop:
    for {
        select {
        case size, ok := <-fileSizes:
            if !ok {
                break loop
            }
            nfiles++
            nbytes += size
        case <-tick:
            fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
        }
    }
    fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
    fmt.Println(time.Since(t1))
}

观察上面代码可以看出我们并不能直接在这个代码的基础上直接给walkDir加上goroutine，这样会导致channel直接被关闭，然后啥也没跑就结束了。

我们需要让主goroutine等待其他goroutine都完成之后再结束，所以主goroutine需要在这里阻塞住，等到得到可以结束的信号之后再结束。

我们可以使用sync.WaitGroup 来对仍旧活跃的walkDir调用进行计数。等到数量为0的时候就算我们可以结束了。

sync.WaitGroup提供了三个方法：

　　Add：添加或减少goroutine的数量。

　　Done：相当于Add(-1)。

　　Wait：阻塞住等待WaitGroup数量变成0.

明白这个道理之后我们改写了一下代码，让它使用sync.WaitGroup来支持同步，最后当所有goroutine都结束之后，关闭channel完成任务。

package main

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

// 入参是一个文件目录，一个INT64的只接收的单向channel
func walkDir(dir string, fileSizes chan<- int64, n *sync.WaitGroup) {
    defer n.Done()
    for _, entry := range dirEntries(dir) {
        if entry.IsDir() {
            n.Add(1)
            subdir := filepath.Join(dir, entry.Name())
            go walkDir(subdir, fileSizes, n)
        } else {
            fileSizes <- entry.Size()
        }
    }
}

func dirEntries(dir string) []os.FileInfo {
    entries, err := ioutil.ReadDir(dir)
    if err != nil {
        fmt.Fprintf(os.Stderr, "du: %v\n", err)
        return nil
    }
    return entries
}

func main() {
    t1 := time.Now()
    root := ""
    verbose := false
    tick := make(<-chan time.Time)
    fileSizes := make(chan int64)

    var n sync.WaitGroup
    var nfiles, nbytes int64

    flag.StringVar(&root, "p", ".", "input dir.")
    flag.BoolVar(&verbose, "v", false, "add verbose if you want")
    flag.Parse()

    if verbose {
        tick = time.Tick(500 * time.Millisecond)
    }

    n.Add(1)
    go walkDir(root, fileSizes, &n)

    go func() {
        n.Wait()
        close(fileSizes)
    }()

loop:
    for {
        select {
        case size, ok := <-fileSizes:
            if !ok {
                break loop
            }
            nfiles++
            nbytes += size
        case <-tick:
            fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
        }
    }
    fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
    fmt.Println(time.Since(t1))
}

随便跑跑。。感觉快得飞起，然而跑不了几秒就会报错，这个程序最大的问题就是我们完全没有办法之后它会自己打开多少个goroutine，感觉会爆炸。所以我们要限制这种夸张的写法，使用channel来做一个并发协程池，把同时开启的goroutine的数量控制一下。

最后上一下完整代码，注意defer关键字，只接收函数，所以我会在释放锁的时候使用匿名函数：

package main

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

var token = make(chan int, 100)

// 入参是一个文件目录，一个INT64的只接收的单向channel
func walkDir(dir string, fileSizes chan<- int64, n *sync.WaitGroup) {
    defer n.Done()
    for _, entry := range dirEntries(dir) {
        if entry.IsDir() {
            n.Add(1)
            subdir := filepath.Join(dir, entry.Name())
            go walkDir(subdir, fileSizes, n)
        } else {
            fileSizes <- entry.Size()
        }
    }
}

func dirEntries(dir string) []os.FileInfo {
    token <- 1
    defer func() {<-token}()
    entries, err := ioutil.ReadDir(dir)
    if err != nil {
        fmt.Fprintf(os.Stderr, "du: %v\n", err)
        return nil
    }
    return entries
}

func main() {
    var nfiles, nbytes int64
    var n sync.WaitGroup

    root := ""
    verbose := false
    t1 := time.Now()
    fileSizes := make(chan int64)
    tick := make(<-chan time.Time)

    flag.StringVar(&root, "p", ".", "input dir.")
    flag.BoolVar(&verbose, "v", false, "add verbose if you want")
    flag.Parse()

    if verbose {
        tick = time.Tick(500 * time.Millisecond)
    }

    n.Add(1)
    go walkDir(root, fileSizes, &n)

    go func() {
        n.Wait()
        close(fileSizes)
    }()

loop:
    for {
        select {
        case size, ok := <-fileSizes:
            if !ok {
                break loop
            }
            nfiles++
            nbytes += size
        case <-tick:
            fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
        }
    }
    fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
    fmt.Println(time.Since(t1))
}

Reference:

https://github.com/gopl-zh/gopl-zh.github.com  The Go Programming Language