Rewrite request-benchmark to not load responses in memory

util-images/request-benchmark/Makefile

@@ -1,6 +1,6 @@
 PROJECT ?= k8s-testimages
 IMG = gcr.io/$(PROJECT)/perf-tests-util/request-benchmark
-TAG = v0.0.6
+TAG = v0.0.7
 
 all: build push
 

util-images/request-benchmark/main.go

@@ -20,8 +20,18 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"io"
 	"log"
+	"net/http"
+	"net/url"
+	"os"
+	"path/filepath"
 	"strings"
+	"time"
+
+	"k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/clientcmd"
+	"k8s.io/client-go/util/flowcontrol"
 )
 
 const (
@@ -30,7 +40,7 @@ const (
 )
 
 var (
-	inflight  = flag.Int("inflight", 0, "Benchmark inflight (number of parallel requests being made to the apiserver")
+	inflight  = flag.Int("inflight", 1, "Benchmark inflight (number of parallel requests being made to the apiserver")
 	namespace = flag.String("namespace", "", "Replace %namespace% in URI with provided namespace")
 	URI       = flag.String("uri", "", "Request URI")
 	verb      = flag.String("verb", "GET", "A verb to be used in requests.")
@@ -42,48 +52,119 @@ func init() {
 }
 
 func main() {
-	client, err := GetClient(float32(*qps))
+	config, err := getConfig()
+	if err != nil {
+		panic(err)
+	}
+	config.QPS = float32(*qps)
+	client, err := rest.HTTPClientFor(config)
 	if err != nil {
 		panic(err)
 	}
+	ctx := context.Background()
 
-	newURI := strings.ReplaceAll(*URI, NamespaceTmpl, *namespace)
-	URI = &newURI
+	serverURL, _, err := rest.DefaultServerUrlFor(config)
+	if err != nil {
+		panic(err)
+	}
+	url, err := url.Parse(strings.ReplaceAll(*URI, NamespaceTmpl, *namespace))
+	if err != nil {
+		panic(err)
+	}
+	url.Host = serverURL.Host
+	url.Scheme = serverURL.Scheme
 
-	if err := healthCheck(client); err != nil {
+	var rateLimiter flowcontrol.RateLimiter
+	if *qps != -1 {
+		rateLimiter = flowcontrol.NewTokenBucketRateLimiter(float32(*qps), 10)
+	}
+
+	if err := healthCheck(ctx, client, *url, rateLimiter); err != nil {
 		panic(err)
 	}
 
-	log.Printf("Sending requests to '%s' with inflight %d. Press Ctrl+C to stop...", *URI, *inflight)
+	log.Printf("Sending requests to '%s' with inflight %d. Press Ctrl+C to stop...", url, *inflight)
 	for i := 0; i < *inflight; i++ {
 		go func() {
 			for {
-				sendRequest(client)
+				sendRequest(ctx, client, *url, rateLimiter)
 			}
 		}()
 	}
 
 	select {} // block main thread from ending
 }
 
-func healthCheck(client *Client) error {
+func healthCheck(ctx context.Context, client *http.Client, url url.URL, rateLimiter flowcontrol.RateLimiter) error {
 	for i := 0; i < HealthCheckRequests; i++ {
-		if sendRequest(client) {
+		if sendRequest(ctx, client, url, rateLimiter) {
 			return nil
 		}
 	}
-	return fmt.Errorf("could not successfully send a request to %s", *URI)
+	return fmt.Errorf("could not successfully send a request to %s", url.String())
 }
 
-func sendRequest(client *Client) bool {
-	request := client.RESTClient().Verb(*verb).RequestURI(*URI)
-	response := client.TimedRequest(context.Background(), request)
-
-	if err := response.Error; err != nil {
-		log.Printf("Got error after sending a request: %v\n%s", err, string(response.Raw))
+func sendRequest(ctx context.Context, client *http.Client, url url.URL, rateLimiter flowcontrol.RateLimiter) bool {
+	req, err := http.NewRequestWithContext(ctx, *verb, url.String(), nil)
+	if err != nil {
+		log.Printf("Got error creating a request: %v\n", err)
 		return false
 	}
-
-	log.Printf("Got response of %d bytes in %v", len(response.Raw), response.Duration)
+	err = tryThrottle(ctx, rateLimiter)
+	if err != nil {
+		log.Printf("Got error throttling a request: %v\n", err)
+		return false
+	}
+	start := time.Now()
+	resp, err := client.Do(req)
+	if err != nil {
+		log.Printf("Got error when sending a request: %v\n", err)
+		return false
+	}
+	if resp.StatusCode < http.StatusOK || resp.StatusCode > http.StatusPartialContent {
+		log.Printf("Got bad status code: %v\n", resp.Status)
+		return false
+	}
+	written, err := io.Copy(io.Discard, resp.Body)
+	if err != nil {
+		log.Printf("Got error when reading response: %v\n", err)
+		return false
+	}
+	log.Printf("Got response of %d bytes in %v", written, time.Since(start))
 	return true
 }
+
+func getConfig() (*rest.Config, error) {
+	if _, ok := os.LookupEnv("KUBERNETES_PORT"); ok {
+		return rest.InClusterConfig()
+	}
+
+	if kubeconfig, ok := os.LookupEnv("KUBECONFIG"); ok {
+		return clientcmd.BuildConfigFromFlags("", kubeconfig)
+	}
+	if home, ok := os.LookupEnv("HOME"); ok {
+		kubeconfig := filepath.Join(home, ".kube", "config")
+		return clientcmd.BuildConfigFromFlags("", kubeconfig)
+	}
+
+	return nil, fmt.Errorf("could not create client-go config")
+}
+
+func tryThrottle(ctx context.Context, rateLimiter flowcontrol.RateLimiter) error {
+	if rateLimiter == nil {
+		return nil
+	}
+
+	now := time.Now()
+	err := rateLimiter.Wait(ctx)
+	if err != nil {
+		err = fmt.Errorf("client rate limiter Wait returned an error: %w", err)
+	}
+	latency := time.Since(now)
+
+	if latency > time.Second {
+		log.Printf("Waited for %v due to client-side throttling, not priority and fairness", latency)
+	}
+
+	return err
+}