Golang并发控制之errgroup使用详解

package main

import (
    "fmt"
    "net/http"
    "sync"
)

func main() {
    var urls = []string{
        "http://www.golang.org/",
        "http://www.google.com/",
        "http://www.somestupidname.com/", // 这是一个错误的 URL，会导致任务失败
    }
    var err error

    var wg sync.WaitGroup // 零值可用，不必显式初始化

    for _, url := range urls {
        wg.Add(1) // 增加 WaitGroup 计数器

        // 启动一个 goroutine 来获取 URL
        go func() {
            defer wg.Done() // 当 goroutine 完成时递减 WaitGroup 计数器

            resp, e := http.Get(url)
            if e != nil { // 发生错误返回，并记录该错误
                err = e
                return
            }
            defer resp.Body.Close()
            fmt.Printf("fetch url %s status %s\n", url, resp.Status)
        }()
    }

    // 等待所有 goroutine 执行完成
    wg.Wait()
    if err != nil { // err 会记录最后一个错误
        fmt.Printf("Error: %s\n", err)
    }
}

var wg sync.WaitGroup

for ... {
    wg.Add(1)

    go func() {
        defer wg.Done()
        // do something
    }()
}

wg.Wait()

package main

import (
    "fmt"
    "net/http"

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

func main() {
    var urls = []string{
        "http://www.golang.org/",
        "http://www.google.com/",
        "http://www.somestupidname.com/", // 这是一个错误的 URL，会导致任务失败
    }

    // 使用 errgroup 创建一个新的 goroutine 组
    var g errgroup.Group // 零值可用，不必显式初始化

    for _, url := range urls {
        // 使用 errgroup 启动一个 goroutine 来获取 URL
        g.Go(func() error {
            resp, err := http.Get(url)
            if err != nil {
                return err // 发生错误，返回该错误
            }
            defer resp.Body.Close()
            fmt.Printf("fetch url %s status %s\n", url, resp.Status)
            return nil // 返回 nil 表示成功
        })
    }

    // 等待所有 goroutine 完成并返回第一个错误（如果有）
    if err := g.Wait(); err != nil {
        fmt.Printf("Error: %s\n", err)
    }
}

package main

import (
    "context"
    "fmt"
    "net/http"
    "sync"

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

func main() {
    var urls = []string{
        "http://www.golang.org/",
        "http://www.google.com/",
        "http://www.somestupidname.com/", // 这是一个错误的 URL，会导致任务失败
    }

    // 创建一个带有 context 的 errgroup
    // 任何一个 goroutine 返回非 nil 的错误，或 Wait() 等待所有 goroutine 完成后，context 都会被取消
    g, ctx := errgroup.WithContext(context.Background())

    // 创建一个 map 来保存结果
    var result sync.Map

    for _, url := range urls {
        // 使用 errgroup 启动一个 goroutine 来获取 URL
        g.Go(func() error {
            req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
            if err != nil {
                return err // 发生错误，返回该错误
            }

            // 发起请求
            resp, err := http.DefaultClient.Do(req)
            if err != nil {
                return err // 发生错误，返回该错误
            }
            defer resp.Body.Close()

            // 保存每个 URL 的响应状态码
            result.Store(url, resp.Status)
            return nil // 返回 nil 表示成功
        })
    }

    // 等待所有 goroutine 完成并返回第一个错误（如果有）
    if err := g.Wait(); err != nil {
        fmt.Println("Error: ", err)
    }

    // 所有 goroutine 都执行完成，遍历并打印成功的结果
    result.Range(func(key, value any) bool {
        fmt.Printf("fetch url %s status %s\n", key, value)
        return true
    })
}

package main

import (
    "fmt"
    "time"

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

func main() {
    // 创建一个 errgroup.Group
    var g errgroup.Group
    // 设置最大并发限制为 3
    g.SetLimit(3)

    // 启动 10 个 goroutine
    for i := 1; i <= 10; i++ {
        g.Go(func() error {
            // 打印正在运行的 goroutine
            fmt.Printf("Goroutine %d is starting\n", i)
            time.Sleep(2 * time.Second) // 模拟任务耗时
            fmt.Printf("Goroutine %d is done\n", i)
            return nil
        })
    }

    // 等待所有 goroutine 完成
    if err := g.Wait(); err != nil {
        fmt.Printf("Encountered an error: %v\n", err)
    }

    fmt.Println("All goroutines complete.")
}

package main

import (
    "fmt"
    "time"

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

func main() {
    // 创建一个 errgroup.Group
    var g errgroup.Group
    // 设置最大并发限制为 3
    g.SetLimit(3)

    // 启动 10 个 goroutine
    for i := 1; i <= 10; i++ {
        if g.TryGo(func() error {
            // 打印正在运行的 goroutine
            fmt.Printf("Goroutine %d is starting\n", i)
            time.Sleep(2 * time.Second) // 模拟工作
            fmt.Printf("Goroutine %d is done\n", i)
            return nil
        }) {
            // 如果成功启动，打印提示
            fmt.Printf("Goroutine %d started successfully\n", i)
        } else {
            // 如果达到并发限制，打印提示
            fmt.Printf("Goroutine %d could not start (limit reached)\n", i)
        }
    }

    // 等待所有 goroutine 完成
    if err := g.Wait(); err != nil {
        fmt.Printf("Encountered an error: %v\n", err)
    }

    fmt.Println("All goroutines complete.")
}

// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package errgroup provides synchronization, error propagation, and Context
// cancelation for groups of goroutines working on subtasks of a common task.
//
// [errgroup.Group] is related to [sync.WaitGroup] but adds handling of tasks
// returning errors.
package errgroup

import (
    "context"
    "fmt"
    "sync"
)

type token struct{}

// A Group is a collection of goroutines working on subtasks that are part of
// the same overall task.
//
// A zero Group is valid, has no limit on the number of active goroutines,
// and does not cancel on error.
type Group struct {
    cancel func(error)

    wg sync.WaitGroup

    sem chan token

    errOnce sync.Once
    err     error
}

func (g *Group) done() {
    if g.sem != nil {
        <-g.sem
    }
    g.wg.Done()
}

// WithContext returns a new Group and an associated Context derived from ctx.
//
// The derived Context is canceled the first time a function passed to Go
// returns a non-nil error or the first time Wait returns, whichever occurs
// first.
func WithContext(ctx context.Context) (*Group, context.Context) {
    ctx, cancel := withCancelCause(ctx)
    return &Group{cancel: cancel}, ctx
}

// Wait blocks until all function calls from the Go method have returned, then
// returns the first non-nil error (if any) from them.
func (g *Group) Wait() error {
    g.wg.Wait()
    if g.cancel != nil {
        g.cancel(g.err)
    }
    return g.err
}

// Go calls the given function in a new goroutine.
// It blocks until the new goroutine can be added without the number of
// active goroutines in the group exceeding the configured limit.
//
// The first call to return a non-nil error cancels the group's context, if the
// group was created by calling WithContext. The error will be returned by Wait.
func (g *Group) Go(f func() error) {
    if g.sem != nil {
        g.sem <- token{}
    }

    g.wg.Add(1)
    go func() {
        defer g.done()

        if err := f(); err != nil {
            g.errOnce.Do(func() {
                g.err = err
                if g.cancel != nil {
                    g.cancel(g.err)
                }
            })
        }
    }()
}

// TryGo calls the given function in a new goroutine only if the number of
// active goroutines in the group is currently below the configured limit.
//
// The return value reports whether the goroutine was started.
func (g *Group) TryGo(f func() error) bool {
    if g.sem != nil {
        select {
        case g.sem <- token{}:
            // Note: this allows barging iff channels in general allow barging.
        default:
            return false
        }
    }

    g.wg.Add(1)
    go func() {
        defer g.done()

        if err := f(); err != nil {
            g.errOnce.Do(func() {
                g.err = err
                if g.cancel != nil {
                    g.cancel(g.err)
                }
            })
        }
    }()
    return true
}

// SetLimit limits the number of active goroutines in this group to at most n.
// A negative value indicates no limit.
//
// Any subsequent call to the Go method will block until it can add an active
// goroutine without exceeding the configured limit.
//
// The limit must not be modified while any goroutines in the group are active.
func (g *Group) SetLimit(n int) {
    if n < 0 {
        g.sem = nil
        return
    }
    if len(g.sem) != 0 {
        panic(fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", len(g.sem)))
    }
    g.sem = make(chan token, n)
}

// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build go1.20

package errgroup

import "context"

func withCancelCause(parent context.Context) (context.Context, func(error)) {
    return context.WithCancelCause(parent)
}

// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build !go1.20

package errgroup

import "context"

func withCancelCause(parent context.Context) (context.Context, func(error)) {
    ctx, cancel := context.WithCancel(parent)
    return ctx, func(error) { cancel() }
}

// 定义一个空结构体类型 token，会作为信号进行传递，用于控制并发数
type token struct{}

// Group 是一组协程的集合，这些协程处理同一整体任务的子任务
//
// 零值 Group 是有效的，对活动协程的数量没有限制，并且不会在出错时取消
type Group struct {
    cancel func(error) // 取消函数，就是 context.CancelCauseFunc 类型

    wg sync.WaitGroup // 内部使用了 sync.WaitGroup

    sem chan token // 信号 channel，可以控制协程并发数量

    errOnce sync.Once // 确保错误仅处理一次
    err     error     // 记录子协程集中返回的第一个错误
}

// SetLimit 限制该 Group 中活动的协程数量最多为 n，负值表示没有限制
//
// 任何后续对 Go 方法的调用都将阻塞，直到可以在不超过限额的情况下添加活动协程
//
// 在 Group 中存在任何活动的协程时，限制不得修改
func (g *Group) SetLimit(n int) { // 传进来的 n 就是 channel 长度，以此来限制协程的并发数量
    if n < 0 { // 这里检查如果小于 0 则不限制协程并发数量。此外，也不要将其设置为 0，会产生死锁
        g.sem = nil
        return
    }
    if len(g.sem) != 0 { // 如果存在活动的协程，调用此方法将产生 panic
        panic(fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", len(g.sem)))
    }
    g.sem = make(chan token, n)
}

// Go 会在新的协程中调用给定的函数
// 它会阻塞，直到可以在不超过配置的活跃协程数量限制的情况下添加新的协程
//
// 首次返回非 nil 错误的调用会取消该 Group 的上下文（context），如果该 context 是通过调用 WithContext 创建的，该错误将由 Wait 返回
func (g *Group) Go(f func() error) {
    if g.sem != nil { // 这个是限制并发数的信号通道
        g.sem <- token{} // 如果超过了配置的活跃协程数量限制，向 channel 发送 token 会阻塞
    }

    g.wg.Add(1) // 转发给 sync.WaitGroup.Add(1)，将活动协程数加一
    go func() {
        defer g.done() // 当一个协程完成时，调用此方法，内部会将调用转发给 sync.WaitGroup.Done()

        if err := f(); err != nil { // f() 就是我们要执行的任务
            g.errOnce.Do(func() { // 仅执行一次，即只处理一次错误，所以会记录第一个非 nil 的错误，与协程启动顺序无关
                g.err = err          // 记录错误
                if g.cancel != nil { // 如果 cancel 不为 nil，则调用取消函数，并设置 cause
                    g.cancel(g.err)
                }
            })
        }
    }()
}

// 当一个协程完成时，调用此方法
func (g *Group) done() {
    // 如果设置了最大并发数，则 sem 不为 nil，从 channel 中消费一个 token，表示一个协程已完成
    if g.sem != nil {
        <-g.sem
    }
    g.wg.Done() // 转发给 sync.WaitGroup.Done()，将活动协程数减一
}

// WithContext 返回一个新的 Group 和一个从 ctx 派生的关联 Context
//
// 派生的 Context 会在传递给 Go 的函数首次返回非 nil 错误或 Wait 首次返回时被取消，以先发生者为准。
func WithContext(ctx context.Context) (*Group, context.Context) {
    ctx, cancel := withCancelCause(ctx)
    return &Group{cancel: cancel}, ctx
}

// 构建约束标识了这个文件是 Go 1.20 版本被加入的
//go:build go1.20

package errgroup

import "context"

// 代理到 context.WithCancelCause
func withCancelCause(parent context.Context) (context.Context, func(error)) {
    return context.WithCancelCause(parent)
}

//go:build !go1.20

package errgroup

import "context"

func withCancelCause(parent context.Context) (context.Context, func(error)) {
    ctx, cancel := context.WithCancel(parent)
    return ctx, func(error) { cancel() }
}

// Wait 会阻塞，直到来自 Go 方法的所有函数调用返回，然后返回它们中的第一个非 nil 错误（如果有的话）
func (g *Group) Wait() error {
    g.wg.Wait()          // 转发给 sync.WaitGroup.Wait()，等待所有协程执行完成
    if g.cancel != nil { // 如果 cancel 不为 nil，则调用取消函数，并设置 cause
        g.cancel(g.err)
    }
    return g.err // 返回错误
}

// TryGo 仅在 Group 中活动的协程数量低于限额时，才在新的协程中调用给定的函数
//
// 返回值标识协程是否启动
func (g *Group) TryGo(f func() error) bool {
    if g.sem != nil { // 如果设置了最大并发数
        select {
        case g.sem <- token{}: // 可以向 channel 写入 token，说明没有达到限额，可以启动协程
            // Note: this allows barging iff channels in general allow barging.
        default: // 如果超过了配置的活跃协程数量限制，会走到这个 case
            return false
        }
    }

    // 接下来的代码与 Go 中的逻辑相同
    g.wg.Add(1)
    go func() {
        defer g.done()

        if err := f(); err != nil {
            g.errOnce.Do(func() {
                g.err = err
                if g.cancel != nil {
                    g.cancel(g.err)
                }
            })
        }
    }()
    return true
}