golang任务拆分errgroup

ErrGroup是 Go 官方提供的一个同步扩展库。可以将一个大任务拆分成几个小任务并发执行，提高程序效率。

主要有三个方法，WithContext、Go、Wait。

func WithContext(ctx context.Context) (*Group, context.Context)

WithContext,返回一个Group实例以及一个Context。如果有一个子任务返回错误，或者Wait调用返回，这个Context就会cancel。

func (g *Group) Go(f func() error)

Go，用于传入子任务，如果成功返回nil，如果失败返回error，同时cancel那个Context

func (g *Group) Wait() error

Wait,类似waitgroup，等所有的子任务完成后返回，如果有多个子任务返回error，则会返回第一个error，所有子任务执行成功则返回nil。

比较常规的用法

package main

import (
    "errors"
    "fmt"
    "time"

    "golang.org/x/sync/errgroup"
)

func main() {
    var g errgroup.Group

    // 启动第一个子任务,它执行成功
    g.Go(func() error {
        time.Sleep(5 * time.Second)
        fmt.Println("exec #1")
        // return errors.New("failed to exec #1")
        return nil
    })
    // 启动第二个子任务，它执行失败
    g.Go(func() error {
        time.Sleep(10 * time.Second)
        fmt.Println("exec #2")
        return errors.New("failed to exec #2")
    })

    // 启动第三个子任务，它执行成功
    g.Go(func() error {
        time.Sleep(15 * time.Second)
        fmt.Println("exec #3")
        return nil
    })
    // 等待三个任务都完成
    if err := g.Wait(); err == nil {
        fmt.Println("Successfully exec all")
    } else {
        fmt.Println("failed:", err)
    }
}

运行结果：

image-20201202163104917.png

在贴一个例子，是官方文档提供的一个pipeline的例子，原文地址：https://godoc.org/golang.org/x/sync/errgroup#example-Group--Pipeline。

package main

import (
    "context"
    "crypto/md5"
    "fmt"
    "io/ioutil"
    "log"
    "os"
    "path/filepath"

    "golang.org/x/sync/errgroup"
)

func main() {
    m, err := MD5All(context.Background(), ".")
    if err != nil {
        log.Fatal(err)
    }

    for k, sum := range m {
        fmt.Printf("%s:\t%x\n", k, sum)
    }
}

type result struct {
    path string
    sum  [md5.Size]byte
}

// 遍历根目录下的所有文件，计算md5值
func MD5All(ctx context.Context, root string) (map[string][md5.Size]byte, error) {
    g, ctx := errgroup.WithContext(ctx)
    //文件路径的channel
    paths := make(chan string)
    //遍历文件，将文件路径放到paths
    g.Go(func() error {
        defer close(paths)
        return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
            if err != nil {
                return err
            }
            if !info.Mode().IsRegular() {
                return nil
            }
            select {
            case paths <- path:
            case <-ctx.Done():
                return ctx.Err()
            }
            return nil
        })
    })

    // 20个goroutine计算md5，从paths获取文件路径
    c := make(chan result) //存储结果
    const numDigesters = 20
    for i := 0; i < numDigesters; i++ {
        g.Go(func() error {
            for path := range paths {
                data, err := ioutil.ReadFile(path)
                if err != nil {
                    return err
                }
                select {
                case c <- result{path, md5.Sum(data)}:
                case <-ctx.Done():
                    return ctx.Err()
                }
            }
            return nil
        })
    }
    go func() {
        g.Wait() //等待执行完
        close(c)
    }()
    //将结果输出到map
    m := make(map[string][md5.Size]byte)
    for r := range c {
        m[r.path] = r.sum
    }
    // 再次调用wait看有没有error
    if err := g.Wait(); err != nil {
        return nil, err
    }
    return m, nil
}