Merge 8cdaef9752 into 320155c608

worker/worker.go

@@ -18,11 +18,15 @@ const (
 // It can connect to multi-server and grab jobs.
 type Worker struct {
 	sync.Mutex
-	agents  []*agent
+	agents       []*agent
-	funcs   jobFuncs
+	funcs        jobFuncs
-	in      chan *inPack
+	in           chan *inPack
-	running bool
+	running      bool
-	ready bool
+	ready        bool
+	shuttingDown bool
+	activeJobs   int
+	// Used during shutdown to wait for all active jobs to finish
+	finishedDraining sync.WaitGroup
 
 	Id           string
 	ErrorHandler ErrorHandler
@@ -43,6 +47,7 @@ func New(limit int) (worker *Worker) {
 		funcs:  make(jobFuncs),
 		in:     make(chan *inPack, queueSize),
 	}
+	worker.finishedDraining.Add(1)
 	if limit != Unlimited {
 		worker.limit = make(chan bool, limit-1)
 	}
@@ -137,7 +142,9 @@ func (worker *Worker) handleInPack(inpack *inPack) {
 	case dtNoJob:
 		inpack.a.PreSleep()
 	case dtNoop:
-		inpack.a.Grab()
+		if !worker.shuttingDown {
+			inpack.a.Grab()
+		}
 	case dtJobAssign, dtJobAssignUniq:
 		go func() {
 			if err := worker.exec(inpack); err != nil {
@@ -147,7 +154,9 @@ func (worker *Worker) handleInPack(inpack *inPack) {
 		if worker.limit != nil {
 			worker.limit <- true
 		}
-		inpack.a.Grab()
+		if !worker.shuttingDown {
+			inpack.a.Grab()
+		}
 	case dtError:
 		worker.err(inpack.Err())
 		fallthrough
@@ -182,11 +191,11 @@ func (worker *Worker) Ready() (err error) {
 // Main loop, block here
 // Most of time, this should be evaluated in goroutine.
 func (worker *Worker) Work() {
-	if ! worker.ready {
+	if !worker.ready {
 		// didn't run Ready beforehand, so we'll have to do it:
 		err := worker.Ready()
 		if err != nil {
-			panic( err )
+			panic(err)
 		}
 	}
 
@@ -224,6 +233,19 @@ func (worker *Worker) Close() {
 	}
 }
 
+// Shutdown server gracefully. This function will block until all active work has finished.
+func (worker *Worker) Shutdown() {
+	worker.Lock()
+	worker.shuttingDown = true
+	if worker.activeJobs == 0 {
+		worker.finishedDraining.Done()
+	}
+	worker.Unlock()
+	// Wait for the mutex
+	worker.finishedDraining.Wait()
+	worker.Close()
+}
+
 // Echo
 func (worker *Worker) Echo(data []byte) {
 	outpack := getOutPack()
@@ -250,6 +272,29 @@ func (worker *Worker) SetId(id string) {
 	worker.broadcast(outpack)
 }
 
+// IncrementActive increments the count of active jobs. This will return false if no more
+// jobs can't be started because the worker is shutting down.
+func (worker *Worker) incrementActive() bool {
+	worker.Lock()
+	defer worker.Unlock()
+	if worker.shuttingDown {
+		return false
+	}
+	worker.activeJobs = worker.activeJobs + 1
+	return true
+}
+
+// DecrementActive decrements the count of active jobs. If the process is shutting down
+// it will set the finishedDraining flag if there are no more active jobs.
+func (worker *Worker) decrementActive() {
+	worker.Lock()
+	defer worker.Unlock()
+	worker.activeJobs = worker.activeJobs - 1
+	if worker.shuttingDown && worker.activeJobs == 0 {
+		worker.finishedDraining.Done()
+	}
+}
+
 // inner job executing
 func (worker *Worker) exec(inpack *inPack) (err error) {
 	defer func() {
@@ -263,7 +308,11 @@ func (worker *Worker) exec(inpack *inPack) (err error) {
 				err = ErrUnknown
 			}
 		}
+		worker.decrementActive()
 	}()
+	if !worker.incrementActive() {
+		return
+	}
 	f, ok := worker.funcs[inpack.fn]
 	if !ok {
 		return fmt.Errorf("The function does not exist: %s", inpack.fn)

worker/worker_test.go

@@ -1,6 +1,7 @@
 package worker
 
 import (
+	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -62,6 +63,7 @@ func TestWorkerRemoveFunc(t *testing.T) {
 }
 
 func TestWork(t *testing.T) {
+	// TODO: Worth looking at this for shutdown (WaitGroup)
 	var wg sync.WaitGroup
 	worker.JobHandler = func(job Job) error {
 		t.Logf("%s", job.Data())
@@ -78,12 +80,11 @@ func TestWork(t *testing.T) {
 	wg.Wait()
 }
 
-
 func TestWorkerClose(t *testing.T) {
 	worker.Close()
 }
 
-func TestWorkWithoutReady(t * testing.T){
+func TestWorkWithoutReady(t *testing.T) {
 	other_worker := New(Unlimited)
 
 	if err := other_worker.AddServer(Network, "127.0.0.1:4730"); err != nil {
@@ -94,13 +95,13 @@ func TestWorkWithoutReady(t * testing.T){
 	}
 
 	timeout := make(chan bool, 1)
-	done := make( chan bool, 1)
+	done := make(chan bool, 1)
 
-	other_worker.JobHandler = func( j Job ) error {
+	other_worker.JobHandler = func(j Job) error {
-		if( ! other_worker.ready ){
+		if !other_worker.ready {
-			t.Error("Worker not ready as expected");
+			t.Error("Worker not ready as expected")
 		}
-		done <-true
+		done <- true
 		return nil
 	}
 	go func() {
@@ -108,15 +109,15 @@ func TestWorkWithoutReady(t * testing.T){
 		timeout <- true
 	}()
 
-	go func(){
+	go func() {
-		other_worker.Work();
+		other_worker.Work()
 	}()
 
 	// With the all-in-one Work() we don't know if the
 	// worker is ready at this stage so we may have to wait a sec:
-	go func(){
+	go func() {
 		tries := 3
-		for( tries > 0 ){
+		for tries > 0 {
 			if other_worker.ready {
 				other_worker.Echo([]byte("Hello"))
 				break
@@ -129,22 +130,22 @@ func TestWorkWithoutReady(t * testing.T){
 	}()
 
 	// determine if we've finished or timed out:
-	select{
+	select {
-	case <- timeout:
+	case <-timeout:
 		t.Error("Test timed out waiting for the worker")
-	case <- done:
+	case <-done:
 	}
 }
 
-func TestWorkWithoutReadyWithPanic(t * testing.T){
+func TestWorkWithoutReadyWithPanic(t *testing.T) {
 	other_worker := New(Unlimited)
 
 	timeout := make(chan bool, 1)
-	done := make( chan bool, 1)
+	done := make(chan bool, 1)
 
 	// Going to work with no worker setup.
 	// when Work (hopefully) calls Ready it will get an error which should cause it to panic()
-	go func(){
+	go func() {
 		defer func() {
 			if err := recover(); err != nil {
 				done <- true
@@ -153,17 +154,122 @@ func TestWorkWithoutReadyWithPanic(t * testing.T){
 			t.Error("Work should raise a panic.")
 			done <- true
 		}()
-		other_worker.Work();
+		other_worker.Work()
 	}()
 	go func() {
 		time.Sleep(2 * time.Second)
 		timeout <- true
 	}()
 
-	select{
+	select {
-	case <- timeout:
+	case <-timeout:
 		t.Error("Test timed out waiting for the worker")
-	case <- done:
+	case <-done:
 	}
 
 }
+
+// initWorker creates a worker and adds the localhost server to it
+func initWorker(t *testing.T) *Worker {
+	otherWorker := New(Unlimited)
+	if err := otherWorker.AddServer(Network, "127.0.0.1:4730"); err != nil {
+		t.Error(err)
+	}
+	return otherWorker
+}
+
+// submitEmptyInPack sends an empty inpack with the specified fn name to the worker. It uses
+// the first agent of the worker.
+func submitEmptyInPack(t *testing.T, worker *Worker, function string) {
+	if l := len(worker.agents); l != 1 {
+		t.Error("The worker has no agents")
+	}
+	inpack := getInPack()
+	inpack.dataType = dtJobAssign
+	inpack.fn = function
+	inpack.a = worker.agents[0]
+	worker.in <- inpack
+}
+
+// TestShutdownSuccessJob tests that shutdown waits for the currently running job to
+// complete.
+func TestShutdownSuccessJob(t *testing.T) {
+	otherWorker := initWorker(t)
+	output := 0
+	var wg sync.WaitGroup
+	successJob := func(job Job) ([]byte, error) {
+		wg.Done()
+		// Sleep for 100ms to ensure that the shutdown waits for this to finish
+		time.Sleep(time.Duration(100 * time.Millisecond))
+		output = 1
+		return nil, nil
+	}
+	if err := otherWorker.AddFunc("test", successJob, 0); err != nil {
+		t.Error(err)
+	}
+	if err := otherWorker.Ready(); err != nil {
+		t.Error(err)
+		return
+	}
+	submitEmptyInPack(t, otherWorker, "test")
+	go otherWorker.Work()
+	// Wait for the success_job to start so that we know we didn't shutdown before even
+	// beginning to process the job.
+	wg.Add(1)
+	wg.Wait()
+	otherWorker.Shutdown()
+	if output != 1 {
+		t.Error("Expected 1, output was: " + strconv.Itoa(output))
+	}
+}
+
+func TestShutdownFailureJob(t *testing.T) {
+	otherWorker := initWorker(t)
+	output := 0
+	var wg sync.WaitGroup
+	failureJob := func(job Job) ([]byte, error) {
+		wg.Done()
+		// Sleep for 100ms to ensure that shutdown waits for this to finish
+		time.Sleep(time.Duration(100 * time.Millisecond))
+		output = 1
+		return nil, nil //new Error()
+	}
+
+	if err := otherWorker.AddFunc("test", failureJob, 0); err != nil {
+		t.Error(err)
+	}
+	if err := otherWorker.Ready(); err != nil {
+		t.Error(err)
+		return
+	}
+	submitEmptyInPack(t, otherWorker, "test")
+	go otherWorker.Work()
+	// Wait for the success_job to start so that we know we didn't shutdown before even
+	// beginning to process the job.
+	wg.Add(1)
+	wg.Wait()
+	otherWorker.Shutdown()
+	if output != 1 {
+		t.Error("Expected 1, output was: " + strconv.Itoa(output))
+	}
+}
+
+func TestSubmitJobAfterShutdown(t *testing.T) {
+	otherWorker := initWorker(t)
+	noRunJob := func(job Job) ([]byte, error) {
+		t.Error("This job shouldn't have been run")
+		return nil, nil
+	}
+	if err := otherWorker.AddFunc("test", noRunJob, 0); err != nil {
+		t.Error(err)
+	}
+	if err := otherWorker.Ready(); err != nil {
+		t.Error(err)
+		return
+	}
+	go otherWorker.Work()
+	otherWorker.Shutdown()
+	submitEmptyInPack(t, otherWorker, "test")
+	// Sleep for 100ms to make sure that the job doesn't actually run
+	time.Sleep(time.Duration(100 * time.Millisecond))
+}