gearman-go/worker/worker.go

351 lines
7.4 KiB
Go
Raw Normal View History

2013-12-25 18:11:01 +08:00
// The worker package helps developers to develop Gearman's worker
// in an easy way.
package worker
import (
2013-12-23 17:05:42 +08:00
"encoding/binary"
2013-08-30 18:01:10 +08:00
"fmt"
"sync"
2013-12-23 17:05:42 +08:00
"time"
)
const (
2013-12-26 12:06:47 +08:00
Unlimited = iota
OneByOne
)
2012-06-01 14:28:10 +08:00
2013-12-25 18:11:01 +08:00
// Worker is the only structure needed by worker side developing.
// It can connect to multi-server and grab jobs.
type Worker struct {
sync.Mutex
2013-12-24 17:41:11 +08:00
agents []*agent
funcs jobFuncs
in chan *inPack
2013-08-30 12:36:57 +08:00
running bool
2015-01-06 11:34:39 +08:00
ready bool
2013-08-30 12:36:57 +08:00
2013-12-25 18:11:01 +08:00
Id string
2013-08-30 18:01:10 +08:00
ErrorHandler ErrorHandler
2013-12-23 17:05:42 +08:00
JobHandler JobHandler
limit chan bool
}
2013-12-25 18:11:01 +08:00
// Return 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) {
2013-08-30 12:36:57 +08:00
worker = &Worker{
2013-12-25 17:01:42 +08:00
agents: make([]*agent, 0, limit),
funcs: make(jobFuncs),
in: make(chan *inPack, queueSize),
2013-08-30 12:36:57 +08:00
}
if limit != Unlimited {
2013-12-25 17:01:42 +08:00
worker.limit = make(chan bool, limit-1)
}
2013-08-30 12:36:57 +08:00
return
}
2013-12-25 18:11:01 +08:00
// inner error handling
2013-08-30 12:36:57 +08:00
func (worker *Worker) err(e error) {
2013-08-30 18:01:10 +08:00
if worker.ErrorHandler != nil {
worker.ErrorHandler(e)
2013-08-30 12:36:57 +08:00
}
}
2013-12-25 18:11:01 +08:00
// Add a Gearman job server.
//
// addr should be formated as 'host:port'.
2013-08-30 18:01:10 +08:00
func (worker *Worker) AddServer(net, addr string) (err error) {
2013-08-30 12:36:57 +08:00
// Create a new job server's client as a agent of server
2013-08-30 18:01:10 +08:00
a, err := newAgent(net, addr, worker)
2013-08-30 12:36:57 +08:00
if err != nil {
return err
}
2013-12-24 17:41:11 +08:00
worker.agents = append(worker.agents, a)
2013-08-30 12:36:57 +08:00
return
}
2013-12-25 18:11:01 +08:00
// Broadcast an outpack to all Gearman server.
func (worker *Worker) broadcast(outpack *outPack) {
2013-08-30 12:36:57 +08:00
for _, v := range worker.agents {
v.write(outpack)
2013-08-30 12:36:57 +08:00
}
}
// Add a function.
2013-12-25 18:11:01 +08:00
// Set timeout as Unlimited(=0) to disable executing timeout.
func (worker *Worker) AddFunc(funcname string,
2013-08-30 12:36:57 +08:00
f JobFunc, timeout uint32) (err error) {
worker.Lock()
defer worker.Unlock()
2013-08-30 12:36:57 +08:00
if _, ok := worker.funcs[funcname]; ok {
2013-08-30 18:01:10 +08:00
return fmt.Errorf("The function already exists: %s", funcname)
2013-08-30 12:36:57 +08:00
}
worker.funcs[funcname] = &jobFunc{f: f, timeout: timeout}
if worker.running {
worker.addFunc(funcname, timeout)
}
return
}
2013-12-25 18:11:01 +08:00
// inner add
func (worker *Worker) addFunc(funcname string, timeout uint32) {
2015-01-06 11:34:39 +08:00
outpack := prepFuncOutpack(funcname, timeout)
worker.broadcast(outpack)
}
2015-01-06 11:34:39 +08:00
func prepFuncOutpack(funcname string, timeout uint32) *outPack {
outpack := getOutPack()
2013-08-30 12:36:57 +08:00
if timeout == 0 {
outpack.dataType = dtCanDo
outpack.data = []byte(funcname)
2013-08-30 12:36:57 +08:00
} else {
outpack.dataType = dtCanDoTimeout
2013-08-30 18:01:10 +08:00
l := len(funcname)
outpack.data = getBuffer(l + 5)
copy(outpack.data, []byte(funcname))
outpack.data[l] = '\x00'
2013-12-23 17:05:42 +08:00
binary.BigEndian.PutUint32(outpack.data[l+1:], timeout)
2013-08-30 12:36:57 +08:00
}
return outpack
}
// Remove a function.
func (worker *Worker) RemoveFunc(funcname string) (err error) {
worker.Lock()
defer worker.Unlock()
2013-08-30 12:36:57 +08:00
if _, ok := worker.funcs[funcname]; !ok {
2013-08-30 18:01:10 +08:00
return fmt.Errorf("The function does not exist: %s", funcname)
2013-08-30 12:36:57 +08:00
}
delete(worker.funcs, funcname)
if worker.running {
worker.removeFunc(funcname)
}
return
}
2013-12-25 18:11:01 +08:00
// inner remove
func (worker *Worker) removeFunc(funcname string) {
outpack := getOutPack()
outpack.dataType = dtCantDo
outpack.data = []byte(funcname)
worker.broadcast(outpack)
}
2013-12-25 18:11:01 +08:00
// inner package handling
func (worker *Worker) handleInPack(inpack *inPack) {
switch inpack.dataType {
case dtNoJob:
2013-12-23 17:01:01 +08:00
inpack.a.PreSleep()
case dtNoop:
2013-12-23 17:01:01 +08:00
inpack.a.Grab()
case dtJobAssign, dtJobAssignUniq:
go func() {
if err := worker.exec(inpack); err != nil {
worker.err(err)
}
}()
2013-12-25 17:01:42 +08:00
if worker.limit != nil {
worker.limit <- true
2013-08-30 12:36:57 +08:00
}
inpack.a.Grab()
case dtError:
2013-12-26 12:06:47 +08:00
worker.err(inpack.Err())
fallthrough
case dtEchoRes:
fallthrough
2013-08-30 12:36:57 +08:00
default:
worker.customeHandler(inpack)
2013-08-30 12:36:57 +08:00
}
2012-12-29 23:54:20 +08:00
}
2013-12-25 18:11:01 +08:00
// Connect to Gearman server and tell every server
// what can this worker do.
2013-12-23 17:01:01 +08:00
func (worker *Worker) Ready() (err error) {
2013-12-26 12:06:47 +08:00
if len(worker.agents) == 0 {
return ErrNoneAgents
}
if len(worker.funcs) == 0 {
return ErrNoneFuncs
}
2013-12-26 15:55:16 +08:00
for _, a := range worker.agents {
if err = a.Connect(); err != nil {
2013-12-23 17:01:01 +08:00
return
}
}
for funcname, f := range worker.funcs {
worker.addFunc(funcname, f.timeout)
}
worker.ready = true
2013-12-23 17:01:01 +08:00
return
}
2013-12-25 18:11:01 +08:00
// Main loop, block here
// Most of time, this should be evaluated in goroutine.
func (worker *Worker) Work() {
2015-01-06 11:34:39 +08:00
if !worker.ready {
// didn't run Ready beforehand, so we'll have to do it:
err := worker.Ready()
if err != nil {
2015-01-06 11:34:39 +08:00
panic(err)
}
}
2013-08-30 12:36:57 +08:00
worker.running = true
2013-12-25 18:11:01 +08:00
for _, a := range worker.agents {
a.Grab()
2013-12-24 14:35:33 +08:00
}
var inpack *inPack
for inpack = range worker.in {
worker.handleInPack(inpack)
2013-08-30 12:36:57 +08:00
}
}
2013-12-25 18:11:01 +08:00
// custome handling warper
func (worker *Worker) customeHandler(inpack *inPack) {
2013-08-30 12:36:57 +08:00
if worker.JobHandler != nil {
if err := worker.JobHandler(inpack); err != nil {
2013-08-30 12:36:57 +08:00
worker.err(err)
}
}
}
2013-12-25 18:11:01 +08:00
// Close connection and exit main loop
func (worker *Worker) Close() {
2013-12-26 12:06:47 +08:00
worker.Lock()
2015-01-06 11:34:39 +08:00
defer worker.Unlock()
2013-12-26 12:06:47 +08:00
if worker.running == true {
2015-01-06 11:34:39 +08:00
for _, a := range worker.agents {
a.Close()
}
2013-12-26 12:06:47 +08:00
worker.running = false
close(worker.in)
}
}
2013-12-25 18:11:01 +08:00
// Echo
func (worker *Worker) Echo(data []byte) {
outpack := getOutPack()
outpack.dataType = dtEchoReq
outpack.data = data
worker.broadcast(outpack)
}
// Remove all of functions.
2013-12-25 18:11:01 +08:00
// 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) {
2013-08-30 12:36:57 +08:00
worker.Id = id
outpack := getOutPack()
outpack.dataType = dtSetClientId
outpack.data = []byte(id)
worker.broadcast(outpack)
}
2013-12-25 18:11:01 +08:00
// inner job executing
func (worker *Worker) exec(inpack *inPack) (err error) {
2013-08-30 12:36:57 +08:00
defer func() {
2013-12-25 17:01:42 +08:00
if worker.limit != nil {
<-worker.limit
}
2013-08-30 12:36:57 +08:00
if r := recover(); r != nil {
if e, ok := r.(error); ok {
err = e
} else {
2013-08-30 18:01:10 +08:00
err = ErrUnknown
2013-08-30 12:36:57 +08:00
}
}
}()
f, ok := worker.funcs[inpack.fn]
2013-08-30 12:36:57 +08:00
if !ok {
return fmt.Errorf("The function does not exist: %s", inpack.fn)
2013-08-30 12:36:57 +08:00
}
var r *result
if f.timeout == 0 {
d, e := f.f(inpack)
2013-08-30 12:36:57 +08:00
r = &result{data: d, err: e}
} else {
r = execTimeout(f.f, inpack, time.Duration(f.timeout)*time.Second)
2013-08-30 12:36:57 +08:00
}
2013-12-24 14:35:33 +08:00
if worker.running {
outpack := getOutPack()
if r.err == nil {
outpack.dataType = dtWorkComplete
2013-08-30 12:36:57 +08:00
} else {
2013-12-24 14:35:33 +08:00
if len(r.data) == 0 {
outpack.dataType = dtWorkFail
2013-12-24 14:35:33 +08:00
} else {
outpack.dataType = dtWorkException
2013-12-24 14:35:33 +08:00
}
err = r.err
2013-08-30 12:36:57 +08:00
}
2013-12-24 14:35:33 +08:00
outpack.handle = inpack.handle
outpack.data = r.data
inpack.a.write(outpack)
2013-08-30 12:36:57 +08:00
}
return
}
2015-01-06 11:34:39 +08:00
func (worker *Worker) reRegisterFuncsForAgent(a *agent) {
worker.Lock()
defer worker.Unlock()
for funcname, f := range worker.funcs {
2015-01-06 11:34:39 +08:00
outpack := prepFuncOutpack(funcname, f.timeout)
a.write(outpack)
}
}
2013-12-25 18:11:01 +08:00
// inner result
type result struct {
2013-08-30 12:36:57 +08:00
data []byte
err error
}
2013-12-25 18:11:01 +08:00
// executing timer
2013-08-30 18:01:10 +08:00
func execTimeout(f JobFunc, job Job, timeout time.Duration) (r *result) {
2013-08-30 12:36:57 +08:00
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):
2013-08-30 18:01:10 +08:00
return &result{err: ErrTimeOut}
2013-08-30 12:36:57 +08:00
}
return r
}
// Error type passed when a worker connection disconnects
2015-01-06 11:34:39 +08:00
type WorkerDisconnectError struct {
err error
agent *agent
}
2015-01-06 11:34:39 +08:00
func (e *WorkerDisconnectError) Error() string {
return e.err.Error()
2014-06-10 10:46:21 +08:00
}
// Responds to the error by asking the worker to reconnect
2015-01-06 11:34:39 +08:00
func (e *WorkerDisconnectError) Reconnect() (err error) {
return e.agent.reconnect()
}
2015-01-06 11:34:39 +08:00
// Which server was this for?
func (e *WorkerDisconnectError) Server() (net string, addr string) {
return e.agent.net, e.agent.addr
}