From 4dd3cbb4ab52a73d9046b721b5b5bdffa9e10922 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Tue, 25 Jan 2022 10:16:10 +0000
Subject: [PATCH] API client: make http reads more efficient (#976)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Replace `io.ReadAll` with `bytes.Buffer.ReadFrom`.
Both need to resize a buffer until they have finished reading;
the former increases by 1.25x each time while the latter uses 2x.

Also added a benchmark to demonstrate the benefit:
name             old time/op    new time/op    delta
Client/4KB-8       35.9µs ± 4%    35.3µs ± 3%     ~     (p=0.310 n=5+5)
Client/50KB-8      83.1µs ± 8%    69.5µs ± 1%  -16.37%  (p=0.008 n=5+5)
Client/1000KB-8     891µs ± 6%     750µs ± 0%  -15.83%  (p=0.016 n=5+4)
Client/2000KB-8    1.74ms ± 2%    1.35ms ± 1%  -22.72%  (p=0.008 n=5+5)

name             old alloc/op   new alloc/op   delta
Client/4KB-8       20.2kB ± 0%    20.4kB ± 0%   +1.26%  (p=0.008 n=5+5)
Client/50KB-8       218kB ± 0%     136kB ± 0%  -37.65%  (p=0.008 n=5+5)
Client/1000KB-8    5.88MB ± 0%    2.11MB ± 0%  -64.10%  (p=0.008 n=5+5)
Client/2000KB-8    11.7MB ± 0%     4.2MB ± 0%  -63.93%  (p=0.008 n=5+5)

name             old allocs/op  new allocs/op  delta
Client/4KB-8         75.0 ± 0%      72.0 ± 0%   -4.00%  (p=0.008 n=5+5)
Client/50KB-8         109 ± 0%        98 ± 0%  -10.09%  (p=0.079 n=4+5)
Client/1000KB-8       617 ± 0%       593 ± 0%   -3.89%  (p=0.008 n=5+5)
Client/2000KB-8     1.13k ± 0%     1.09k ± 0%   -3.27%  (p=0.008 n=5+5)

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>
---
 api/client.go      |  6 ++++--
 api/client_test.go | 51 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/api/client.go b/api/client.go
index f7ca60b67..1413f65fe 100644
--- a/api/client.go
+++ b/api/client.go
@@ -15,8 +15,8 @@
 package api
 
 import (
+	"bytes"
 	"context"
-	"io/ioutil"
 	"net"
 	"net/http"
 	"net/url"
@@ -111,7 +111,9 @@ func (c *httpClient) Do(ctx context.Context, req *http.Request) (*http.Response,
 	var body []byte
 	done := make(chan struct{})
 	go func() {
-		body, err = ioutil.ReadAll(resp.Body)
+		var buf bytes.Buffer
+		_, err = buf.ReadFrom(resp.Body)
+		body = buf.Bytes()
 		close(done)
 	}()
 
diff --git a/api/client_test.go b/api/client_test.go
index 47094fccd..4215c73db 100644
--- a/api/client_test.go
+++ b/api/client_test.go
@@ -14,7 +14,11 @@
 package api
 
 import (
+	"bytes"
+	"context"
+	"fmt"
 	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"testing"
 )
@@ -111,3 +115,50 @@ func TestClientURL(t *testing.T) {
 		}
 	}
 }
+
+// Serve any http request with a response of N KB of spaces.
+type serveSpaces struct {
+	sizeKB int
+}
+
+func (t serveSpaces) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	kb := bytes.Repeat([]byte{' '}, 1024)
+	for i := 0; i < t.sizeKB; i++ {
+		w.Write(kb)
+	}
+}
+
+func BenchmarkClient(b *testing.B) {
+	b.ReportAllocs()
+	ctx := context.Background()
+
+	for _, sizeKB := range []int{4, 50, 1000, 2000} {
+		b.Run(fmt.Sprintf("%dKB", sizeKB), func(b *testing.B) {
+
+			testServer := httptest.NewServer(serveSpaces{sizeKB})
+			defer testServer.Close()
+
+			client, err := NewClient(Config{
+				Address: testServer.URL,
+			})
+			if err != nil {
+				b.Fatalf("Failed to initialize client: %v", err)
+			}
+			url, err := url.Parse(testServer.URL + "/prometheus/api/v1/query?query=up")
+			if err != nil {
+				b.Fatalf("Failed to parse url: %v", err)
+			}
+			req := &http.Request{
+				URL: url,
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				_, _, err := client.Do(ctx, req)
+				if err != nil {
+					b.Fatalf("Query failed: %v", err)
+				}
+			}
+			b.StopTimer()
+		})
+	}
+}