gearman-go/worker/worker.go

290 lines
7.0 KiB
Go
Raw Normal View History

// Copyright 2011 Xing Xing <mikespook@gmail.com> All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.
package worker
import (
"bytes"
"bitbucket.org/mikespook/gearman-go/common"
)
const (
Unlimited = 0
OneByOne = 1
)
// Job handler
type JobHandler func(*Job) error
type JobFunc func(job *Job) ([]byte, error)
// The definition of the callback function.
type jobFunc struct {
f JobFunc
timeout uint32
}
// Map for added function.
type JobFuncs map[string]*jobFunc
/*
Worker side api for gearman
usage:
w = worker.New(worker.Unlimited)
w.AddFunction("foobar", foobar)
w.AddServer("127.0.0.1:4730")
w.Work() // Enter the worker's main loop
The definition of the callback function 'foobar' should suit for the type 'JobFunction'.
It looks like this:
func foobar(job *Job) (data []byte, err os.Error) {
//sth. here
//plaplapla...
return
}
*/
type Worker struct {
agents []*agent
funcs JobFuncs
in chan *Job
running bool
limit chan bool
Id string
// assign a ErrFunc to handle errors
ErrHandler common.ErrorHandler
JobHandler JobHandler
}
// Get a new worker
func New(l int) (worker *Worker) {
worker = &Worker{
agents: make([]*agent, 0),
funcs: make(JobFuncs),
in: make(chan *Job, common.QUEUE_SIZE),
}
if l != Unlimited {
worker.limit = make(chan bool, l)
for i := 0; i < l; i ++ {
worker.limit <- true
}
}
return
}
//
func (worker *Worker)err(e error) {
if worker.ErrHandler != nil {
worker.ErrHandler(e)
}
}
// Add a server. The addr should be 'host:port' format.
// The connection is established at this time.
func (worker *Worker) AddServer(addr string) (err error) {
// Create a new job server's client as a agent of server
server, err := newAgent(addr, worker)
if err != nil {
return err
}
worker.agents = append(worker.agents, server)
return
}
// Write a job to job server.
// Here, the job's mean is not the oraginal mean.
// Just looks like a network package for job's result or tell job server, there was a fail.
func (worker *Worker) broadcast(job *Job) {
for _, v := range worker.agents {
v.WriteJob(job)
}
}
// Add a function.
// Plz added job servers first, then functions.
// The API will tell every connected job server that 'I can do this'
func (worker *Worker) AddFunc(funcname string,
f JobFunc, timeout uint32) (err error) {
if _, ok := worker.funcs[funcname]; ok {
return common.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 function
func (worker *Worker) addFunc(funcname string, timeout uint32) {
var datatype uint32
var data []byte
if timeout == 0 {
datatype = common.CAN_DO
data = []byte(funcname)
} else {
datatype = common.CAN_DO_TIMEOUT
data = []byte(funcname + "\x00")
t := common.Uint32ToBytes(timeout)
data = append(data, t[:]...)
}
job := newJob(common.REQ, datatype, data)
worker.broadcast(job)
}
// Remove a function.
// Tell job servers 'I can not do this now' at the same time.
func (worker *Worker) RemoveFunc(funcname string) (err error) {
if _, ok := worker.funcs[funcname]; !ok {
return common.Errorf("The function does not exist: %s", funcname)
}
delete(worker.funcs, funcname)
if worker.running {
worker.removeFunc(funcname)
}
return
}
// inner remove function
func (worker *Worker) removeFunc(funcname string) {
job := newJob(common.REQ, common.CANT_DO, []byte(funcname))
worker.broadcast(job)
}
// Main loop
func (worker *Worker) Work() {
defer func() {
worker.running = false
2012-05-28 10:34:16 +08:00
for _, v := range worker.agents {
v.Close()
}
}()
for funcname, f := range worker.funcs {
worker.addFunc(funcname, f.timeout)
}
worker.running = true
for _, v := range worker.agents {
go v.Work()
}
ok := true
2012-05-28 10:34:16 +08:00
var job *Job
for ok {
2012-05-28 10:34:16 +08:00
if job, ok = <-worker.in; ok {
switch job.DataType {
case common.ERROR:
go func() {
_, err := common.GetError(job.Data)
worker.err(err)
}()
case common.JOB_ASSIGN, common.JOB_ASSIGN_UNIQ:
go func() {
if err := worker.exec(job); err != nil {
worker.err(err)
}
}()
default:
go worker.handleJob(job)
}
}
}
}
// job handler
func (worker *Worker) handleJob(job *Job) {
if worker.JobHandler != nil {
if err := worker.JobHandler(job); err != nil {
worker.err(err)
}
}
}
// Close.
func (worker *Worker) Close() {
close(worker.in)
if worker.limit != nil {
close(worker.limit)
}
}
// Send a something out, get the samething back.
func (worker *Worker) Echo(data []byte) {
job := newJob(common.REQ, common.ECHO_REQ, data)
worker.broadcast(job)
}
// Remove all of functions.
// Both from the worker or job servers.
func (worker *Worker) Reset() {
job := newJob(common.REQ, common.RESET_ABILITIES, nil)
worker.broadcast(job)
worker.funcs = make(JobFuncs)
}
// Set the worker's unique id.
func (worker *Worker) SetId(id string) {
worker.Id = id
job := newJob(common.REQ, common.SET_CLIENT_ID, []byte(id))
worker.broadcast(job)
}
// Execute the job. And send back the result.
func (worker *Worker) exec(job *Job) (err error) {
if worker.limit != nil {
<-worker.limit
defer func() {
worker.limit <- true
}()
}
var limit int
if job.DataType == common.JOB_ASSIGN {
limit = 3
} else {
limit = 4
}
jobdata := bytes.SplitN(job.Data, []byte{'\x00'}, limit)
job.Handle = string(jobdata[0])
funcname := string(jobdata[1])
if job.DataType == common.JOB_ASSIGN {
job.Data = jobdata[2]
} else {
job.UniqueId = string(jobdata[2])
job.Data = jobdata[3]
}
f, ok := worker.funcs[funcname]
if !ok {
return common.Errorf("The function does not exist: %s", funcname)
}
result, err := f.f(job)
var datatype uint32
if err == nil {
datatype = common.WORK_COMPLETE
} else {
if result == nil {
datatype = common.WORK_FAIL
} else {
datatype = common.WORK_EXCEPTION
}
}
job.magicCode = common.REQ
job.DataType = datatype
job.Data = result
job.agent.WriteJob(job)
return
}
func (worker *Worker) removeAgent(a *agent) {
for k, v := range worker.agents {
if v == a {
worker.agents = append(worker.agents[:k], worker.agents[k + 1:] ...)
}
}
if len(worker.agents) == 0 {
2012-05-28 10:34:16 +08:00
worker.err(common.ErrNoActiveAgent)
}
}