sync.WaitGroup 是 Go 语言标准库中的一个并发原语,用于等待一组并发操作的完成。它提供了一种简单的方式来跟踪一组 goroutine 的执行状态,并在所有 goroutine 完成后恢复执行。
下面是关于 sync.WaitGroup 的实现细节的详细解释:
- 创建 WaitGroup
可以通过创建 sync.WaitGroup 类型的变量来创建 WaitGroup:
var wg sync.WaitGroup
- 添加任务
使用 Add 方法将要等待的任务数量加一。每个任务都应该在启动之前调用 Add,以确保 WaitGroup 知道要等待的任务数量。
wg.Add(1) // 添加一个任务
- 完成任务
在每个任务完成时,应调用 Done 方法来通知 WaitGroup 该任务已完成。
wg.Done() // 完成一个任务
等待任务完成:
使用 Wait 方法来阻塞当前 goroutine,直到所有的任务都完成。
wg.Wait() // 等待所有任务完成
如果在调用 Wait
之前已经调用了 Add
,那么 Wait
将会阻塞并等待所有任务完成。一旦所有任务完成,Wait
将返回,允许当前 goroutine 继续执行。
注意,Wait
方法可以在任何地方调用,但是需要确保在所有添加任务的地方都已经调用了 Add
方法,以避免出现死锁。
需要注意的是,WaitGroup 是通过内部计数器来实现的。每次调用 Add 方法增加计数器的值,每次调用 Done 方法减少计数器的值。当计数器的值为零时,等待的任务被认为已经完成。
下面是一个简单的示例,演示如何使用 WaitGroup:
package main
import (
"fmt"
"sync"
"time"
)
func main() {
var wg sync.WaitGroup
wg.Add(2) // 添加两个任务
go func() {
defer wg.Done() // 标记任务完成
time.Sleep(1 * time.Second)
fmt.Println("Task 1 completed")
}()
go func() {
defer wg.Done() // 标记任务完成
time.Sleep(2 * time.Second)
fmt.Println("Task 2 completed")
}()
wg.Wait() // 等待所有任务完成
fmt.Println("All tasks completed")
}
在上面的示例中,我们创建了一个 WaitGroup,并添加了两个任务。每个任务使用匿名函数表示,其中包含了任务的具体逻辑。在每个任务的最后,我们使用 defer wg.Done() 来标记任务的完成。最后,我们调用 wg.Wait() 来等待所有的任务完成,并在所有任务完成后打印 “All tasks completed”。
通过使用 WaitGroup,我们可以轻松地跟踪一组并发操作的完成状态,以便在需要时等待它们完成。这对于需要等待多个 goroutine 完成的并发任务非常有用,它包含一个计数器和两个方法:Add和Done。
Add方法用于增加计数器的值,表示有多少个goroutine需要等待。Done方法用于减少计数器的值,表示一个goroutine已经完成了它的工作。当计数器的值变为0时,Wait方法将返回,表示所有的goroutine都已经完成了它们的工作。
下面这个示例代码不会打印,思考一下为什么?
package main
import "fmt"
func main() {
var dog = make(chan string)
var cat = make(chan string)
go func() {
dog <- "dog"
fmt.Println("fog")
}()
go func() {
cat <- "cat"
fmt.Println("cat")
}()
<-dog
<-cat
}
创建了两个无缓冲的通道 dog 和 cat,并在两个匿名的 goroutine 中分别向这两个通道发送了字符串,然后使用 <-dog 和 <-cat 从通道中接收数据,但没有对接收到的数据进行任何处理。
问题出在这里:两个 goroutine 发送完数据之后,主 goroutine 就会继续执行 <-dog 和 <-cat 后面的代码,即关闭通道。然而,由于通道是无缓冲的,发送和接收操作是同步的,即发送操作会阻塞直到有对应的接收操作。因此,当主 goroutine 尝试接收数据时,由于没有 goroutine 在接收数据,发送操作也会被阻塞,导致主 goroutine 无法继续执行,从而没有打印任何内容。
为了解决这个问题,可以使用 sync.WaitGroup 来等待两个 goroutine 完成发送操作,然后再进行接收操作。
package main
import (
"fmt"
"sync"
)
func main() {
var wg sync.WaitGroup
wg.Add(2)
var dog = make(chan string)
var cat = make(chan string)
go func() {
defer wg.Done()
dog <- "dog"
fmt.Println("dog")
}()
go func() {
defer wg.Done()
cat <- "cat"
fmt.Println("cat")
}()
defer wg.Wait()
<-dog
<-cat
}
下面演示一下Add,Done,Wait的实现
package main
import (
"fmt"
"sync"
)
type WaitGroup struct {
counter int32
wait chan struct{}
lock sync.Mutex
}
func (wg *WaitGroup) Add(delta int) {
wg.lock.Lock()
defer wg.lock.Unlock()
wg.counter += int32(delta)
}
func (wg *WaitGroup) Done() {
wg.Add(-1)
}
func (wg *WaitGroup) Wait() {
wg.lock.Lock()
if wg.counter == 0 {
wg.lock.Unlock()
return
}
wg.wait = make(chan struct{})
defer wg.lock.Unlock()
<-wg.wait
}
func (wg *WaitGroup) DoneAndWait() {
wg.Done()
wg.Wait()
}
func main() {
var wg sync.WaitGroup
// 设置等待的 goroutine 数量
wg.Add(5)
// 模拟并发任务
for i := 0; i < 5; i++ {
go func(index int) {
defer wg.Done()
fmt.Printf("Goroutine %d started\n", index)
// 执行一些任务
// ...
fmt.Printf("Goroutine %d finished\n", index)
}(i)
}
wg.Wait()
fmt.Println("All goroutines finished")
}
在这个实现中,WaitGroup包含一个计数器和一个等待通道。Add方法使用互斥锁来保护计数器的并发访问。Done方法简单地调用Add方法并将delta设置为-1。Wait方法首先使用互斥锁来检查计数器的值是否为0。如果计数器的值为0,则立即返回。否则,它创建一个新的等待通道,并将其存储在WaitGroup中。最后,它释放互斥锁并等待等待通道上的信号。
wait chan struct{}
wait 参数是一个无缓冲通道,用于在 Wait 方法中阻塞等待信号的接收。<-wg.wait 表示从通道接收信号,阻塞当前 goroutine 直到接收到信号。这种机制允许等待的 goroutine 在条件满足时被唤醒,从而实现等待并发任务完成的效果。
在 Wait 方法中,当计数器 counter 不为零时,会创建一个新的无缓冲通道并将其赋值给 wait 字段。然后,在 <-wg.wait 这一行代码中,当前的 goroutine 会阻塞,直到从 wg.wait 通道接收到一个值。
<-wg.wait
表示从 wg.wait 通道接收值。在这里,我们不关心接收到的具体值是什么,因此我们使用了空的 struct{} 类型,这个类型不占用任何内存空间,只是为了作为一个信号使用。
使用 <-wg.wait
的目的是让当前的 goroutine 阻塞,直到最后一个 goroutine 执行 Done 方法并将计数器 counter 减少到零。当计数器为零时,最后一个 goroutine 会通过向 wg.wait 通道发送一个值,这个值会被当前的 goroutine 接收到,从而解除阻塞。
通过这种方式,我们可以实现等待所有 goroutine 完成的效果:每个 goroutine 在完成任务后调用 Done 方法,计数器 counter 减少,直到最后一个 goroutine 完成任务并将计数器减少到零,从而释放阻塞在 <-wg.wait 处的等待 goroutine。