golang worker channel 模式

• 大概流程就是job -> JobQueue
• 调度器循环获取JobQueue ，获取到的job ，再去异步获取等待可用的 worker，取出 chan Job,将job 写入改worker的 chan Job
• worker 处理任务，先处理 case job := <-w.JobChannel: 处理完成后再将 chan Job 写入到worker 里面，等待调度去取调用

package main

import (
	"log"
	"os"
	"strconv"
	"sync"
	"time"
)

var (
	MaxWorker int
	MaxQueue  int
	JobQueue  chan Job
)

func init() {
	var err error
	MaxWorker, err = strconv.Atoi(os.Getenv("MAX_WORKERS"))
	if err != nil {
		MaxWorker = 5 // 默认值
	}

	MaxQueue, err = strconv.Atoi(os.Getenv("MAX_QUEUE"))
	if err != nil {
		MaxQueue = 10 // 默认值
	}

	JobQueue = make(chan Job, MaxQueue)
}

type Payload struct {
	// Payload的属性
}

func (p *Payload) UploadToS3() error {
	// 模拟上传操作
	log.Println("Uploading to S3")
	return nil
}

type Job struct {
	Payload Payload
}

type Worker struct {
	WorkerPool chan chan Job
	JobChannel chan Job
	quit       chan bool
}

func NewWorker(workerPool chan chan Job) Worker {
	return Worker{
		WorkerPool: workerPool,
		JobChannel: make(chan Job),
		quit:       make(chan bool)}
}

func (w Worker) Start() {
	go func() {
		for {
			w.WorkerPool <- w.JobChannel
			select {
			case job := <-w.JobChannel:
				if err := job.Payload.UploadToS3(); err != nil {
					log.Printf("Error uploading to S3: %s", err)
				}
			case <-w.quit:
				return
			}
		}
	}()
}

func (w *Worker) Stop() {
	go func() {
		w.quit <- true // 通知工作线程停止
	}()
}

type Dispatcher struct {
	WorkerPool chan chan Job
	maxWorkers int
	workers    []Worker  // 新增：用于跟踪所有工作线程
	quit       chan bool // 用于停止dispatch循环
}

func NewDispatcher(maxWorkers int) *Dispatcher {
	return &Dispatcher{
		WorkerPool: make(chan chan Job, maxWorkers),
		maxWorkers: maxWorkers,
		workers:    make([]Worker, 0, maxWorkers),
	}
}

func (d *Dispatcher) Runs() {
	for i := 0; i < d.maxWorkers; i++ {
		worker := NewWorker(d.WorkerPool)
		d.workers = append(d.workers, worker) // 跟踪新创建的工作线程
		worker.Start()
	}

	go d.dispatch()
}

func (d *Dispatcher) dispatch() {
	for {
		select {
		// 从JobQueue中获取一个job
		case job := <-JobQueue:
			go func(job Job) {
				// 尝试获取一个可用的worker job channel，阻塞直到有可用的worker
				jobChannel := <-d.WorkerPool
				// 分发job到worker job channel中
				jobChannel <- job
			}(job)
		case <-d.quit:
			// 退出
			return
		}
	}
}

func (d *Dispatcher) StopAllWorkers() {
	var wg sync.WaitGroup
	for _, worker := range d.workers {
		wg.Add(1)
		go func(w Worker) {
			w.Stop() // 停止工作线程
			wg.Done()
		}(worker)
	}
	wg.Wait() // 等待所有工作线程安全退出
}

func (d *Dispatcher) Stop() {
	d.quit <- true
	d.StopAllWorkers()
}

func main() {
	dispatcher := NewDispatcher(MaxWorker)
	dispatcher.Runs()

	// 模拟作业提交
	for i := 0; i < 20; i++ {
		payload := Payload{ /* ... */ }
		job := Job{Payload: payload}
		JobQueue <- job
	}
	// 等待一段时间，以便可以看到工作的完成
	time.Sleep(10 * time.Second)
}