Golang开发--sync.WaitGroup

sync.WaitGroup 是 Go 语言标准库中的一个并发原语，用于等待一组并发操作的完成。它提供了一种简单的方式来跟踪一组 goroutine 的执行状态，并在所有 goroutine 完成后恢复执行。

下面是关于 sync.WaitGroup 的实现细节的详细解释：

• 创建 WaitGroup
  可以通过创建 sync.WaitGroup 类型的变量来创建 WaitGroup：

var wg sync.WaitGroup
1

• 添加任务
  使用 Add 方法将要等待的任务数量加一。每个任务都应该在启动之前调用 Add，以确保 WaitGroup 知道要等待的任务数量。

wg.Add(1) // 添加一个任务
1

• 完成任务
  在每个任务完成时，应调用 Done 方法来通知 WaitGroup 该任务已完成。

wg.Done() // 完成一个任务
1

等待任务完成：
使用 Wait 方法来阻塞当前 goroutine，直到所有的任务都完成。

wg.Wait() // 等待所有任务完成
1

如果在调用 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")
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

在上面的示例中，我们创建了一个 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
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

创建了两个无缓冲的通道 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
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

下面演示一下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")
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

在这个实现中，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。