func (b *Broker) HandleTaskResult(uuid string) (*task.Reply, error) {
if len(uuid) == 0 {
return nil, errors.ErrInvalidArgument
}
key := fmt.Sprintf("r_%s", uuid)
result, err := b.redisClient.HMGet(key,
"is_success",
"result",
).Result()
if err != nil {
golog.Error("Broker", "HandleFailTask", err.Error(), 0, "key", key)
return nil, err
}
//key不存在
if result[0] == nil {
return nil, errors.ErrResultNotExist
}
isSuccess, err := strconv.Atoi(result[0].(string))
if err != nil {
return nil, err
}
ret := result[1].(string)
return &task.Reply{
IsResultExist: 1,
IsSuccess: isSuccess,
Result: ret,
}, nil
}
func (b *Broker) resetTaskRequest(args []interface{}) error {
var err error
if len(args) == 0 || len(args) != config.TaskRequestItemCount {
return errors.ErrInvalidArgument
}
request := new(task.TaskRequest)
request.Uuid = args[0].(string)
request.BinName = args[1].(string)
request.Args = args[2].(string)
request.StartTime, err = strconv.ParseInt(args[3].(string), 10, 64)
if err != nil {
return err
}
request.TimeInterval = args[4].(string)
request.Index, err = strconv.Atoi(args[5].(string))
if err != nil {
return err
}
vec := strings.Split(request.TimeInterval, " ")
request.Index++
if request.Index < len(vec) {
timeLater, err := strconv.Atoi(vec[request.Index])
if err != nil {
return err
}
afterTime := time.Second * time.Duration(timeLater)
b.timer.NewTimer(afterTime, b.AddRequestToRedis, request)
} else {
golog.Error("Broker", "HandleFailTask", "retry max time", 0,
"key", fmt.Sprintf("t_%s", request.Uuid))
return errors.ErrTryMaxTimes
}
return nil
}
func NewWorker(cfg *config.WorkerConfig) (*Worker, error) {
var err error
w := new(Worker)
w.cfg = cfg
vec := strings.SplitN(cfg.RedisAddr, "/", 2)
if len(vec) == 2 {
w.redisAddr = vec[0]
w.redisDB, err = strconv.Atoi(vec[1])
if err != nil {
return nil, err
}
} else {
w.redisAddr = vec[0]
w.redisDB = config.DefaultRedisDB
}
w.redisClient = redis.NewClient(
&redis.Options{
Addr: w.redisAddr,
Password: "", // no password set
DB: int64(w.redisDB),
},
)
_, err = w.redisClient.Ping().Result()
if err != nil {
golog.Error("worker", "NewWorker", "ping redis fail", 0, "err", err.Error())
return nil, err
}
return w, nil
}
func (w *Worker) DoScriptTaskRequest(req *task.TaskRequest) (string, error) {
var output string
var err error
var maxRunTime int64
binPath := path.Clean(w.cfg.BinPath + "/" + req.BinName)
_, err = os.Stat(binPath)
if err != nil && os.IsNotExist(err) {
golog.Error("worker", "DoScrpitTaskRequest", "File not exist", 0,
"key", fmt.Sprintf("t_%s", req.Uuid),
"bin_path", binPath,
)
return "", errors.ErrFileNotExist
}
if req.MaxRunTime == 0 {
maxRunTime = w.cfg.TaskRunTime
} else {
maxRunTime = req.MaxRunTime
}
if len(req.Args) == 0 {
output, err = w.ExecBin(binPath, nil, maxRunTime)
} else {
argsVec := strings.Split(req.Args, " ")
output, err = w.ExecBin(binPath, argsVec, maxRunTime)
}
return output, err
}
func (w *Worker) CmdRunWithTimeout(cmd *exec.Cmd, timeout time.Duration) (error, bool) {
done := make(chan error)
go func() {
done <- cmd.Wait()
}()
var err error
select {
case <-time.After(timeout):
// timeout
if err = cmd.Process.Kill(); err != nil {
golog.Error("worker", "CmdRunTimeout", "kill error", 0,
"path", cmd.Path,
"error", err.Error(),
)
}
golog.Info("worker", "CmdRunWithTimeout", "kill process", 0,
"path", cmd.Path,
"error", errors.ErrExecTimeout.Error(),
)
go func() {
<-done // allow goroutine to exit
}()
return errors.ErrExecTimeout, true
case err = <-done:
return err, false
}
}
func main() {
runtime.GOMAXPROCS(runtime.NumCPU())
flag.Parse()
if len(*configFile) == 0 {
fmt.Println("must use a config file")
return
}
cfg, err := config.ParseWorkerConfigFile(*configFile)
if err != nil {
fmt.Printf("parse config file error:%v\n", err.Error())
return
}
//when the log file size greater than 1GB, kingtask will generate a new file
if len(cfg.LogPath) != 0 {
sysFilePath := path.Join(cfg.LogPath, sysLogName)
sysFile, err := golog.NewRotatingFileHandler(sysFilePath, MaxLogSize, 1)
if err != nil {
fmt.Printf("new log file error:%v\n", err.Error())
return
}
golog.GlobalLogger = golog.New(sysFile, golog.Lfile|golog.Ltime|golog.Llevel)
}
if *logLevel != "" {
setLogLevel(*logLevel)
} else {
setLogLevel(cfg.LogLevel)
}
var w *worker.Worker
w, err = worker.NewWorker(cfg)
if err != nil {
golog.Error("main", "main", err.Error(), 0)
golog.GlobalLogger.Close()
w.Close()
return
}
sc := make(chan os.Signal, 1)
signal.Notify(sc,
syscall.SIGHUP,
syscall.SIGINT,
syscall.SIGTERM,
syscall.SIGQUIT)
go func() {
sig := <-sc
golog.Info("main", "main", "Got signal", 0, "signal", sig)
golog.GlobalLogger.Close()
w.Close()
}()
golog.Info("main", "main", "Worker start!", 0)
w.Run()
}
func (b *Broker) handleConn(c net.Conn) error {
defer func() {
r := recover()
if err, ok := r.(error); ok {
const size = 4096
buf := make([]byte, size)
buf = buf[:runtime.Stack(buf, false)]
golog.Error("Broker", "handleConn",
err.Error(), 0,
"stack", string(buf))
}
c.Close()
}()
var CloseConn bool
reader := bufio.NewReaderSize(c, 1024)
for {
msgType := []byte{0}
if _, err := io.ReadFull(reader, msgType); err != nil {
return errors.ErrBadConn
}
switch msgType[0] {
case config.TypeRequestTask:
b.HandleRequest(reader, c)
case config.TypeGetTaskResult:
b.HandleTaskResult(reader, c)
case config.TypeCloseConn:
CloseConn = true
default:
golog.Error("Broker", "handleConn", "msgType error", 0, "msg_type", msgType[0])
CloseConn = true
}
if CloseConn {
break
}
}
return nil
}
//处理失败的任务
func (b *Broker) HandleFailTask() error {
var uuid string
var err error
for b.running {
uuid, err = b.redisClient.SPop(config.FailResultUuidSet).Result()
//没有结果，直接返回
if err == redis.Nil {
time.Sleep(time.Second)
continue
}
if err != nil {
golog.Error("Broker", "HandleFailTask", "spop error", 0, "error", err.Error())
continue
}
key := fmt.Sprintf("r_%s", uuid)
timeInterval, err := b.redisClient.HGet(key, "time_interval").Result()
if err != nil {
golog.Error("Broker", "HandleFailTask", err.Error(), 0, "key", key)
continue
}
//没有超时重试机制
if len(timeInterval) == 0 {
b.SetFailTaskCount(fmt.Sprintf("t_%s", uuid))
continue
}
//获取结果中所有值,改为逐个获取
results, err := b.redisClient.HMGet(key,
"uuid",
"bin_name",
"args",
"start_time",
"time_interval",
"index",
"max_run_time",
"task_type",
).Result()
if err != nil {
golog.Error("Broker", "HandleFailTask", err.Error(), 0, "key", key)
continue
}
//key已经过期
if results[0] == nil {
golog.Error("Broker", "HandleFailTask", "result expired", 0, "key", key)
continue
}
//删除结果
_, err = b.redisClient.Del(key).Result()
if err != nil {
golog.Error("Broker", "HandleFailTask", "delete result failed", 0, "key", key)
}
err = b.resetTaskRequest(results)
if err != nil {
golog.Error("Broker", "HandleFailTask", err.Error(), 0, "key", key)
b.SetFailTaskCount(fmt.Sprintf("t_%s", uuid))
}
}
return nil
}
func (w *Worker) SetSuccessTaskCount(reqKey string) error {
successTaskKey := fmt.Sprintf(config.SuccessTaskKey,
time.Now().Format(config.TimeFormat))
count, err := w.redisClient.Incr(successTaskKey).Result()
if err != nil {
golog.Error("Worker", "SetSuccessTaskCount", "Incr", 0, "err", err.Error(),
"req_key", reqKey)
return err
}
//第一次设置该key
if count == 1 {
//保存一个月
expireTime := time.Second * time.Duration(60*60*24*30)
_, err = w.redisClient.Expire(successTaskKey, expireTime).Result()
if err != nil {
golog.Error("Worker", "SetSuccessTaskCount", "Expire", 0, "err", err.Error(),
"req_key", reqKey)
return err
}
}
return nil
}
func (w *Worker) Run() error {
w.running = true
for w.running {
uuid, err := w.redisClient.SPop(config.RequestUuidSet).Result()
//没有请求
if err == redis.Nil {
time.Sleep(time.Second)
continue
}
if err != nil {
golog.Error("Worker", "run", "spop error", 0, "error", err.Error())
continue
}
reqKey := fmt.Sprintf("t_%s", uuid)
//获取请求中所有值
request, err := w.redisClient.HMGet(reqKey,
"uuid",
"bin_name",
"args",
"start_time",
"time_interval",
"index",
).Result()
if err != nil {
golog.Error("Worker", "run", err.Error(), 0, "req_key", reqKey)
continue
}
//key不存在
if request[0] == nil {
golog.Error("Worker", "run", "Key is not exist", 0, "req_key", reqKey)
continue
}
_, err = w.redisClient.Del(reqKey).Result()
if err != nil {
golog.Error("Worker", "run", "delete result failed", 0, "req_key", reqKey)
}
taskResult, err := w.DoTaskRequest(request)
if err != nil {
golog.Error("Worker", "run", "DoTaskRequest", 0, "err", err.Error(),
"req_key", reqKey)
}
if taskResult != nil {
err = w.SetTaskResult(taskResult)
if err != nil {
golog.Error("Worker", "run", "DoTaskRequest", 0,
"err", err.Error(), "req_key", reqKey)
}
golog.Info("worker", "run", "do task success", 0, "req_key", reqKey,
"result", taskResult.Result)
}
if w.cfg.Peroid != 0 {
time.Sleep(time.Second * time.Duration(w.cfg.Peroid))
}
}
return nil
}
func (b *Broker) AddRequestToRedis(tr interface{}) error {
r, ok := tr.(*task.TaskRequest)
if !ok {
return errors.ErrInvalidArgument
}
key := fmt.Sprintf("t_%s", r.Uuid)
setCmd := b.redisClient.HMSet(key,
"uuid", r.Uuid,
"bin_name", r.BinName,
"args", r.Args,
"start_time", strconv.FormatInt(r.StartTime, 10),
"time_interval", r.TimeInterval,
"index", strconv.Itoa(r.Index),
"max_run_time", strconv.FormatInt(r.MaxRunTime, 10),
"task_type", strconv.Itoa(r.TaskType),
)
err := setCmd.Err()
if err != nil {
golog.Error("Broker", "AddRequestToRedis", "HMSET error", 0,
"set", config.RequestUuidSet,
"uuid", r.Uuid,
"err", err.Error(),
)
return err
}
saddCmd := b.redisClient.SAdd(config.RequestUuidSet, r.Uuid)
err = saddCmd.Err()
if err != nil {
golog.Error("Broker", "AddRequestToRedis", "SADD error", 0,
"set", config.RequestUuidSet,
"uuid", r.Uuid,
"err", err.Error(),
)
return err
}
return nil
}
func (w *Worker) DoTaskRequest(args []interface{}) (*task.TaskResult, error) {
var err error
var output string
req := new(task.TaskRequest)
ret := new(task.TaskResult)
req.Uuid = args[0].(string)
req.BinName = args[1].(string)
req.Args = args[2].(string)
req.StartTime, err = strconv.ParseInt(args[3].(string), 10, 64)
if err != nil {
return nil, err
}
req.TimeInterval = args[4].(string)
req.Index, err = strconv.Atoi(args[5].(string))
if err != nil {
return nil, err
}
req.MaxRunTime, err = strconv.ParseInt(args[6].(string), 10, 64)
if err != nil {
return nil, err
}
req.TaskType, err = strconv.Atoi(args[7].(string))
if err != nil {
return nil, err
}
switch req.TaskType {
case task.ScriptTask:
output, err = w.DoScriptTaskRequest(req)
case task.RpcTaskGET, task.RpcTaskPOST, task.RpcTaskPUT, task.RpcTaskDELETE:
output, err = w.DoRpcTaskRequest(req)
default:
err = errors.ErrInvalidArgument
golog.Error("Worker", "DoTaskRequest", "task type error", 0, "task_type", req.TaskType)
}
ret.TaskRequest = *req
//执行任务失败，
if err != nil {
ret.IsSuccess = int64(0)
ret.Result = err.Error()
return ret, nil
}
ret.IsSuccess = int64(1)
ret.Result = output
return ret, nil
}
func (b *Broker) Run() error {
b.running = true
go b.HandleFailTask()
for b.running {
conn, err := b.listener.Accept()
if err != nil {
golog.Error("server", "Run", err.Error(), 0)
continue
}
//处理客户端请求
go b.handleConn(conn)
}
return nil
}
func (w *Worker) DoTaskRequest(args []interface{}) (*task.TaskResult, error) {
var err error
var output string
req := new(task.TaskRequest)
ret := new(task.TaskResult)
req.Uuid = args[0].(string)
req.BinName = args[1].(string)
req.Args = args[2].(string)
req.StartTime, err = strconv.ParseInt(args[3].(string), 10, 64)
if err != nil {
return nil, err
}
req.TimeInterval = args[4].(string)
req.Index, err = strconv.Atoi(args[5].(string))
if err != nil {
return nil, err
}
binPath := path.Clean(w.cfg.BinPath + "/" + req.BinName)
_, err = os.Stat(binPath)
if err != nil && os.IsNotExist(err) {
golog.Error("worker", "DoTaskRequest", "File not exist", 0,
"key", fmt.Sprintf("t_%s", req.Uuid),
"bin_path", binPath,
)
return nil, errors.ErrFileNotExist
}
if len(req.Args) == 0 {
output, err = w.ExecBin(binPath, nil)
} else {
argsVec := strings.Split(req.Args, " ")
output, err = w.ExecBin(binPath, argsVec)
}
ret.TaskRequest = *req
//执行任务失败
if err != nil {
ret.IsSuccess = int64(0)
ret.Result = err.Error()
return ret, nil
}
ret.IsSuccess = int64(1)
ret.Result = output
return ret, nil
}
func NewBroker(cfg *config.BrokerConfig) (*Broker, error) {
var err error
broker := new(Broker)
broker.cfg = cfg
broker.addr = cfg.Addr
if len(broker.addr) == 0 {
return nil, errors.ErrInvalidArgument
}
vec := strings.SplitN(cfg.RedisAddr, "/", 2)
if len(vec) == 2 {
broker.redisAddr = vec[0]
broker.redisDB, err = strconv.Atoi(vec[1])
if err != nil {
return nil, err
}
} else {
broker.redisAddr = vec[0]
broker.redisDB = config.DefaultRedisDB
}
broker.web = echo.New()
broker.timer = timer.New(time.Millisecond * 10)
go broker.timer.Start()
broker.redisClient = redis.NewClient(
&redis.Options{
Addr: broker.redisAddr,
Password: "", // no password set
DB: int64(broker.redisDB),
},
)
_, err = broker.redisClient.Ping().Result()
if err != nil {
golog.Error("broker", "NewBroker", "ping redis fail", 0, "err", err.Error())
return nil, err
}
return broker, nil
}
func (b *Broker) HandleTaskResult(rb *bufio.Reader, c net.Conn) error {
args := struct {
Key string `json:"key"`
}{}
buf := make([]byte, 128)
readLen, err := rb.Read(buf)
if err != nil {
b.WriteError(err, c)
return err
}
err = json.Unmarshal(buf[:readLen], &args)
if err != nil {
b.WriteError(err, c)
return err
}
result, err := b.redisClient.HMGet(args.Key,
"is_success",
"result",
).Result()
if err != nil {
golog.Error("Broker", "HandleFailTask", err.Error(), 0, "key", args.Key)
b.WriteError(err, c)
return err
}
//key不存在
if result[0] == nil {
err = b.WriteResult(config.ResultNotExist, "0", "", c)
return err
}
isSuccess := result[0].(string)
ret := result[1].(string)
return b.WriteResult(config.ResultIsExist, isSuccess, ret, c)
}