Go语言如何保证context监听协程在主程序退出前执行完成

问题原因 #

你当前代码存在两个核心问题导致解锁逻辑无法保证执行完成：

1. 周期调度协程死循环：收到context取消信号后仅停止了定时器，没有退出for循环，wg.Done()永远不会被触发，且程序退出时会直接强制终止所有协程
2. 解锁协程无同步机制：process函数内启动的监听context取消的协程是游离状态，没有任何等待机制，process函数本身执行完就会直接返回，不会等待清理协程执行完毕

修复方案 #

核心修改点 #

• 给周期调度协程增加退出逻辑，收到取消信号后直接退出循环触发wg.Done()
• 在process函数内增加sync.WaitGroup，强制等待内部解锁协程执行完成后再返回
• 不管是process正常执行结束触发的context取消，还是外部传入的context取消，解锁逻辑都会被完整执行后才会向上返回

修复后完整代码 #

package main

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

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

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		fmt.Println("Started periodic process")
		defer wg.Done()
		ticker := time.NewTicker(20 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				process(ctx)
			case <-ctx.Done():
				// 收到取消信号直接退出循环
				return
			}
		}
		fmt.Println("Finished periodic process")
	}()

	time.Sleep(100 * time.Millisecond)

	// cancel context
	cancel()

	wg.Wait()
}

func process(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	var cleanWg sync.WaitGroup
	cleanWg.Add(1)
	go func() {
		defer cleanWg.Done()
		<-ctx.Done()

		// unlock resources
		fmt.Println("Unlocking resources")

		time.Sleep(20 * time.Millisecond)

		fmt.Println("Unlocked resources")
	}()
	// 函数返回前必须等清理协程执行完成
	defer cleanWg.Wait()

	// do some work
	fmt.Println("Started process")

	// acquire lock, in actual process
	// acquiring lock will fail unless
	// it was released be previous process
	// run
	fmt.Println("Acquired lock")

	time.Sleep(10 * time.Millisecond)

	fmt.Println("Finished process")
}

验证说明 #

运行修改后的代码，你会发现所有Unlocked resources日志都会被正常打印，不存在程序退出时强制杀掉清理协程的问题。主协程的sync.WaitGroup会等待周期调度协程完全退出，而每次调用process都会等待内部解锁逻辑全部执行完成后才返回，整个链路的退出逻辑是完全同步可控的。

内容的提问来源于stack exchange，提问作者septerr