go 并发处理脚本

package script

import (
	"fmt"
	"time"
	"errors"
	"flag"
	"pigcome/utils"
)

// 实现此接口即可
type model interface {
  	// 每个goroutine运行的函数
		RunOne(id int64) (params []int64, errInfos map[int64]string, err error)
  	// 最后收集到结果后的处理函数
		HandleResult(res Result)
}

type RunParam struct {
	Start    int64			//起始ID
	End      int64			//结束ID
	Step     int			//同时进行的goroutine数目
	Ids      []int64		//自定义IDs
	Deadline time.Duration	//goroutine超时限制
}

type ResultType int

const (
	_       ResultType = iota
	Success
	Panic
	Timeout
	Error
)

type Result struct {
	Type     ResultType			//结果类型
	ID       int64				//每个goroutine的ID
  	Ress     []int64			//结果的参数(自定义)
	ErrInfos map[int64]string	 //runOne函数自定义的错误信息
	ErrInfo  error				//整个goroutine的错误信息
}

func Run(params *RunParam, m model) (err error) {
	beginAt := time.Now()
	err = CheckParams(params)
	if err != nil {
		return
	}
	// 1000并没有什么特殊意义，只是为了有缓冲，提高速度
	ch := make(chan Result, 1000) 
  // 令牌，只有持有令牌才能运行，为了控制goroutine同时进行的数目
	token := make(chan struct{}, params.Step)
  // 以此判断结果是否都处理完成
	done := make(chan struct{})

	go collectResult(params, ch, done, m)
	if len(params.Ids) > 0 {
		for _, id := range params.Ids {
			token <- struct{}{}
			go runOne(params, ch, m, token, id)
		}
	} else {
		for id := params.Start; id < params.End; id++ {
			token <- struct{}{}
			go runOne(params, ch, m, token, id)
		}
	}
	<-done

	since := time.Since(beginAt)
	utils.PrintlnColorful(utils.Yellow, "耗时：", since.String())
	return
}

func collectResult(params *RunParam, ch chan Result, done chan struct{}, m model) {
	var finishCount, count int64
	idsLength := len(params.Ids)
	if idsLength > 0 {
		count = int64(idsLength)
	} else {
		count = params.End - params.Start
	}
	for {
		result := <-ch
		m.HandleResult(result)
		finishCount++
		if result.ErrInfo != nil {
			utils.PrintfColorful(utils.Red, "%d: %+v\n", finishCount, result)
		}
		if finishCount >= count {
			done <- struct{}{}
			return
		}
	}
}

func runOne(params *RunParam, ch chan Result, m model, token chan struct{}, id int64) {
	defer func() {
		<-token
		if err := recover(); err != nil {
			ch <- Result{Panic, id, nil, nil, fmt.Errorf("panic: %v", err)}
		}
	}()

	type runOneCh struct {
		ress     []int64
		errInfos map[int64]string
		err      error
	}
	errCh := make(chan runOneCh)

	go func(ch chan runOneCh, id int64) {
		ress, errInfos, err := m.RunOne(id)
		errCh <- runOneCh{ress, errInfos, err}
	}(errCh, id)

	select {
	case runOneRes := <-errCh:
		if runOneRes.err != nil {
			ch <- Result{Error, id, nil, nil, runOneRes.err}
		} else {
			ch <- Result{Success, id, runOneRes.ress, nil, nil}
		}
	case <-time.After(params.Deadline):
		ch <- Result{Timeout, id, nil, nil, errors.New("timeout")}
	}
}

// 初始化运行参数
func NewParams() (params *RunParam) {
	params = new(RunParam)
	flag.Int64Var(&params.Start, "start", 0, "")
	flag.Int64Var(&params.End, "end", 0, "")
	flag.IntVar(&params.Step, "step", 100, "")
	var second int
	flag.IntVar(&second, "second", 5, "")
	flag.Parse()
	params.Deadline = time.Second * time.Duration(second)
	return
}

// 运行参数检测
func CheckParams(params *RunParam) error {
	if params.End <= params.Start {
		return fmt.Errorf("ress range err, %+v", *params)
	}
	if params.Deadline <= 0 {
		return fmt.Errorf("ress deadline err, %+v", *params)
	}
	return nil
}

const (
	Red    = "\033[31m"
	Yellow = "\033[33m"
	Green  = "\033[32m"
)

func PrintlnColorful(color string, vals ...interface{}) {
	fmt.Printf(color)
	fmt.Println(vals...)
	fmt.Printf("\033[0m")
}

func PrintfColorful(color string, format string, vals ...interface{}) {
	fmt.Printf(color)
	fmt.Printf(format, vals...)
	fmt.Printf("\033[0m")
}