From 0bae6649688937352e6cad3c5cbe5528558dcba2 Mon Sep 17 00:00:00 2001
From: galih rivanto <galih.rivanto@gmail.com>
Date: Tue, 28 May 2019 12:45:24 +0700
Subject: [PATCH] add synchronization to prevent race condition on worker
 shutdown

---
 worker/agent.go  | 3 +++
 worker/worker.go | 7 +++++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/worker/agent.go b/worker/agent.go
index 18c8f67..9fc5716 100644
--- a/worker/agent.go
+++ b/worker/agent.go
@@ -111,6 +111,9 @@ func (a *agent) work() {
 }
 
 func (a *agent) disconnect_error(err error) {
+	a.Unlock()
+	defer a.Unlock()
+
 	if a.conn != nil {
 		err = &WorkerDisconnectError{
 			err:   err,
diff --git a/worker/worker.go b/worker/worker.go
index 8ee04e6..04f4914 100644
--- a/worker/worker.go
+++ b/worker/worker.go
@@ -4,9 +4,9 @@ package worker
 
 import (
 	"fmt"
+	"strconv"
 	"sync"
 	"time"
-	"strconv"
 )
 
 const (
@@ -186,7 +186,7 @@ func (worker *Worker) Ready() (err error) {
 	return
 }
 
-// Main loop, block here
+// Work start main loop (blocking)
 // Most of time, this should be evaluated in goroutine.
 func (worker *Worker) Work() {
 	if !worker.ready {
@@ -197,7 +197,10 @@ func (worker *Worker) Work() {
 		}
 	}
 
+	worker.Lock()
 	worker.running = true
+	worker.Unlock()
+
 	for _, a := range worker.agents {
 		a.Grab()
 	}