package schedule
import (
"container/list"
"context"
"fmt"
"sync"
)
/**********************************\
Task任务接口
\**********************************/
type Task interface {
// Process 任务处理逻辑. 通过ctx查知worker是否关闭.
// 若有业务侧的"cancel/ack"关闭,请在process封装!
Process(ctx context.Context) error
}
/**********************************\
Worker结构
\**********************************/
type Worker struct {
wctx context.Context
wcnf context.CancelFunc
}
func NewWorker(parent context.Context, tasks chan Task) *Worker {
worker := new(Worker)
worker.wctx, worker.wcnf = context.WithCancel(parent)
go worker.consume(tasks)
return worker
}
func (w *Worker) Cancel() {
w.wcnf()
}
func (w *Worker) consume(tasks chan Task) {
for {
select {
case <-w.wctx.Done():
return
default:
task, ok := <-tasks
if !ok {
return
}
if !w.process(task) {
// 未处理成功需返还tasks
tasks <- task
}
}
}
}
func (w *Worker) process(task Task) bool {
defer func() {
if prr := recover(); prr != nil {
// 集中处理panic
fmt.Printf("worker process panic: %v", prr)
// 最好堆栈一块打印方便排查
}
}()
// 二次校验worker是否结束
select {
case <-w.wctx.Done():
return false
default:
// process task
err := task.Process(w.wctx)
if err != nil {
// 集中处理error
fmt.Printf("worker process error: %v", err)
}
return true
}
}
/**********************************\
Master结构
\**********************************/
// Master 管理者数据结构
type Master struct {
mutex sync.Mutex // 同步控制
works *list.List // 工作者列表
tasks chan Task // 任务管道
mctx context.Context
mcnf context.CancelFunc
}
func NewMaster(parent context.Context, tasks chan Task, works int) *Master {
master := new(Master)
master.works = list.New()
master.tasks = tasks
master.mctx, master.mcnf = context.WithCancel(parent)
master.Stretch(works)
return master
}
func (m *Master) Stretch(tasks int) {
m.mutex.Lock()
defer m.mutex.Unlock()
for m.works.Len() < tasks {
// 新增
work := NewWorker(m.mctx, m.tasks)
m.works.PushBack(work)
}
for m.works.Len() > tasks {
// 缩减
elem := m.works.Back()
if work, ok := elem.Value.(*Worker); ok {
work.Cancel()
}
m.works.Remove(elem)
}
}
func (m *Master) Cancel() {
m.mcnf()
}
