golang常驻后台类worker模板

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)
}