diff options
author | Martin Polden <mpolden@mpolden.no> | 2023-05-15 15:21:59 +0200 |
---|---|---|
committer | Martin Polden <mpolden@mpolden.no> | 2023-05-16 11:23:04 +0200 |
commit | 8a55207eecd3f7612c64918171bef599a2dbf5c7 (patch) | |
tree | 0111c5603ddf70d2fef3d114482acd65858909a9 /client | |
parent | 11b2506e5c9463fa748441f3cbad85e24272796c (diff) |
Reuse headers map
Diffstat (limited to 'client')
-rw-r--r-- | client/go/internal/mock/http.go | 2 | ||||
-rw-r--r-- | client/go/internal/util/http.go | 23 | ||||
-rw-r--r-- | client/go/internal/vespa/document/http.go | 39 | ||||
-rw-r--r-- | client/go/internal/vespa/document/http_test.go | 58 |
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) |