golang常驻后台类worker模板

╰+哭是因爲堅強的太久メ 2022-12-01 15:58 124阅读 0赞

golang中经常会用到常驻后台类的worker，实现例如消费队列、定期执行任务、定期统计数据等功能。

这里自己实现了通用的worker模板，主要有以下功能：

1.  panic自动重启，最大重启次数可自定义。
2.  optional参数，有默认参数和支持自定义参数。
3.  busy模式和idle模式，执行完任务后睡眠不同的时间。
4.  支持以一定的频率执行，例如每5分钟执行一次任务。
5.  可动态的通过channel控制任务的启动和暂停。ps：此处可结合etcd实现分布式选主机制。
6.  支持用context优雅的停止，worker停止后waitGroup通知，超时强停。

package worker
    
    import (
    	"fmt"
    	"strings"
    	"sync"
    	"time"
    
    	"context"
    
    )
    
    const (
    	defaultWorkerName        = "worker"
    	defaultBusySleepTime     = time.Millisecond * 100
    	defaultIdleSleepTime     = time.Second * 1
    	defaultPanicRecoverCount = 5
    	defaultWaitCloseTime     = time.Second * 5
    )
    
    type TaskFunc func() (idle bool, err error)
    type WorkerOption func(*Worker)
    
    type Worker struct {
    	name              string
    	busySleepTime     time.Duration
    	idleSleepTime     time.Duration
    	task              TaskFunc
    	panicRecoverCount int
    	ctx               context.Context
    	wg                *sync.WaitGroup
    	taskFrequency     time.Duration
    	checkMaster       <-chan struct{}
    	waitCloseTime     time.Duration
    }
    
    func (w *Worker) Work() {
    	defer func() {
    		if r := recover(); r != nil {
    			DealRecover(r)
    			if w.panicRecoverCount == 0 {
    				fmt.Println(w.name+" panic too many", r)
    			} else {
    				w.panicRecoverCount--
    				go w.Work() //restart goroutine
    				return
    			}
    		}
    		if w.wg != nil {
    			w.wg.Done()
    		}
    	}()
    	var (
    		timer     = time.NewTimer(time.Millisecond)
    		sleepTime time.Duration
    	)
    	defer timer.Stop()
    	for {
    		select {
    		case <-w.ctx.Done():
    			return
    		case <-w.checkMaster:
    		}
    		var (
    			done = make(chan struct{})
    			idle bool
    			err  error
    		)
    		go func() {
    			idle, err = w.task()
    			if err != nil {
    				dlog.Error(w.name+" err", err)
    			}
    			close(done)
    		}()
    		// 优雅的停止，但同时不能无限期的等。（有些任务可能执行很久）
    		select {
    		case <-w.ctx.Done():
    			select {
    			case <-done:
    			case <-time.After(w.waitCloseTime):
    			}
    			return
    		case <-done:
    		}
    		if w.taskFrequency > 0 {
    			sleepTime = w.taskFrequency
    		} else {
    			sleepTime = w.idleSleepTime
    			if !idle {
    				sleepTime = w.busySleepTime
    			}
    		}
    		// 停止计时器，清空channel
    		if !timer.Stop() && len(timer.C) > 0 {
    			<-timer.C
    		}
    		timer.Reset(sleepTime)
    		select {
    		case <-w.ctx.Done():
    			return
    		case <-timer.C: // time.sleep改为channel阻塞的形式, time.After每次都会创建一个timer
    		}
    	}
    }
    
    func NewWorker(task TaskFunc, opts ...WorkerOption) (w *Worker) {
    	w = initWorker(task)
    	for _, opt := range opts {
    		opt(w)
    	}
    	return
    }
    
    func initWorker(task TaskFunc) *Worker {
    	w := &Worker{
    		name:              defaultWorkerName,
    		busySleepTime:     defaultBusySleepTime,
    		idleSleepTime:     defaultIdleSleepTime,
    		task:              task,
    		panicRecoverCount: defaultPanicRecoverCount,
    		ctx:               context.TODO(),
    		taskFrequency:     0,
    		waitCloseTime:     defaultWaitCloseTime,
    	}
    	tmp := make(chan struct{})
    	w.checkMaster = tmp
    	close(tmp)
    	return w
    }
    
    func WithName(name string) WorkerOption {
    	return func(worker *Worker) {
    		worker.name = name
    	}
    }
    
    // 传入的channel应该是阻塞式的，如果是leader节点，需要每次判断都不阻塞。
    func WithCheckMasterChannel(checkMasterCh <-chan struct{}) WorkerOption {
    	return func(worker *Worker) {
    		worker.checkMaster = checkMasterCh
    	}
    }
    
    func WithBusySleepTime(t time.Duration) WorkerOption {
    	return func(worker *Worker) {
    		worker.busySleepTime = t
    	}
    }
    
    func WithIdleSleepTime(t time.Duration) WorkerOption {
    	return func(worker *Worker) {
    		worker.idleSleepTime = t
    	}
    }
    
    func WithWaitCloseTime(t time.Duration) WorkerOption {
    	return func(worker *Worker) {
    		worker.waitCloseTime = t
    	}
    }
    
    func WithPanicRecoverCount(c int) WorkerOption {
    	return func(worker *Worker) {
    		worker.panicRecoverCount = c
    	}
    }
    
    func WithContext(ctx context.Context) WorkerOption {
    	return func(worker *Worker) {
    		worker.ctx = ctx
    	}
    }
    
    // 传入的waitGroup，请在外面调用wg.Add()
    func WithWaitGroup(wg *sync.WaitGroup) WorkerOption {
    	return func(worker *Worker) {
    		worker.wg = wg
    	}
    }
    
    func WithTaskFrequency(d time.Duration) WorkerOption {
    	return func(worker *Worker) {
    		worker.taskFrequency = d
    	}
    }
    
    func DealRecover(r interface{}) {
    	fmt.Println("worker panic, now recover", r)
    }