Reuse headers map

4 files changed, 82 insertions, 40 deletions

diff --git a/client/go/internal/mock/http.go b/client/go/internal/mock/http.go
index 8a17d9996d6..f176870a940 100644
--- a/client/go/internal/mock/http.go
+++ b/client/go/internal/mock/http.go
@@ -61,6 +61,8 @@ func (c *HTTPClient) Do(request *http.Request, timeout time.Duration) (*http.Res
 			return nil, err
 		}
 		c.LastBody = body
+	} else {
+		c.LastBody = nil
 	}
 	c.Requests = append(c.Requests, request)
 	if response.Header == nil {
diff --git a/client/go/internal/util/http.go b/client/go/internal/util/http.go
index 546098a204d..35e35b16720 100644
--- a/client/go/internal/util/http.go
+++ b/client/go/internal/util/http.go
@@ -19,17 +19,20 @@ type HTTPClient interface {
 }
 
 type defaultHTTPClient struct {
-	client *http.Client
+	client       *http.Client
+	setUserAgent bool
 }
 
 func (c *defaultHTTPClient) Do(request *http.Request, timeout time.Duration) (response *http.Response, error error) {
 	if c.client.Timeout != timeout { // Set wanted timeout
 		c.client.Timeout = timeout
 	}
-	if request.Header == nil {
-		request.Header = make(http.Header)
+	if c.setUserAgent {
+		if request.Header == nil {
+			request.Header = make(http.Header)
+		}
+		request.Header.Set("User-Agent", fmt.Sprintf("Vespa CLI/%s", build.Version))
 	}
-	request.Header.Set("User-Agent", fmt.Sprintf("Vespa CLI/%s", build.Version))
 	return c.client.Do(request)
 }
 
@@ -65,6 +68,7 @@ func ForceHTTP2(client HTTPClient, certificates []tls.Certificate, caCertificate
 	if !ok {
 		return
 	}
+	c.setUserAgent = false // Let caller control all request headers
 	var dialFunc func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error)
 	if certificates == nil {
 		// No certificate, so force H2C (HTTP/2 over clear-text) by using a non-TLS Dialer
@@ -86,8 +90,11 @@ func ForceHTTP2(client HTTPClient, certificates []tls.Certificate, caCertificate
 }
 
 func CreateClient(timeout time.Duration) HTTPClient {
-	return &defaultHTTPClient{client: &http.Client{
-		Timeout:   timeout,
-		Transport: http.DefaultTransport,
-	}}
+	return &defaultHTTPClient{
+		client: &http.Client{
+			Timeout:   timeout,
+			Transport: http.DefaultTransport,
+		},
+		setUserAgent: true,
+	}
 }
diff --git a/client/go/internal/vespa/document/http.go b/client/go/internal/vespa/document/http.go
index c1ecf38940a..699229f3e73 100644
--- a/client/go/internal/vespa/document/http.go
+++ b/client/go/internal/vespa/document/http.go
@@ -17,6 +17,7 @@ import (
 
 	"github.com/klauspost/compress/gzip"
 
+	"github.com/vespa-engine/vespa/client/go/internal/build"
 	"github.com/vespa-engine/vespa/client/go/internal/util"
 )
 
@@ -31,6 +32,16 @@ const (
 var (
 	fieldsPrefix = []byte(`{"fields":`)
 	fieldsSuffix = []byte("}")
+
+	defaultHeaders http.Header = map[string][]string{
+		"User-Agent":   {fmt.Sprintf("Vespa CLI/%s", build.Version)},
+		"Content-Type": {"application/json; charset=utf-8"},
+	}
+	gzipHeaders http.Header = map[string][]string{
+		"User-Agent":       {fmt.Sprintf("Vespa CLI/%s", build.Version)},
+		"Content-Type":     {"application/json; charset=utf-8"},
+		"Content-Encoding": {"gzip"},
+	}
 )
 
 // Client represents a HTTP client for the /document/v1/ API.
@@ -230,14 +241,26 @@ func (c *Client) prepare(document Document) (*http.Request, *bytes.Buffer, error
 	return pd.request, pd.buf, pd.err
 }
 
+func newRequest(method, url string, body io.Reader, gzipped bool) (*http.Request, error) {
+	req, err := http.NewRequest(method, url, body)
+	if err != nil {
+		return nil, err
+	}
+	if gzipped {
+		req.Header = gzipHeaders
+	} else {
+		req.Header = defaultHeaders
+	}
+	return req, nil
+}
+
 func (c *Client) createRequest(method, url string, body []byte, buf *bytes.Buffer) (*http.Request, error) {
+	buf.Reset()
 	if len(body) == 0 {
-		req, err := http.NewRequest(method, url, nil)
-		return req, err
+		return newRequest(method, url, nil, false)
 	}
 	bodySize := len(fieldsPrefix) + len(body) + len(fieldsSuffix)
 	useGzip := c.options.Compression == CompressionGzip || (c.options.Compression == CompressionAuto && bodySize > 512)
-	buf.Reset()
 	buf.Grow(min(1024, bodySize))
 	if useGzip {
 		zw := c.gzipWriter(buf)
@@ -253,15 +276,7 @@ func (c *Client) createRequest(method, url string, body []byte, buf *bytes.Buffe
 			return nil, err
 		}
 	}
-	req, err := http.NewRequest(method, url, buf)
-	if err != nil {
-		return nil, err
-	}
-	if useGzip {
-		req.Header.Set("Content-Encoding", "gzip")
-	}
-	req.Header.Set("Content-Type", "application/json; charset=utf-8")
-	return req, nil
+	return newRequest(method, url, buf, useGzip)
 }
 
 func (c *Client) clientTimeout() time.Duration {
diff --git a/client/go/internal/vespa/document/http_test.go b/client/go/internal/vespa/document/http_test.go
index fae43ce24e5..7d636aa8d5c 100644
--- a/client/go/internal/vespa/document/http_test.go
+++ b/client/go/internal/vespa/document/http_test.go
@@ -57,10 +57,23 @@ func assertLeastBusy(t *testing.T, id int, client *Client) {
 }
 
 func TestClientSend(t *testing.T) {
-	docs := []Document{
-		{Create: true, Id: mustParseId("id:ns:type::doc1"), Operation: OperationUpdate, Fields: []byte(`{"foo": "123"}`)},
-		{Create: true, Id: mustParseId("id:ns:type::doc2"), Operation: OperationUpdate, Fields: []byte(`{"foo": "456"}`)},
-		{Create: true, Id: mustParseId("id:ns:type::doc3"), Operation: OperationUpdate, Fields: []byte(`{"baz": "789"}`)},
+	var tests = []struct {
+		in     Document
+		method string
+		url    string
+	}{
+		{Document{Create: true, Id: mustParseId("id:ns:type::doc1"), Operation: OperationUpdate, Fields: []byte(`{"foo": "123"}`)},
+			"PUT",
+			"https://example.com:1337/document/v1/ns/type/docid/doc1?timeout=5000ms&create=true"},
+		{Document{Id: mustParseId("id:ns:type::doc2"), Operation: OperationUpdate, Fields: []byte(`{"foo": "456"}`)},
+			"PUT",
+			"https://example.com:1337/document/v1/ns/type/docid/doc2?timeout=5000ms"},
+		{Document{Id: mustParseId("id:ns:type::doc3"), Operation: OperationRemove},
+			"DELETE",
+			"https://example.com:1337/document/v1/ns/type/docid/doc3?timeout=5000ms"},
+		{Document{Condition: "foo", Id: mustParseId("id:ns:type::doc4"), Operation: OperationUpdate, Fields: []byte(`{"baz": "789"}`)},
+			"PUT",
+			"https://example.com:1337/document/v1/ns/type/docid/doc4?timeout=5000ms&condition=foo"},
 	}
 	httpClient := mock.HTTPClient{ReadBody: true}
 	client, _ := NewClient(ClientOptions{
@@ -70,7 +83,8 @@ func TestClientSend(t *testing.T) {
 	clock := manualClock{t: time.Now(), tick: time.Second}
 	client.now = clock.now
 	var stats Stats
-	for i, doc := range docs {
+	for i, tt := range tests {
+		doc := tt.in
 		wantRes := Result{
 			Id: doc.Id,
 			Stats: Stats{
@@ -81,7 +95,8 @@ func TestClientSend(t *testing.T) {
 				MaxLatency:   time.Second,
 			},
 		}
-		if i < 2 {
+		var wantBody bytes.Buffer
+		if i < 3 {
 			httpClient.NextResponseString(200, `{"message":"All good!"}`)
 			wantRes.Status = StatusSuccess
 			wantRes.HTTPStatus = 200
@@ -97,6 +112,11 @@ func TestClientSend(t *testing.T) {
 			wantRes.Stats.Errors = 1
 			wantRes.Stats.BytesRecv = 36
 		}
+		if tt.method == http.MethodPut {
+			wantBody.WriteString(`{"fields":`)
+			wantBody.Write(doc.Fields)
+			wantBody.WriteString("}")
+		}
 		res := client.Send(doc)
 		wantRes.Stats.BytesSent = int64(len(httpClient.LastBody))
 		if !reflect.DeepEqual(res, wantRes) {
@@ -104,35 +124,33 @@ func TestClientSend(t *testing.T) {
 		}
 		stats.Add(res.Stats)
 		r := httpClient.LastRequest
-		if r.Method != http.MethodPut {
-			t.Errorf("got r.Method = %q, want %q", r.Method, http.MethodPut)
+		if r.Method != tt.method {
+			t.Errorf("got r.Method = %q, want %q", r.Method, tt.method)
 		}
-		wantURL := fmt.Sprintf("https://example.com:1337/document/v1/ns/type/docid/%s?timeout=5000ms&create=true", doc.Id.UserSpecific)
-		if r.URL.String() != wantURL {
-			t.Errorf("got r.URL = %q, want %q", r.URL, wantURL)
+		if !reflect.DeepEqual(r.Header, defaultHeaders) {
+			t.Errorf("got r.Header = %v, want %v", r.Header, defaultHeaders)
+		}
+		if r.URL.String() != tt.url {
+			t.Errorf("got r.URL = %q, want %q", r.URL, tt.url)
 		}
-		var wantBody bytes.Buffer
-		wantBody.WriteString(`{"fields":`)
-		wantBody.Write(doc.Fields)
-		wantBody.WriteString("}")
 		if !bytes.Equal(httpClient.LastBody, wantBody.Bytes()) {
 			t.Errorf("got r.Body = %q, want %q", string(httpClient.LastBody), wantBody.String())
 		}
 	}
 	want := Stats{
-		Requests:  3,
-		Responses: 3,
+		Requests:  4,
+		Responses: 4,
 		ResponsesByCode: map[int]int64{
-			200: 2,
+			200: 3,
 			502: 1,
 		},
 		Errors:       1,
 		Inflight:     0,
-		TotalLatency: 3 * time.Second,
+		TotalLatency: 4 * time.Second,
 		MinLatency:   time.Second,
 		MaxLatency:   time.Second,
 		BytesSent:    75,
-		BytesRecv:    82,
+		BytesRecv:    105,
 	}
 	if !reflect.DeepEqual(want, stats) {
 		t.Errorf("got %+v, want %+v", stats, want)