Merge pull request #27083 from vespa-engine/mpolden/use-correct-size

Include prefix and suffix in body size

3 files changed, 21 insertions, 11 deletions

diff --git a/client/go/internal/vespa/document/document.go b/client/go/internal/vespa/document/document.go
index 89127d82dce..ce8b22b24f0 100644
--- a/client/go/internal/vespa/document/document.go
+++ b/client/go/internal/vespa/document/document.go
@@ -265,7 +265,7 @@ func NewGenerator(size int, deadline time.Time) *Generator {
 	return &Generator{Size: size, Deadline: deadline, nowFunc: time.Now}
 }
 
-func (g *Generator) randString(size int) string {
+func randString(size int) string {
 	b := make([]byte, size)
 	for i := range b {
 		b[i] = byte('a' + rand.Intn('z'-'a'+1))
@@ -278,7 +278,7 @@ func (g *Generator) Read(p []byte) (int, error) {
 		if !g.nowFunc().Before(g.Deadline) {
 			return 0, io.EOF
 		}
-		fmt.Fprintf(&g.buf, "{\"put\": \"id:test:test::%s\", \"fields\": {\"test\": \"%s\"}}\n", g.randString(8), g.randString(g.Size))
+		fmt.Fprintf(&g.buf, "{\"put\": \"id:test:test::%s\", \"fields\": {\"test\": \"%s\"}}\n", randString(8), randString(g.Size))
 	}
 	return g.buf.Read(p)
 }
diff --git a/client/go/internal/vespa/document/http.go b/client/go/internal/vespa/document/http.go
index 8ca8ca1e93d..319512458c7 100644
--- a/client/go/internal/vespa/document/http.go
+++ b/client/go/internal/vespa/document/http.go
@@ -219,10 +219,11 @@ func (c *Client) createRequest(method, url string, body []byte) (*http.Request,
 		req, err := http.NewRequest(method, url, nil)
 		return req, nil, err
 	}
-	useGzip := c.options.Compression == CompressionGzip || (c.options.Compression == CompressionAuto && len(body) > 512)
+	bodySize := len(fieldsPrefix) + len(body) + len(fieldsSuffix)
+	useGzip := c.options.Compression == CompressionGzip || (c.options.Compression == CompressionAuto && bodySize > 512)
 	pr, pw := io.Pipe()
 	go func() {
-		bw := bufio.NewWriterSize(pw, min(1024, len(fieldsPrefix)+len(body)+len(fieldsSuffix)))
+		bw := bufio.NewWriterSize(pw, min(1024, bodySize))
 		defer func() {
 			bw.Flush()
 			pw.Close()
diff --git a/client/go/internal/vespa/document/http_test.go b/client/go/internal/vespa/document/http_test.go
index 9a47b4f45fe..8fb0b1b2b95 100644
--- a/client/go/internal/vespa/document/http_test.go
+++ b/client/go/internal/vespa/document/http_test.go
@@ -288,13 +288,22 @@ func benchmarkClientSend(b *testing.B, compression Compression, document Documen
 	}
 }
 
-func BenchmarkClientSend(b *testing.B) {
-	doc := Document{Create: true, Id: mustParseId("id:ns:type::doc1"), Operation: OperationUpdate, Fields: []byte(`{"foo": "my document"}`)}
-	benchmarkClientSend(b, CompressionNone, doc)
+func makeDocument(size int) Document {
+	return Document{Id: mustParseId("id:ns:type::doc1"), Operation: OperationUpdate, Fields: []byte(fmt.Sprintf(`{"foo": "%s"}`, randString(size)))}
 }
 
-func BenchmarkClientSendCompressed(b *testing.B) {
-	body := fmt.Sprintf(`{"foo": "%s"}`, strings.Repeat("my document", 100))
-	doc := Document{Create: true, Id: mustParseId("id:ns:type::doc1"), Operation: OperationUpdate, Fields: []byte(body)}
-	benchmarkClientSend(b, CompressionGzip, doc)
+func BenchmarkClientSendSmallUncompressed(b *testing.B) {
+	benchmarkClientSend(b, CompressionNone, makeDocument(10))
+}
+
+func BenchmarkClientSendMediumUncompressed(b *testing.B) {
+	benchmarkClientSend(b, CompressionNone, makeDocument(10))
+}
+
+func BenchmarkClientSendMediumGzip(b *testing.B) {
+	benchmarkClientSend(b, CompressionGzip, makeDocument(1000))
+}
+
+func BenchmarkClientSendLargeGzip(b *testing.B) {
+	benchmarkClientSend(b, CompressionGzip, makeDocument(10000))
 }