最近经常涉及到脚本的编写。本身项目数据量较大,所以经常编写的脚本需要高并发,干脆就提取出来。
如果有地方用到,只需要实现接口即可。
谨以此文抛砖引玉,不喜勿喷
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(¶ms.Start, "start", 0, "")
flag.Int64Var(¶ms.End, "end", 0, "")
flag.IntVar(¶ms.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
}
另外,用到了utils的的相关代码,下面贴一下
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")
}
原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。
原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。