Go面试题：worker pool

题面

1. master负责分发计算任务，任务总数1000
2. master启动worker协程，master退出前，必须通知所有worker先退出，自己再退出（假定master在所有计算任务完成前不会退出）
3. worker协程不超过3个；worker超时5s自动退出（假定要归还任务）
4. worker计算结果返回给master

思路

master：

1. 初始化：缓冲大小==worker数量的dieCh；无缓冲任务通道jobCh，无缓冲结束通告doneCh（只负责通告monitor）；缓冲大小==总任务数量的结果通道resultCh；worker、monitor的2个waitGroup（全局变量，含义见下文）
2. for启动所有worker，Add waitGroup
3. for启动monitor，Add waitGroup（与前一步的顺序可颠倒）
4. for分派任务——写入jobCh
5. for读取结果——从resultCh读
6. 关闭jobCh：通告worker们退出（由于上一步读完了结果，此刻必然——所有worker阻塞，因为没任务可做；monitor阻塞，因为没worker会超时）
7. 关闭（或写入）doneCh，通告monitor退出。（与前一步顺序可颠倒）
8. 等待waitGroup（worker的、monitor的，共2个）
9. return

monitor：

1. 无限for包裹
2. select检查只读doneCh、只读dieCh：前者则return；后者则对worker的waitGroup Add，并重启一个worker

worker：

1. defer：worker的waitGroup要Done
2. for range通道jobCh：启动定时器，select定时器、计算结果的返回——超时则只写dieCh并向jobCh归还job并return；否则只写resultCh

时序图

对一个共享变量，最后一个写操作happened-before读操作，就不容易出问题

master与worker同步

monitor与worker同步，monitor与master同步

代码

环境：go1.22.4，windows，amd64

package main

import (
	"log"
	"math/rand"
	"sync"
	"time"
)

type Job struct{ id int } // 任务
type Result struct{} // 计算结果

const jobNum = 1000
const workerNum = 3
const expire = 5 // 超时时间,单位秒
// 返回一个计算结果
func GetResult() *Result {
	// 测试worker超时
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	time.Sleep(time.Duration(r.Intn(1+expire) * 1000000000) * time.Second)
	return &Result{}
}

func worker(id int, wg *sync.WaitGroup, jobCh chan *Job, resultCh chan<- *Result, dieCh chan<- int) {
	defer wg.Done()
	for job := range jobCh {
		timer := time.NewTimer(expire * time.Second)
		select {
		case <-timer.C: //处理超时
			dieCh <- id  //通告monitor重启一个worker
			jobCh <- job // 归还失败任务
			return
		case resultCh <- GetResult(): //返回计算结果
			log.Printf("worker#%v success job#%v\n", id, job.id)
		}
	}
}

func monitor(doneCh <-chan struct{}, dieCh chan int, jobCh chan *Job, resultCh chan<- *Result, waitMonitor, waitWorkers *sync.WaitGroup) {
	defer func() {
		waitMonitor.Done()
		log.Println(">>monitor exit")
	}()
	for {
		select {
		case <-doneCh: // master通告结束
			return
		case id := <-dieCh: // 重启一个worker
			waitWorkers.Add(1)
			log.Printf("recover worker#%v\n", id)
			go worker(id, waitWorkers, jobCh, resultCh, dieCh)
		}
	}
}

func main() {
	// 1.初始化
	dieCh := make(chan int, workerNum) // worker报告超时退出
	jobCh := make(chan *Job, jobNum)
	doneCh := make(chan struct{}) // 通告monitor
	resultCh := make(chan *Result, jobNum)

	waitWorkers := &sync.WaitGroup{}
	waitMonitor := &sync.WaitGroup{}
	// 2.启动monitor
	waitMonitor.Add(1)
	go monitor(doneCh, dieCh, jobCh, resultCh, waitMonitor, waitWorkers)
	// 3.启动worker
	for i := 0; i < workerNum; i++ {
		waitWorkers.Add(1)
		go worker(i, waitWorkers, jobCh, resultCh, dieCh)
	}
	// 4.分派任务
	for i := 0; i < jobNum; i++ {
		jobCh <- &Job{i}
	}
	// 5.读取结果
	for i := 0; i < jobNum; i++ {
		<-resultCh
	}
	// 此刻worker全阻塞在jobCh，monitor阻塞在dieCh
	// 6.通告、等待worker、monitor退出
	close(jobCh)
	close(doneCh)
	waitWorkers.Wait()
	waitMonitor.Wait()

	log.Println(">>master exit")
}