From dcbc1ada9029aecdebe3bd4fd158a525460fcb20 Mon Sep 17 00:00:00 2001 From: David Veszelovszki Date: Mon, 16 Dec 2024 10:28:30 +0100 Subject: [PATCH] Add --use-special-header to gateway benchmark --- cmd/src/gateway_benchmark.go | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/cmd/src/gateway_benchmark.go b/cmd/src/gateway_benchmark.go index 0c1c78570d..3420225da5 100644 --- a/cmd/src/gateway_benchmark.go +++ b/cmd/src/gateway_benchmark.go @@ -40,6 +40,7 @@ Examples: $ src gateway benchmark --requests 50 --sgp $ src gateway benchmark --gateway http://localhost:9992 --sourcegraph http://localhost:3082 --sgp $ src gateway benchmark --requests 50 --csv results.csv --sgp + $ src gateway benchmark --gateway https://cody-gateway.sourcegraph.com --sourcegraph https://sourcegraph.com --sgp --use-special-header ` flagSet := flag.NewFlagSet("benchmark", flag.ExitOnError) @@ -50,6 +51,7 @@ Examples: gatewayEndpoint = flagSet.String("gateway", "https://cody-gateway.sourcegraph.com", "Cody Gateway endpoint") sgEndpoint = flagSet.String("sourcegraph", "https://sourcegraph.com", "Sourcegraph endpoint") sgpToken = flagSet.String("sgp", "", "Sourcegraph personal access token for the called instance") + useSpecialHeader = flagSet.Bool("use-special-header", false, "Use special header to test the gateway") ) handler := func(args []string) error { @@ -61,6 +63,10 @@ Examples: return cmderrors.Usage("additional arguments not allowed") } + if *useSpecialHeader { + fmt.Println("Using special header 'cody-core-gc-test'") + } + var ( httpClient = &http.Client{} endpoints = map[string]any{} // Values: URL `string`s or `*webSocketClient`s @@ -80,12 +86,17 @@ Examples: return cmderrors.Usage("must specify --sgp ") } fmt.Println("Benchmarking Sourcegraph instance:", *sgEndpoint) + headers := http.Header{ + "Authorization": []string{"token " + *sgpToken}, + } + if *useSpecialHeader { + headers.Set("cody-core-gc-test", "M2R{+6VI?1,M3n&") + } + endpoints["ws(s): sourcegraph"] = &webSocketClient{ - conn: nil, - URL: strings.Replace(fmt.Sprint(*sgEndpoint, "/.api/gateway/websocket"), "http", "ws", 1), - headers: http.Header{ - "Authorization": []string{"token " + *sgpToken}, - }, + conn: nil, + URL: strings.Replace(fmt.Sprint(*sgEndpoint, "/.api/gateway/websocket"), "http", "ws", 1), + headers: headers, } endpoints["http(s): sourcegraph"] = fmt.Sprint(*sgEndpoint, "/.api/gateway/http") endpoints["http(s): http-then-ws"] = fmt.Sprint(*sgEndpoint, "/.api/gateway/http-then-websocket") @@ -107,7 +118,7 @@ Examples: durations = append(durations, duration) } } else if url, ok := clientOrURL.(string); ok { - duration := benchmarkEndpointHTTP(httpClient, url, *sgpToken) + duration := benchmarkEndpointHTTP(httpClient, url, *sgpToken, *useSpecialHeader) if duration > 0 { durations = append(durations, duration) } @@ -190,7 +201,7 @@ type endpointResult struct { successful int } -func benchmarkEndpointHTTP(client *http.Client, url, accessToken string) time.Duration { +func benchmarkEndpointHTTP(client *http.Client, url, accessToken string, useSpecialHeader bool) time.Duration { start := time.Now() req, err := http.NewRequest("POST", url, strings.NewReader("ping")) if err != nil { @@ -199,6 +210,9 @@ func benchmarkEndpointHTTP(client *http.Client, url, accessToken string) time.Du } req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "token "+accessToken) + if useSpecialHeader { + req.Header.Set("cody-core-gc-test", "M2R{+6VI?1,M3n&") + } resp, err := client.Do(req) if err != nil { fmt.Printf("Error calling %s: %v\n", url, err)