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
}

另外,用到了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")
}

原创声明,本文系作者授权云+社区发表,未经许可,不得转载。

如有侵权,请联系 yunjia_community@tencent.com 删除。

编辑于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏张善友的专栏

Mix 10 上的asp.net mvc 2的相关Session

Beyond File | New Company: From Cheesy Sample to Social Platform Scott Hansel...

2567
来自专栏陈仁松博客

ASP.NET Core 'Microsoft.Win32.Registry' 错误修复

今天在发布Asp.net Core应用到Azure的时候出现错误InvalidOperationException: Cannot find compilati...

4848
来自专栏hbbliyong

WPF Trigger for IsSelected in a DataTemplate for ListBox items

<DataTemplate DataType="{x:Type vm:HeaderSlugViewModel}"> <vw:HeaderSlug...

4064
来自专栏张善友的专栏

LINQ via C# 系列文章

LINQ via C# Recently I am giving a series of talk on LINQ. the name “LINQ via C...

2645
来自专栏魂祭心

原 canvas绘制clock

4064
来自专栏pangguoming

Spring Boot集成JasperReports生成PDF文档

由于工作需要,要实现后端根据模板动态填充数据生成PDF文档,通过技术选型,使用Ireport5.6来设计模板,结合JasperReports5.6工具库来调用渲...

1.2K7
来自专栏菩提树下的杨过

Flash/Flex学习笔记(23):运动学原理

先写一个公用的小球类Ball: package{ import flash.display.Sprite; //小球 类 public class B...

25310
来自专栏大内老A

The .NET of Tomorrow

Ed Charbeneau(http://developer.telerik.com/featured/the-net-of-tomorrow/) Exciti...

31610
来自专栏我和未来有约会

Silverlight第三方控件专题

这里我收集整理了目前网上silverlight第三方控件的专题,若果有所遗漏请告知我一下。 名称 简介 截图 telerik 商 RadC...

4015
来自专栏转载gongluck的CSDN博客

cocos2dx 打灰机

#include "GamePlane.h" #include "PlaneSprite.h" #include "BulletNode.h" #include...

5426

扫码关注云+社区