// The worker package helps developers to develop Gearman's worker // in an easy way. package worker import ( "fmt" "strconv" "sync" "time" ) const ( Unlimited = iota OneByOne ) // Worker is the only structure needed by worker side developing. // It can connect to multi-server and grab jobs. type Worker struct { sync.Mutex agents []*agent funcs jobFuncs in chan *inPack running bool ready bool jobLeftNum int64 Id string ErrorHandler ErrorHandler JobHandler JobHandler limit chan bool closed chan struct{} leftJobs chan struct{} } // New returns a worker. // // If limit is set to Unlimited(=0), the worker will grab all jobs // and execute them parallelly. // If limit is greater than zero, the number of paralled executing // jobs are limited under the number. If limit is assgined to // OneByOne(=1), there will be only one job executed in a time. func New(limit int) (worker *Worker) { worker = &Worker{ agents: make([]*agent, 0, limit), funcs: make(jobFuncs), in: make(chan *inPack, queueSize), } if limit != Unlimited { worker.limit = make(chan bool, limit-1) } return } // inner error handling func (worker *Worker) err(e error) { if worker.ErrorHandler != nil { worker.ErrorHandler(e) } } // AddServer adds a Gearman job server. // // addr should be formated as 'host:port'. func (worker *Worker) AddServer(net, addr string) (err error) { // Create a new job server's client as a agent of server a, err := newAgent(net, addr, worker) if err != nil { return err } worker.agents = append(worker.agents, a) return } // Broadcast an outpack to all Gearman server. func (worker *Worker) broadcast(outpack *outPack) { for _, v := range worker.agents { v.Write(outpack) } } // AddFunc adds a function. // Set timeout as Unlimited(=0) to disable executing timeout. func (worker *Worker) AddFunc(funcname string, f JobFunc, timeout uint32) (err error) { worker.Lock() defer worker.Unlock() if _, ok := worker.funcs[funcname]; ok { return fmt.Errorf("The function already exists: %s", funcname) } worker.funcs[funcname] = &jobFunc{f: f, timeout: timeout} if worker.running { worker.addFunc(funcname, timeout) } return } // inner add func (worker *Worker) addFunc(funcname string, timeout uint32) { outpack := prepFuncOutpack(funcname, timeout) worker.broadcast(outpack) } func prepFuncOutpack(funcname string, timeout uint32) *outPack { outpack := getOutPack() if timeout == 0 { outpack.dataType = dtCanDo outpack.data = []byte(funcname) } else { outpack.dataType = dtCanDoTimeout l := len(funcname) timeoutString := strconv.FormatUint(uint64(timeout), 10) outpack.data = getBuffer(l + len(timeoutString) + 1) copy(outpack.data, []byte(funcname)) outpack.data[l] = '\x00' copy(outpack.data[l+1:], []byte(timeoutString)) } return outpack } // RemoveFunc removes a function. func (worker *Worker) RemoveFunc(funcname string) (err error) { worker.Lock() defer worker.Unlock() if _, ok := worker.funcs[funcname]; !ok { return fmt.Errorf("The function does not exist: %s", funcname) } delete(worker.funcs, funcname) if worker.running { worker.removeFunc(funcname) } return } // inner remove func (worker *Worker) removeFunc(funcname string) { outpack := getOutPack() outpack.dataType = dtCantDo outpack.data = []byte(funcname) worker.broadcast(outpack) } // inner package handling func (worker *Worker) handleInPack(inpack *inPack) { switch inpack.dataType { case dtNoJob: inpack.a.PreSleep() case dtNoop: inpack.a.Grab() case dtJobAssign, dtJobAssignUniq: go func() { go func() { worker.incrExecJobNum() defer func() { worker.decrExecJobNum() }() if err := worker.exec(inpack); err != nil { worker.err(err) } }() if worker.limit != nil { worker.limit <- true } inpack.a.Grab() }() case dtError: worker.err(inpack.Err()) fallthrough case dtEchoRes: fallthrough default: worker.customeHandler(inpack) } } // Connect to Gearman server and tell every server // what can this worker do. func (worker *Worker) Ready() (err error) { if len(worker.agents) == 0 { return ErrNoneAgents } if len(worker.funcs) == 0 { return ErrNoneFuncs } for _, a := range worker.agents { if err = a.Connect(); err != nil { return } } for funcname, f := range worker.funcs { worker.addFunc(funcname, f.timeout) } worker.ready = true return } // Work start main loop (blocking) // Most of time, this should be evaluated in goroutine. func (worker *Worker) Work() { if !worker.ready { // didn't run Ready beforehand, so we'll have to do it: err := worker.Ready() if err != nil { panic(err) } } worker.Lock() worker.running = true worker.Unlock() for _, a := range worker.agents { a.Grab() } // 执行任务(阻塞) var inpack *inPack for inpack = range worker.in { worker.handleInPack(inpack) } // 关闭Worker进程后 等待任务完成后退出 worker.Lock() leftJobNum := int(worker.jobLeftNum) worker.Unlock() if worker.leftJobs != nil { for i := 0; i < leftJobNum; i++ { <-worker.leftJobs } } worker.Reset() worker.close() } // custome handling warper func (worker *Worker) customeHandler(inpack *inPack) { if worker.JobHandler != nil { if err := worker.JobHandler(inpack); err != nil { worker.err(err) } } } // Close connection and exit main loop func (worker *Worker) Close() { worker.Lock() defer worker.Unlock() if worker.running == true && worker.closed == nil { worker.closed = make(chan struct{}, 1) worker.closed <- struct{}{} worker.running = false close(worker.in) // 创建关闭后执行中的任务列表 if worker.jobLeftNum != 0 { worker.leftJobs = make(chan struct{}, worker.jobLeftNum+int64(len(worker.in))) } } } func (worker *Worker) close() { for _, a := range worker.agents { a.Close() } } // Echo func (worker *Worker) Echo(data []byte) { outpack := getOutPack() outpack.dataType = dtEchoReq outpack.data = data worker.broadcast(outpack) } // Reset removes all of functions. // Both from the worker and job servers. func (worker *Worker) Reset() { outpack := getOutPack() outpack.dataType = dtResetAbilities worker.broadcast(outpack) worker.funcs = make(jobFuncs) } // Set the worker's unique id. func (worker *Worker) SetId(id string) { worker.Id = id outpack := getOutPack() outpack.dataType = dtSetClientId outpack.data = []byte(id) worker.broadcast(outpack) } func (worker *Worker) incrExecJobNum() int64 { worker.Lock() defer worker.Unlock() worker.jobLeftNum++ return worker.jobLeftNum } func (worker *Worker) decrExecJobNum() int64 { worker.Lock() defer worker.Unlock() worker.jobLeftNum-- if worker.jobLeftNum < 0 { worker.jobLeftNum = 0 } return worker.jobLeftNum } // inner job executing func (worker *Worker) exec(inpack *inPack) (err error) { defer func() { if worker.limit != nil { <-worker.limit } if r := recover(); r != nil { if e, ok := r.(error); ok { err = e } else { err = ErrUnknown } } }() f, ok := worker.funcs[inpack.fn] if !ok { return fmt.Errorf("The function does not exist: %s", inpack.fn) } var r *result if f.timeout == 0 { d, e := f.f(inpack) r = &result{data: d, err: e} } else { r = execTimeout(f.f, inpack, time.Duration(f.timeout)*time.Second) } //if worker.running { outpack := getOutPack() if r.err == nil { outpack.dataType = dtWorkComplete } else { if len(r.data) == 0 { outpack.dataType = dtWorkFail } else { outpack.dataType = dtWorkException } err = r.err } outpack.handle = inpack.handle outpack.data = r.data _ = inpack.a.Write(outpack) if worker.leftJobs != nil { worker.leftJobs <- struct{}{} } //} return } func (worker *Worker) reRegisterFuncsForAgent(a *agent) { worker.Lock() defer worker.Unlock() for funcname, f := range worker.funcs { outpack := prepFuncOutpack(funcname, f.timeout) a.write(outpack) } } // inner result type result struct { data []byte err error } // executing timer func execTimeout(f JobFunc, job Job, timeout time.Duration) (r *result) { rslt := make(chan *result) defer close(rslt) go func() { defer func() { recover() }() d, e := f(job) rslt <- &result{data: d, err: e} }() select { case r = <-rslt: case <-time.After(timeout): return &result{err: ErrTimeOut} } return r } // Error type passed when a worker connection disconnects type WorkerDisconnectError struct { err error agent *agent } func (e *WorkerDisconnectError) Error() string { return e.err.Error() } // Responds to the error by asking the worker to reconnect func (e *WorkerDisconnectError) Reconnect() (err error) { return e.agent.reconnect() } // Which server was this for? func (e *WorkerDisconnectError) Server() (net string, addr string) { return e.agent.net, e.agent.addr }