From 2dbf199260063e9ec79a9cda6cbba477d6d2f71f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christoffer=20Fjellstr=C3=B6m?= <c.fjellstrom@gmail.com>
Date: Wed, 27 Sep 2017 11:57:57 +0200
Subject: [PATCH] Add lock on job handler assignment

Fixes race condition on jobs being done before handler is set.
---
 client/client.go | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/client/client.go b/client/client.go
index b243802..ca52b13 100644
--- a/client/client.go
+++ b/client/client.go
@@ -31,7 +31,7 @@ type Client struct {
 }
 
 type responseHandlerMap struct {
-	sync.RWMutex
+	sync.Mutex
 	holder map[string]ResponseHandler
 }
 
@@ -46,17 +46,22 @@ func (r *responseHandlerMap) remove(key string) {
 }
 
 func (r *responseHandlerMap) get(key string) (ResponseHandler, bool) {
-	r.RLock()
+	r.Lock()
 	rh, b := r.holder[key]
-	r.RUnlock()
+	r.Unlock()
 	return rh, b
 }
+
 func (r *responseHandlerMap) put(key string, rh ResponseHandler) {
 	r.Lock()
 	r.holder[key] = rh
 	r.Unlock()
 }
 
+func (r *responseHandlerMap) putNoLock(key string, rh ResponseHandler) {
+	r.holder[key] = rh
+}
+
 // Return a client.
 func New(network, addr string) (client *Client, err error) {
 	client = &Client{
@@ -266,9 +271,12 @@ func (client *Client) Do(funcname string, data []byte,
 	default:
 		datatype = dtSubmitJob
 	}
+
+	client.respHandler.Lock()
+	defer client.respHandler.Unlock()
 	handle, err = client.do(funcname, data, datatype)
 	if err == nil && h != nil {
-		client.respHandler.put(handle, h)
+		client.respHandler.putNoLock(handle, h)
 	}
 	return
 }