// Copyright 2011 Xing Xing 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 ( "bitbucket.org/mikespook/gearman-go/gearman" "bytes" "sync" ) // The definition of the callback function. type JobFunction func(job *WorkerJob) ([]byte, error) // Map for added function. type JobFunctionMap map[string]JobFunction /* Worker side api for gearman. usage: worker = NewWorker() worker.AddFunction("foobar", foobar) worker.AddServer("127.0.0.1:4730") worker.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 *WorkerJob) (data []byte, err os.Error) { //sth. here //plaplapla... return } */ type Worker struct { clients []*jobAgent functions JobFunctionMap running bool incoming chan *WorkerJob mutex sync.Mutex JobQueue chan *WorkerJob ErrQueue chan error } // Get a new worker func New() (worker *Worker) { worker = &Worker{ // job server list clients: make([]*jobAgent, 0, gearman.WORKER_SERVER_CAP), // function list functions: make(JobFunctionMap), incoming: make(chan *WorkerJob, gearman.QUEUE_CAP), JobQueue: make(chan *WorkerJob, gearman.QUEUE_CAP), ErrQueue: make(chan error, gearman.QUEUE_CAP), running: true, } return } // 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) { worker.mutex.Lock() defer worker.mutex.Unlock() if len(worker.clients) == cap(worker.clients) { return gearman.ErrOutOfCap } // Create a new job server's client as a agent of server server, err := newJobAgent(addr, worker) if err != nil { return err } n := len(worker.clients) worker.clients = worker.clients[0 : n+1] worker.clients[n] = server return } // 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) AddFunction(funcname string, f JobFunction, timeout uint32) (err error) { if len(worker.clients) < 1 { return gearman.ErrNotConn } worker.mutex.Lock() defer worker.mutex.Unlock() worker.functions[funcname] = f var datatype uint32 var data []byte if timeout == 0 { datatype = gearman.CAN_DO data = []byte(funcname) } else { datatype = gearman.CAN_DO_TIMEOUT data = []byte(funcname + "\x00") t := gearman.Uint32ToBytes(timeout) data = append(data, t[:]...) } job := NewWorkerJob(gearman.REQ, datatype, data) worker.WriteJob(job) return } // Remove a function. // Tell job servers 'I can not do this now' at the same time. func (worker *Worker) RemoveFunction(funcname string) (err error) { worker.mutex.Lock() defer worker.mutex.Unlock() if worker.functions[funcname] == nil { return gearman.ErrFuncNotFound } delete(worker.functions, funcname) job := NewWorkerJob(gearman.REQ, gearman.CANT_DO, []byte(funcname)) worker.WriteJob(job) return } // Main loop func (worker *Worker) Work() { for _, v := range worker.clients { go v.Work() } for worker.running || len(worker.incoming) > 0{ select { case job := <-worker.incoming: if job == nil { break } switch job.DataType { case gearman.NO_JOB: // do nothing case gearman.ERROR: _, err := gearman.GetError(job.Data) worker.ErrQueue <- err case gearman.JOB_ASSIGN, gearman.JOB_ASSIGN_UNIQ: go func() { if err := worker.exec(job); err != nil { worker.ErrQueue <- err } }() default: worker.JobQueue <- job } } } close(worker.incoming) } // Get the last job in queue. // If there are more than one job in the queue, // the last one will be returned, // the others will be lost. func (worker *Worker) LastJob() (job *WorkerJob) { if l := len(worker.JobQueue); l != 1 { if l == 0 { return } for i := 0; i < l-1; i++ { <-worker.JobQueue } } return <-worker.JobQueue } // Close. func (worker *Worker) Close() (err error) { for _, v := range worker.clients { err = v.Close() } worker.running = false return err } // 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) WriteJob(job *WorkerJob) (err error) { e := make(chan error) for _, v := range worker.clients { go func() { e <- v.WriteJob(job) }() } return <-e } // Send a something out, get the samething back. func (worker *Worker) Echo(data []byte) (err error) { job := NewWorkerJob(gearman.REQ, gearman.ECHO_REQ, data) return worker.WriteJob(job) } // Remove all of functions. // Both from the worker or job servers. func (worker *Worker) Reset() (err error) { job := NewWorkerJob(gearman.REQ, gearman.RESET_ABILITIES, nil) err = worker.WriteJob(job) worker.functions = make(JobFunctionMap) return } // Set the worker's unique id. func (worker *Worker) SetId(id string) (err error) { job := NewWorkerJob(gearman.REQ, gearman.SET_CLIENT_ID, []byte(id)) return worker.WriteJob(job) } // Execute the job. And send back the result. func (worker *Worker) exec(job *WorkerJob) (err error) { var limit int if job.DataType == gearman.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 == gearman.JOB_ASSIGN { job.Data = jobdata[2] } else { job.UniqueId = string(jobdata[2]) job.Data = jobdata[3] } f, ok := worker.functions[funcname] if !ok { return gearman.ErrFuncNotFound } result, err := f(job) var datatype uint32 if err == nil { datatype = gearman.WORK_COMPLETE } else { if result == nil { datatype = gearman.WORK_FAIL } else { datatype = gearman.WORK_EXCEPTION } } job.magicCode = gearman.REQ job.DataType = datatype job.Data = result worker.WriteJob(job) return }