并发编程技巧：Go WaitGroup的高级用法

并发编程技巧：Go WaitGroup的高级用法

在并发编程中，协调和管理多个并发任务的执行是一项重要的任务。Go语言提供了一个非常实用的并发原语——WaitGroup，它可以帮助我们优雅地实现并发控制。本文将介绍WaitGroup的基本用法，并重点讨论其高级用法，通过具体的代码示例来帮助读者更好地理解和应用。

WaitGroup是Go语言内置的一个并发原语，它能够帮助我们等待并发任务的完成。它提供了三个方法：Add、Done和Wait。Add方法用于设置等待任务的数量，Done方法用于减少等待任务的数量，Wait方法用于阻塞当前协程，直到所有等待任务完成。

下面是一个简单的示例，展示了WaitGroup的基本用法：

package main

import (
    "fmt"
    "sync"
    "time"
)

func main() {
    var wg sync.WaitGroup

    for i := 0; i < 5; i++ {
        wg.Add(1)
        go func(num int) {
            defer wg.Done()

            time.Sleep(time.Second)
            fmt.Println("Task", num, "done")
        }(i)
    }

    wg.Wait()
    fmt.Println("All tasks done")
}

在上面的代码中，我们创建了一个WaitGroup对象wg，并通过循环创建了5个并发任务。在每个任务的执行过程中，我们使用Add方法增加等待任务的数量，并在任务结束时通过Done方法减少等待任务的数量。最后，我们调用Wait方法阻塞了主协程，直到所有等待任务完成。

除了基本的用法外，WaitGroup还提供了一些高级的用法，可以更加灵活地控制并发任务的执行。下面我们将详细介绍几个常用的高级用法。

1. 执行一组任务并设置最大并发数

如果我们需要同时执行一组任务，但又希望限制最大并发数，可以使用带缓冲通道结合WaitGroup来实现。下面的代码展示了如何同时执行一组任务，但最多只允许3个任务并发执行：

package main

import (
    "fmt"
    "sync"
    "time"
)

func main() {
    var wg sync.WaitGroup
    maxConcurrency := 3
    tasks := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}

    sem := make(chan struct{}, maxConcurrency)

    for _, task := range tasks {
        wg.Add(1)
        sem <- struct{}{} // 获取令牌，控制最大并发数

        go func(num int) {
            defer wg.Done()

            time.Sleep(time.Second)
            fmt.Println("Task", num, "done")

            <-sem // 释放令牌，允许新的任务执行
        }(task)
    }

    wg.Wait()
    fmt.Println("All tasks done")
}

在上面的代码中，我们创建了一个带缓冲的通道sem，并将其大小设置为最大并发数。在每个任务开始前，我们通过sem <- struct{}{}语句获取一个令牌，当任务完成后，使用<-sem语句释放令牌。通过控制令牌的获取和释放，我们就能够限制最大并发数。

2. 超时控制并发任务的执行

有时候我们希望对并发任务的执行时间进行控制，并在超时时终止任务的执行。通过使用带缓冲通道和定时器，我们可以轻松实现这个功能。下面的代码展示了如何设置并发任务的超时时间为3秒：

package main

import (
    "fmt"
    "sync"
    "time"
)

func main() {
    var wg sync.WaitGroup
    tasks := []int{1, 2, 3, 4, 5, 6, 7}

    timeout := 3 * time.Second
    done := make(chan struct{})

    for _, task := range tasks {
        wg.Add(1)

        go func(num int) {
            defer wg.Done()

            // 模拟任务执行时间不定
            time.Sleep(time.Duration(num) * time.Second)
            fmt.Println("Task", num, "done")

            // 判断任务是否超时
            select {
            case <-done:
                // 任务在超时前完成，正常退出
                return
            default:
                // 任务超时，向通道发送信号
                close(done)
            }
        }(task)
    }



    wg.Wait()
    fmt.Println("All tasks done")
}

在上面的代码中，我们创建了一个通道done，并在任务执行过程中判断通道是否关闭来判断任务是否超时。当一个任务完成时，我们使用close(done)语句向done通道发送信号，表示任务已经超时。通过select语句选择不同的分支来处理不同的情况。

通过上面的示例代码，我们可以看到WaitGroup的高级用法在实际的并发编程中非常实用。掌握了这些技巧，我们能够更好地控制并发任务的执行，提高代码的性能和可维护性。希望读者能通过本文的介绍和示例代码，深入理解WaitGroup的用法，从而应用到实际的项目中。