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

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

如有侵权,请联系 zhuanlan_guanli@qq.com 删除。

编辑于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏bluesummer

.Net Core 学习之路-AutoFac的使用

本文不介绍IoC和DI的概念,如果你对Ioc之前没有了解的话,建议先去搜索一下相关的资料 这篇文章将简单介绍一下AutoFac的基本使用以及在asp .net ...

2057
来自专栏java思维导图

自己手写一个Mybatis框架(简化)

继上一篇手写SpringMVC之后,我最近趁热打铁,研究了一下Mybatis。MyBatis框架的核心功能其实不难,无非就是动态代理和jdbc的操作,难的是写出...

2756
来自专栏NetCore

一个让人遗忘的角落—Exception(二)

在上一篇中"一个被人遗忘的角落--Exception(一)"中,跟大家简单介绍了一下Exception,也使大家充分的了解了Exception管理在一个项目中的...

1739
来自专栏java初学

关于mybatis的思考(1)——mybatis的使用实例

33413
来自专栏电光石火

springboot整合websocket实现消息推送

1763
来自专栏WebApiClient

WebApiClient的JsonPatch局部更新

随着WebApiClient的不断完善,越来越多开发者选择WebApiClient替换原生的HttpClient,本文将介绍使用WebApiClient来完成J...

1514
来自专栏Seebug漏洞平台

Mirai源码分析

1. 背景概述 最近的德国断网事件让Mirai恶意程序再次跃入公众的视线,相对而言,目前的IoT领域对于恶意程序还是一片蓝海,因此吸引了越来越多的人开始涉足这趟...

3637
来自专栏阿杜的世界

Java Web技术经验总结(九)

这个配置的含义有二:(1)对于所有符合"/"模式的请求(除"/admin/"之外)要应用ThemeChangeInterceptor拦截器;(2)对于所有符合"...

793
来自专栏郭霖

Android图片加载框架最全解析(三),深入探究Glide的缓存机制

在本系列的上一篇文章中,我带着大家一起阅读了一遍Glide的源码,初步了解了这个强大的图片加载框架的基本执行流程。 不过,上一篇文章只能说是比较粗略地阅读了Gl...

39810
来自专栏云瓣

深入Redux架构

关于redux 之前写了一篇通过一个demo了解Redux,但对于redux的核心方法没有进行深入剖析,在此重新总结学习,完整的代码看这里。(参考了React ...

2476

扫描关注云+社区