added counting running jobs, disable/enable worker to recive jobs

.gitignore

@@ -20,3 +20,4 @@ _cgo_export.*
 _testmain.go
 
 *.exe
+.idea/

worker/agent.go

@@ -55,6 +55,7 @@ func (a *agent) work() {
 	var err error
 	var data, leftdata []byte
 	for {
+		if !a.worker.IsDisabled() {
 			if data, err = a.read(); err != nil {
 				if opErr, ok := err.(*net.OpError); ok {
 					if opErr.Temporary() {
@@ -109,6 +110,7 @@ func (a *agent) work() {
 			}
 		}
 	}
+}
 
 func (a *agent) disconnect_error(err error) {
 	if a.conn != nil {

worker/common.go

@@ -3,7 +3,7 @@ package worker
 const (
 	Network = "tcp"
 	// queue size
-	queueSize = 8
+	queueSize int = 8
 	// read buffer size
 	bufferSize = 1024
 	// min packet length

worker/worker.go

@@ -23,11 +23,15 @@ type Worker struct {
 	in       chan *inPack
 	running  bool
 	ready    bool
+	disabled bool
+	once     sync.Once
 
 	Id           string
 	ErrorHandler ErrorHandler
 	JobHandler   JobHandler
 	limit        chan bool
+
+	runningJobs int
 }
 
 // Return a worker.
@@ -35,13 +39,14 @@ type Worker struct {
 // 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
+// jobs are limited under the number. If limit is assigned 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),
+		runningJobs: 0,
 	}
 	if limit != Unlimited {
 		worker.limit = make(chan bool, limit-1)
@@ -58,7 +63,7 @@ func (worker *Worker) err(e error) {
 
 // Add a Gearman job server.
 //
-// addr should be formated as 'host:port'.
+// addr should be formatted 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)
@@ -159,7 +164,7 @@ func (worker *Worker) handleInPack(inpack *inPack) {
 	case dtEchoRes:
 		fallthrough
 	default:
-		worker.customeHandler(inpack)
+		worker.customHandler(inpack)
 	}
 }
 
@@ -177,9 +182,13 @@ func (worker *Worker) Ready() (err error) {
 			return
 		}
 	}
+
+	// `once` protects registering worker functions multiple times.
+	worker.once.Do(func() {
 		for funcname, f := range worker.funcs {
 			worker.addFunc(funcname, f.timeout)
 		}
+	})
 	worker.ready = true
 	return
 }
@@ -205,8 +214,8 @@ func (worker *Worker) Work() {
 	}
 }
 
-// custome handling warper
-func (worker *Worker) customeHandler(inpack *inPack) {
+// custom handling wrapper
+func (worker *Worker) customHandler(inpack *inPack) {
 	if worker.JobHandler != nil {
 		if err := worker.JobHandler(inpack); err != nil {
 			worker.err(err)
@@ -227,6 +236,19 @@ func (worker *Worker) Close() {
 	}
 }
 
+func (worker *Worker) Reconnect() error {
+	worker.Lock()
+	defer worker.Unlock()
+	if worker.running == true {
+		for _, a := range worker.agents {
+			if err := a.reconnect(); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
 // Echo
 func (worker *Worker) Echo(data []byte) {
 	outpack := getOutPack()
@@ -266,11 +288,19 @@ func (worker *Worker) exec(inpack *inPack) (err error) {
 				err = ErrUnknown
 			}
 		}
+		if worker.runningJobs > 0 {
+			worker.Lock()
+			worker.runningJobs--
+			worker.Unlock()
+		}
 	}()
 	f, ok := worker.funcs[inpack.fn]
 	if !ok {
 		return fmt.Errorf("The function does not exist: %s", inpack.fn)
 	}
+	worker.Lock()
+	worker.runningJobs++
+	worker.Unlock()
 	var r *result
 	if f.timeout == 0 {
 		d, e := f.f(inpack)
@@ -306,6 +336,37 @@ func (worker *Worker) reRegisterFuncsForAgent(a *agent) {
 
 }
 
+// Counts running jobs
+func (worker *Worker) Count() int {
+	worker.Lock()
+	defer worker.Unlock()
+
+	return worker.runningJobs
+}
+
+// Stops accepting new jobs
+func (worker *Worker) Disable() {
+	worker.Lock()
+	defer worker.Unlock()
+
+	worker.disabled = true
+}
+
+// Renewable disabled workers
+func (worker *Worker) Enable() {
+	worker.Lock()
+	defer worker.Unlock()
+
+	worker.disabled = false
+}
+
+func (worker *Worker) IsDisabled() bool {
+	worker.Lock()
+	defer worker.Unlock()
+
+	return worker.disabled
+}
+
 // inner result
 type result struct {
 	data []byte

worker/worker_disconnect_test.go

@@ -15,7 +15,7 @@ var gearman_ready chan bool
 var kill_gearman chan bool
 var bye chan bool
 
-func init() {
+func init2() {
 
 	if check_gearman_present() {
 		panic(`Something already listening on our testing port. Chickening out of testing with it!`)

worker/worker_test.go

@@ -5,6 +5,8 @@ import (
 	"sync"
 	"testing"
 	"time"
+	"github.com/mikespook/gearman-go/client"
+	"fmt"
 )
 
 var worker *Worker
@@ -223,3 +225,98 @@ func TestWorkWithoutReadyWithPanic(t *testing.T) {
 	}
 
 }
+
+func TestDisableWorkersAndCountRunningJobs(t *testing.T) {
+	worker := New(Unlimited)
+	defer worker.Close()
+
+	if err := worker.AddServer(Network, "127.0.0.1:4730"); err != nil {
+		t.Error(err)
+	}
+	worker.Ready()
+
+	var wg sync.WaitGroup
+	handler := func(job Job) ([]byte, error) {
+		fmt.Println("running job")
+		time.Sleep(time.Second*20)
+		fmt.Println("done")
+		wg.Done()
+		return nil, nil
+	}
+
+	if err := worker.AddFunc("handler", handler, 0); err != nil {
+		wg.Done()
+		t.Error(err)
+	}
+	//worker.JobHandler = handler
+
+	worker.ErrorHandler = func(err error) {
+		t.Fatal("shouldn't have received an error")
+	}
+
+	if err := worker.Ready(); err != nil {
+		t.Error(err)
+		return
+	}
+	go worker.Work()
+
+	var cli *client.Client
+	var err error
+	if cli, err = client.New(client.Network, "127.0.0.1:4730"); err != nil {
+		t.Fatal(err)
+	}
+	cli.ErrorHandler = func(e error) {
+		t.Error(e)
+	}
+
+	worker.Disable()
+	if worker.IsDisabled() {
+		wg.Add(1)
+		_, err = cli.Do("handler", bytes.Repeat([]byte("a"), 50), client.JobHigh, func(res *client.Response) {
+		})
+		if err != nil {
+			t.Error(err)
+		}
+
+		wg.Add(1)
+		_, err = cli.Do("handler", bytes.Repeat([]byte("a"), 50), client.JobHigh, func(res *client.Response) {
+		})
+		if err != nil {
+			t.Error(err)
+		}
+
+		go func () {
+			for {
+				time.Sleep(time.Second*10)
+				if worker.Count() > 0 {
+					fmt.Println("worker enabled", worker.Count())
+					break;
+				} else {
+					fmt.Println("worker do not have any jobs")
+				}
+			}
+		}()
+
+		time.Sleep(time.Second*50)
+		wg.Add(1)
+		_, err = cli.Do("handler", bytes.Repeat([]byte("a"), 50), client.JobHigh, func(res *client.Response) {
+		})
+		if err != nil {
+			t.Error(err)
+		}
+
+		worker.Enable()
+		if !worker.IsDisabled() {
+			fmt.Println("worker is enabled")
+			time.Sleep(time.Second)
+			for i := 1; i < 10; i++ {
+				fmt.Println("Running Job", worker.Count())
+			}
+		} else {
+			t.Fatal("worker should enabled now")
+		}
+		wg.Wait()
+	} else {
+		t.Fatal("worker should disabled")
+	}
+}