From 2dbf199260063e9ec79a9cda6cbba477d6d2f71f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoffer=20Fjellstr=C3=B6m?= 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 }