Golang等待Goroutine完成

当我们在Golang中启动了Goroutine，那么我们如何才能够确保这些Goroutine在执行结束后再退出程序呢？本文将会分享一些方法来等待Goroutine完成。

使用WaitGroup

Go语言内置了sync包，其中的WaitGroup就是为了处理一组Goroutine的同步问题而生的。

package main

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

func main() {
	var wg sync.WaitGroup

	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done() // 通知WaitGroup任务已完成
			time.Sleep(1 * time.Second)
			fmt.Printf("Goroutine %d 完成\n", i)
		}(i)
	}

	wg.Wait() // 等待所有Goroutine完成

	fmt.Println("All Goroutine Done")
}

由于WaitGroup的计数器是非负整数，因此我们通过Add方法增加协程数量，然后执行具体任务时将计数器减1。最后，我们使用Wait方法，阻塞直到所有Goroutine执行完毕。

使用Channel

另一种等待Goroutine完成的方法是通过Channel进行通信。在Goroutine完成任务后，我们向Channel发送一个消息，主线程等待所有Goroutine发送完消息后再退出。

package main

import (
	"fmt"
	"time"
)

func main() {
	done := make(chan struct{}) // 创建一个无缓冲的Channel

	for i := 0; i < 5; i++ {
		go func(i int) {
			time.Sleep(1 * time.Second)
			fmt.Printf("Goroutine %d 完成\n", i)
			done <- struct{}{} // 发送一个空结构体通知主线程
		}(i)
	}

	for i := 0; i < 5; i++ {
		<-done // 阻塞，等待所有Goroutine完成
	}

	fmt.Println("All Goroutine Done")
}

在上面的例子中，我们创建了一个无缓冲的Channel，在每个Goroutine结束后发送一个空结构体。最后，我们在主线程阻塞等待所有结构体都被接收到。

使用context

在Golang中，我们可以使用context包来取消或超时一个Goroutine或函数调用。通过WithCancel方法，我们可以创建一个父Context和用于取消Context的函数。在每个Goroutine中，我们监听这个Context的Done信号。如果该信号被触发，我们将退出Goroutine。

package main

import (
	"context"
	"fmt"
	"time"
)

func worker(id int, ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			fmt.Printf("Worker %d 已退出\n", id)
			return
		default:
			time.Sleep(200 * time.Millisecond)
			fmt.Printf("Worker %d is working\n", id)
		}
	}
}

func main() {
	ctx, cancel := context.WithCancel(context.Background())

	for i := 0; i < 5; i++ {
		go worker(i, ctx)
	}

	time.Sleep(2 * time.Second)

	cancel() // 通知所有Goroutine退出

	time.Sleep(2 * time.Second)
	fmt.Println("All workers Done")
}

在worker函数中，我们使用select语句监听Context的Done信号。如果该信号被触发，我们将退出协程。由于我们在主线程中延时了2秒后cancel Context，因此所有Goroutine将退出。

总结

等待Goroutine完成是Go语言中常见的问题，我们有三种方法可以解决：使用WaitGroup、使用Channel、使用Context。每个方法都有自己的优点和适用场景，程序员应该根据需要选择合适的方法来解决这个问题。