package main

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

// 带 WaitGroup 的工作函数
func workerWaitGroup(ctx context.Context, wg *sync.WaitGroup, id int) {
	defer wg.Done() // 通知 WaitGroup 该 Goroutine 已完成
	defer fmt.Printf("worker %d：清理完成\n", id)

	for {
		select {
		case <-ctx.Done():
			fmt.Printf("worker %d：收到取消信号，退出\n", id)
			return
		default:
			fmt.Printf("worker %d：处理任务...\n", id)
			time.Sleep(400 * time.Millisecond)
		}
	}
}

func main() {
	var wg sync.WaitGroup
	ctx, cancel := context.WithCancel(context.Background())

	// 启动 3 个工作 Goroutine
	for i := 0; i < 3; i++ {
		wg.Add(1) // 注册一个 Goroutine
		go workerWaitGroup(ctx, &wg, i)
	}

	// 2 秒后触发取消
	time.Sleep(2 * time.Second)
	fmt.Println("主程序：发送取消信号")
	cancel()

	// 等待所有 Goroutine 退出（阻塞直到 wg 计数为 0）
	wg.Wait()
	fmt.Println("所有工作已完成，主程序退出")
}