Use the new URLs

Author: Emi Gutekanst

cmd/src/gateway_benchmark.go

@@ -11,6 +11,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/gorilla/websocket"
+
 	"github.com/sourcegraph/src-cli/internal/cmderrors"
 )
 
@@ -42,8 +44,9 @@ Examples:
 	flagSet := flag.NewFlagSet("benchmark", flag.ExitOnError)
 
 	var (
-		requestCount = flagSet.Int("requests", 1000, "Number of requests to make per endpoint")
-		csvOutput    = flagSet.String("csv", "", "Export results to CSV file (provide filename)")
+		requestCount    = flagSet.Int("requests", 1000, "Number of requests to make per endpoint")
+		csvOutput       = flagSet.String("csv", "", "Export results to CSV file (provide filename)")
+		gatewayEndpoint = flagSet.String("gateway", "https://cody-gateway.sourcegraph.com", "Cody Gateway endpoint")
 	)
 
 	handler := func(args []string) error {
@@ -55,26 +58,49 @@ Examples:
 			return cmderrors.Usage("additional arguments not allowed")
 		}
 
-		// Create HTTP client with TLS skip verify
-		client := &http.Client{Transport: &http.Transport{}}
-
-		endpoints := map[string]string{
-			"HTTP":                fmt.Sprintf("%s/gateway", cfg.Endpoint),
-			"HTTP then WebSocket": fmt.Sprintf("%s/gateway/http-then-websocket", cfg.Endpoint),
+		var (
+			gatewayWebsocket, sourcegraphWebsocket *websocket.Conn
+			err                                    error
+			httpClient                             = &http.Client{}
+			endpoints                              = map[string]any{} // Values: URL `string`s or `*websocket.Conn`s
+		)
+		if *gatewayEndpoint != "" {
+			wsURL := strings.Replace(fmt.Sprint(*gatewayEndpoint, "/v2/websocket"), "http", "ws", 1)
+			gatewayWebsocket, _, err = websocket.DefaultDialer.Dial(wsURL, nil)
+			if err != nil {
+				return fmt.Errorf("WebSocket dial(%s): %v", wsURL, err)
+			}
+			endpoints["ws(s): gateway"] = gatewayWebsocket
+			endpoints["http(s): gateway"] = fmt.Sprint(*gatewayEndpoint, "/v2/http")
+		}
+		if cfg.Endpoint != "" {
+			wsURL := strings.Replace(fmt.Sprint(cfg.Endpoint, "/.api/gateway/websocket"), "http", "ws", 1)
+			sourcegraphWebsocket, _, err = websocket.DefaultDialer.Dial(wsURL, nil)
+			if err != nil {
+				return fmt.Errorf("WebSocket dial(%s): %v", wsURL, err)
+			}
+			endpoints["ws(s): sourcegraph"] = sourcegraphWebsocket
+			endpoints["http(s): sourcegraph"] = fmt.Sprint(*gatewayEndpoint, "/.api/gateway/http")
 		}
 
 		fmt.Printf("Starting benchmark with %d requests per endpoint...\n", *requestCount)
 
 		var results []endpointResult
-
-		for name, url := range endpoints {
+		for name, clientOrURL := range endpoints {
 			durations := make([]time.Duration, 0, *requestCount)
 			fmt.Printf("\nTesting %s...", name)
 
 			for i := 0; i < *requestCount; i++ {
-				duration := benchmarkEndpoint(client, url)
-				if duration > 0 {
-					durations = append(durations, duration)
+				if ws, ok := clientOrURL.(*websocket.Conn); ok {
+					duration := benchmarkEndpointWebSocket(ws)
+					if duration > 0 {
+						durations = append(durations, duration)
+					}
+				} else if url, ok := clientOrURL.(string); ok {
+					duration := benchmarkEndpointHTTP(httpClient, url)
+					if duration > 0 {
+						durations = append(durations, duration)
+					}
 				}
 			}
 			fmt.Println()
@@ -133,15 +159,15 @@ type endpointResult struct {
 	successful int
 }
 
-func benchmarkEndpoint(client *http.Client, url string) time.Duration {
+func benchmarkEndpointHTTP(client *http.Client, url string) time.Duration {
 	start := time.Now()
 	resp, err := client.Get(url)
 	if err != nil {
 		fmt.Printf("Error calling %s: %v\n", url, err)
 		return 0
 	}
-	defer func(Body io.ReadCloser) {
-		err := Body.Close()
+	defer func(body io.ReadCloser) {
+		err := body.Close()
 		if err != nil {
 			fmt.Printf("Error closing response body: %v\n", err)
 		}
@@ -152,10 +178,42 @@ func benchmarkEndpoint(client *http.Client, url string) time.Duration {
 		fmt.Printf("Error reading response body: %v\n", err)
 		return 0
 	}
+	if resp.StatusCode != http.StatusOK {
+		fmt.Printf("non-200 response: %v\n", resp.Status)
+		return 0
+	}
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		fmt.Printf("Error reading response body: %v\n", err)
+		return 0
+	}
+	if string(body) != "pong" {
+		fmt.Printf("Expected 'pong' response, got: %q\n", string(body))
+		return 0
+	}
 
 	return time.Since(start)
 }
 
+func benchmarkEndpointWebSocket(conn *websocket.Conn) time.Duration {
+	start := time.Now()
+	err := conn.WriteMessage(websocket.TextMessage, []byte("ping"))
+	if err != nil {
+		fmt.Printf("WebSocket write error: %v\n", err)
+		return 0
+	}
+	_, message, err := conn.ReadMessage()
+	if err != nil {
+		fmt.Printf("WebSocket read error: %v\n", err)
+		return 0
+	}
+	if string(message) != "pong" {
+		fmt.Printf("Expected 'pong' response, got: %q\n", string(message))
+		return 0
+	}
+	return time.Since(start)
+}
+
 func calculateStats(durations []time.Duration) Stats {
 	if len(durations) == 0 {
 		return Stats{0, 0, 0, 0, 0, 0, 0}

go.mod

@@ -49,6 +49,7 @@ require (
 require (
 	github.com/google/pprof v0.0.0-20230602150820-91b7bce49751 // indirect
 	github.com/google/s2a-go v0.1.4 // indirect
+	github.com/gorilla/websocket v1.5.3 // indirect
 	github.com/jackc/chunkreader/v2 v2.0.1 // indirect
 	github.com/jackc/pgconn v1.14.3 // indirect
 	github.com/jackc/pgio v1.0.0 // indirect
@@ -196,4 +197,5 @@ require (
 
 // See: https://github.com/ghodss/yaml/pull/65
 replace github.com/ghodss/yaml => github.com/sourcegraph/yaml v1.0.1-0.20200714132230-56936252f152
+
 replace github.com/sourcegraph/sourcegraph/lib => github.com/sourcegraph/sourcegraph-public-snapshot/lib v0.0.0-20240709083501-1af563b61442

go.sum

@@ -211,6 +211,8 @@ github.com/googleapis/gax-go/v2 v2.12.0/go.mod h1:y+aIqrI5eb1YGMVJfuV3185Ts/D7qK
 github.com/gorilla/css v1.0.0 h1:BQqNyPTi50JCFMTw/b67hByjMVXZRwGha6wxVGkeihY=
 github.com/gorilla/css v1.0.0/go.mod h1:Dn721qIggHpt4+EFCcTLTU/vk5ySda2ReITrtgBl60c=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grafana/regexp v0.0.0-20221123153739-15dc172cd2db h1:7aN5cccjIqCLTzedH7MZzRZt5/lsAHch6Z3L2ZGn5FA=
 github.com/grafana/regexp v0.0.0-20221123153739-15dc172cd2db/go.mod h1:M5qHK+eWfAv8VR/265dIuEpL3fNfeC21tXXp9itM24A=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=