Properly synchronizing access to the headers, to prevent concurrent map iteration and map write error (#50)

Author: VitalyChait

presto/client.go

@@ -11,6 +11,7 @@ import (
 	"pbench/log"
 	"reflect"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -42,6 +43,7 @@ type Client struct {
 	baseHeader    http.Header
 	isTrino       bool
 	forceHttps    bool
+	headerMutex   sync.RWMutex
 }
 
 func NewClient(serverUrl string, isTrino bool) (*Client, error) {
@@ -105,20 +107,26 @@ func (c *Client) GetHost() string {
 }
 
 func (c *Client) setHeader(key, value string) {
+	c.headerMutex.Lock()
+	defer c.headerMutex.Unlock()
 	if c.isTrino {
 		key = strings.Replace(key, "X-Presto", "X-Trino", 1)
 	}
 	c.baseHeader.Set(key, value)
 }
 
 func (c *Client) delHeader(key string) {
+	c.headerMutex.Lock()
+	defer c.headerMutex.Unlock()
 	if c.isTrino {
 		key = strings.Replace(key, "X-Presto", "X-Trino", 1)
 	}
 	c.baseHeader.Del(key)
 }
 
 func (c *Client) getHeader(key string) string {
+	c.headerMutex.RLock()
+	defer c.headerMutex.RUnlock()
 	if c.isTrino {
 		key = strings.Replace(key, "X-Presto", "X-Trino", 1)
 	}
@@ -280,7 +288,10 @@ func (c *Client) NewRequest(method, urlStr string, body interface{}, opts ...Req
 		return nil, err
 	}
 
+	c.headerMutex.RLock()
 	req.Header = c.baseHeader.Clone()
+	c.headerMutex.RUnlock()
+
 	if password, ok := c.userInfo.Password(); ok {
 		req.SetBasicAuth(c.userInfo.Username(), password)
 	}