go context控制协程(goroutine)

疑问

在select多路复用中，如何优雅退出goroutine

方法

使用context

package main

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

var wg sync.WaitGroup

func main() {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	ch := make(chan int, 5)
	wg.Add(2)
	go send(ctx, ch)
	go receive(ctx, ch)
	time.Sleep(1 * time.Second)
	cancel()
	wg.Wait()
	fmt.Println("stop")
}

func send(ctx context.Context, ch chan int) {
	defer wg.Done()
	i := 0
	for {
		select {
		case v := <-ctx.Done():
			fmt.Println("send Done", v)
			return
		default:
			if i >= 10 {
				fmt.Println("send finished")
				return
			}
			ch<-i
			fmt.Println("send",i)
		}
		i++
		time.Sleep(200 * time.Millisecond)
	}
}

func receive(ctx context.Context, ch chan int) {
	defer wg.Done()
	for {
		select {
		case v := <-ctx.Done():
			fmt.Println("receive Done", v)
			return
		case v := <-ch:
			fmt.Println("receive", v)
		default:
			fmt.Println("receive wait")
			time.Sleep(1 * time.Second)
		}
	}
}

结果

可以看到当调用cancel时，会向Done发送一个信号，利用这个信号我们可以退出goroutine