Skip to content

Commit 177b25b

Browse files
authored
feat/gateway: Add src gateway benchmark command (#1124)
1 parent b9d5d30 commit 177b25b

File tree

4 files changed

+378
-1
lines changed

4 files changed

+378
-1
lines changed

cmd/src/colors.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,14 @@ func bg256Color(code int) string {
2222

2323
// See https://i.stack.imgur.com/KTSQa.png or https://jonasjacek.github.io/colors/
2424
var ansiColors = map[string]string{
25-
"nc": "\033[0m",
25+
// Simple colors.
26+
"blue": "\033[34m",
27+
"green": "\033[32m",
28+
"yellow": "\033[33m",
29+
"red": "\033[31m",
30+
"nc": "\033[0m", // reset
31+
32+
// Custom colors.
2633
"logo": fg256Color(57),
2734
"warning": fg256Color(124),
2835
"success": fg256Color(2),

cmd/src/gateway.go

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package main
2+
3+
import (
4+
"flag"
5+
"fmt"
6+
)
7+
8+
var gatewayCommands commander
9+
10+
func init() {
11+
usage := `'src gateway' interacts with Cody Gateway (directly or through a Sourcegraph instance).
12+
13+
Usage:
14+
15+
src gateway command [command options]
16+
17+
The commands are:
18+
19+
benchmark runs benchmarks against Cody Gateway
20+
21+
Use "src gateway [command] -h" for more information about a command.
22+
23+
`
24+
25+
flagSet := flag.NewFlagSet("gateway", flag.ExitOnError)
26+
handler := func(args []string) error {
27+
gatewayCommands.run(flagSet, "src gateway", usage, args)
28+
return nil
29+
}
30+
31+
// Register the command.
32+
commands = append(commands, &command{
33+
flagSet: flagSet,
34+
aliases: []string{}, // No aliases for gateway command
35+
handler: handler,
36+
usageFunc: func() { fmt.Println(usage) },
37+
})
38+
}

cmd/src/gateway_benchmark.go

Lines changed: 331 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,331 @@
1+
package main
2+
3+
import (
4+
"encoding/csv"
5+
"flag"
6+
"fmt"
7+
"io"
8+
"net/http"
9+
"os"
10+
"sort"
11+
"strings"
12+
"time"
13+
14+
"github.com/sourcegraph/src-cli/internal/cmderrors"
15+
)
16+
17+
type Stats struct {
18+
Avg time.Duration
19+
P5 time.Duration
20+
P75 time.Duration
21+
P80 time.Duration
22+
P95 time.Duration
23+
Median time.Duration
24+
Total time.Duration
25+
}
26+
27+
func init() {
28+
usage := `
29+
'src gateway benchmark' runs performance benchmarks against Cody Gateway endpoints.
30+
31+
Usage:
32+
33+
src gateway benchmark [flags]
34+
35+
Examples:
36+
37+
$ src gateway benchmark
38+
$ src gateway benchmark --requests 50
39+
$ src gateway benchmark --requests 50 --csv results.csv
40+
`
41+
42+
flagSet := flag.NewFlagSet("benchmark", flag.ExitOnError)
43+
44+
var (
45+
requestCount = flagSet.Int("requests", 1000, "Number of requests to make per endpoint")
46+
csvOutput = flagSet.String("csv", "", "Export results to CSV file (provide filename)")
47+
)
48+
49+
handler := func(args []string) error {
50+
if err := flagSet.Parse(args); err != nil {
51+
return err
52+
}
53+
54+
if len(flagSet.Args()) != 0 {
55+
return cmderrors.Usage("additional arguments not allowed")
56+
}
57+
58+
// Create HTTP client with TLS skip verify
59+
client := &http.Client{Transport: &http.Transport{}}
60+
61+
endpoints := map[string]string{
62+
"HTTP": fmt.Sprintf("%s/gateway", cfg.Endpoint),
63+
"HTTP then WebSocket": fmt.Sprintf("%s/gateway/http-then-websocket", cfg.Endpoint),
64+
}
65+
66+
fmt.Printf("Starting benchmark with %d requests per endpoint...\n", *requestCount)
67+
68+
var results []endpointResult
69+
70+
for name, url := range endpoints {
71+
durations := make([]time.Duration, 0, *requestCount)
72+
fmt.Printf("\nTesting %s...", name)
73+
74+
for i := 0; i < *requestCount; i++ {
75+
duration := benchmarkEndpoint(client, url)
76+
if duration > 0 {
77+
durations = append(durations, duration)
78+
}
79+
}
80+
fmt.Println()
81+
82+
stats := calculateStats(durations)
83+
84+
results = append(results, endpointResult{
85+
name: name,
86+
avg: stats.Avg,
87+
median: stats.Median,
88+
p5: stats.P5,
89+
p75: stats.P75,
90+
p80: stats.P80,
91+
p95: stats.P95,
92+
total: stats.Total,
93+
successful: len(durations),
94+
})
95+
}
96+
97+
printResults(results, requestCount)
98+
99+
if *csvOutput != "" {
100+
if err := writeResultsToCSV(*csvOutput, results, requestCount); err != nil {
101+
return fmt.Errorf("failed to export CSV: %v", err)
102+
}
103+
fmt.Printf("\nResults exported to %s\n", *csvOutput)
104+
}
105+
106+
return nil
107+
}
108+
109+
gatewayCommands = append(gatewayCommands, &command{
110+
flagSet: flagSet,
111+
aliases: []string{},
112+
handler: handler,
113+
usageFunc: func() {
114+
_, err := fmt.Fprintf(flag.CommandLine.Output(), "Usage of 'src gateway %s':\n", flagSet.Name())
115+
if err != nil {
116+
return
117+
}
118+
flagSet.PrintDefaults()
119+
fmt.Println(usage)
120+
},
121+
})
122+
}
123+
124+
type endpointResult struct {
125+
name string
126+
avg time.Duration
127+
median time.Duration
128+
p5 time.Duration
129+
p75 time.Duration
130+
p80 time.Duration
131+
p95 time.Duration
132+
total time.Duration
133+
successful int
134+
}
135+
136+
func benchmarkEndpoint(client *http.Client, url string) time.Duration {
137+
start := time.Now()
138+
resp, err := client.Get(url)
139+
if err != nil {
140+
fmt.Printf("Error calling %s: %v\n", url, err)
141+
return 0
142+
}
143+
defer func(Body io.ReadCloser) {
144+
err := Body.Close()
145+
if err != nil {
146+
fmt.Printf("Error closing response body: %v\n", err)
147+
}
148+
}(resp.Body)
149+
150+
_, err = io.ReadAll(resp.Body)
151+
if err != nil {
152+
fmt.Printf("Error reading response body: %v\n", err)
153+
return 0
154+
}
155+
156+
return time.Since(start)
157+
}
158+
159+
func calculateStats(durations []time.Duration) Stats {
160+
if len(durations) == 0 {
161+
return Stats{0, 0, 0, 0, 0, 0, 0}
162+
}
163+
164+
// Sort durations in ascending order
165+
sort.Slice(durations, func(i, j int) bool {
166+
return durations[i] < durations[j]
167+
})
168+
169+
var sum time.Duration
170+
for _, d := range durations {
171+
sum += d
172+
}
173+
avg := sum / time.Duration(len(durations))
174+
175+
return Stats{
176+
Avg: avg,
177+
P5: durations[int(float64(len(durations))*0.05)],
178+
P75: durations[int(float64(len(durations))*0.75)],
179+
P80: durations[int(float64(len(durations))*0.80)],
180+
P95: durations[int(float64(len(durations))*0.95)],
181+
Median: durations[(len(durations) / 2)],
182+
Total: sum,
183+
}
184+
}
185+
186+
func formatDuration(d time.Duration, best bool, worst bool) string {
187+
value := fmt.Sprintf("%.2fms", float64(d.Microseconds())/1000)
188+
if best {
189+
return ansiColors["green"] + value + ansiColors["nc"]
190+
}
191+
if worst {
192+
return ansiColors["red"] + value + ansiColors["nc"]
193+
}
194+
return ansiColors["yellow"] + value + ansiColors["nc"]
195+
}
196+
197+
func formatSuccessRate(successful, total int, best bool, worst bool) string {
198+
value := fmt.Sprintf("%d/%d", successful, total)
199+
if best {
200+
return ansiColors["green"] + value + ansiColors["nc"]
201+
}
202+
if worst {
203+
return ansiColors["red"] + value + ansiColors["nc"]
204+
}
205+
return ansiColors["yellow"] + value + ansiColors["nc"]
206+
}
207+
208+
func printResults(results []endpointResult, requestCount *int) {
209+
// Print header
210+
headerFmt := ansiColors["blue"] + "%-20s | %-10s | %-10s | %-10s | %-10s | %-10s | %-10s | %-10s | %-10s" + ansiColors["nc"] + "\n"
211+
fmt.Printf("\n"+headerFmt,
212+
"Endpoint ", "Average", "Median", "P5", "P75", "P80", "P95", "Total", "Success")
213+
fmt.Println(ansiColors["blue"] + strings.Repeat("-", 121) + ansiColors["nc"])
214+
215+
// Find best/worst values for each metric
216+
var bestAvg, worstAvg time.Duration
217+
var bestMedian, worstMedian time.Duration
218+
var bestP5, worstP5 time.Duration
219+
var bestP75, worstP75 time.Duration
220+
var bestP80, worstP80 time.Duration
221+
var bestP95, worstP95 time.Duration
222+
var bestTotal, worstTotal time.Duration
223+
var bestSuccess, worstSuccess int
224+
225+
for i, r := range results {
226+
if i == 0 || r.avg < bestAvg {
227+
bestAvg = r.avg
228+
}
229+
if i == 0 || r.avg > worstAvg {
230+
worstAvg = r.avg
231+
}
232+
if i == 0 || r.median < bestMedian {
233+
bestMedian = r.median
234+
}
235+
if i == 0 || r.median > worstMedian {
236+
worstMedian = r.median
237+
}
238+
if i == 0 || r.p5 < bestP5 {
239+
bestP5 = r.p5
240+
}
241+
if i == 0 || r.p5 > worstP5 {
242+
worstP5 = r.p5
243+
}
244+
if i == 0 || r.p75 < bestP75 {
245+
bestP75 = r.p75
246+
}
247+
if i == 0 || r.p75 > worstP75 {
248+
worstP75 = r.p75
249+
}
250+
if i == 0 || r.p80 < bestP80 {
251+
bestP80 = r.p80
252+
}
253+
if i == 0 || r.p80 > worstP80 {
254+
worstP80 = r.p80
255+
}
256+
if i == 0 || r.p95 < bestP95 {
257+
bestP95 = r.p95
258+
}
259+
if i == 0 || r.p95 > worstP95 {
260+
worstP95 = r.p95
261+
}
262+
if i == 0 || r.total < bestTotal {
263+
bestTotal = r.total
264+
}
265+
if i == 0 || r.total > worstTotal {
266+
worstTotal = r.total
267+
}
268+
if i == 0 || r.successful > bestSuccess {
269+
bestSuccess = r.successful
270+
}
271+
if i == 0 || r.successful < worstSuccess {
272+
worstSuccess = r.successful
273+
}
274+
}
275+
276+
// Print each row
277+
for _, r := range results {
278+
fmt.Printf("%-20s | %-19s | %-19s | %-19s | %-19s | %-19s | %-19s | %-19s | %s\n",
279+
r.name,
280+
formatDuration(r.avg, r.avg == bestAvg, r.avg == worstAvg),
281+
formatDuration(r.median, r.median == bestMedian, r.median == worstMedian),
282+
formatDuration(r.p5, r.p5 == bestP5, r.p5 == worstP5),
283+
formatDuration(r.p75, r.p75 == bestP75, r.p75 == worstP75),
284+
formatDuration(r.p80, r.p80 == bestP80, r.p80 == worstP80),
285+
formatDuration(r.p95, r.p95 == bestP95, r.p95 == worstP95),
286+
formatDuration(r.total, r.total == bestTotal, r.total == worstTotal),
287+
formatSuccessRate(r.successful, *requestCount, r.successful == bestSuccess, r.successful == worstSuccess))
288+
}
289+
}
290+
291+
func writeResultsToCSV(filename string, results []endpointResult, requestCount *int) error {
292+
file, err := os.Create(filename)
293+
if err != nil {
294+
return fmt.Errorf("failed to create CSV file: %v", err)
295+
}
296+
defer func() {
297+
err := file.Close()
298+
if err != nil {
299+
return
300+
}
301+
}()
302+
303+
writer := csv.NewWriter(file)
304+
defer writer.Flush()
305+
306+
// Write header
307+
header := []string{"Endpoint", "Average (ms)", "Median (ms)", "P5 (ms)", "P75 (ms)", "P80 (ms)", "P95 (ms)", "Total (ms)", "Success Rate"}
308+
if err := writer.Write(header); err != nil {
309+
return fmt.Errorf("failed to write CSV header: %v", err)
310+
}
311+
312+
// Write data rows
313+
for _, r := range results {
314+
row := []string{
315+
r.name,
316+
fmt.Sprintf("%.2f", float64(r.avg.Microseconds())/1000),
317+
fmt.Sprintf("%.2f", float64(r.median.Microseconds())/1000),
318+
fmt.Sprintf("%.2f", float64(r.p5.Microseconds())/1000),
319+
fmt.Sprintf("%.2f", float64(r.p75.Microseconds())/1000),
320+
fmt.Sprintf("%.2f", float64(r.p80.Microseconds())/1000),
321+
fmt.Sprintf("%.2f", float64(r.p95.Microseconds())/1000),
322+
fmt.Sprintf("%.2f", float64(r.total.Microseconds())/1000),
323+
fmt.Sprintf("%d/%d", r.successful, *requestCount),
324+
}
325+
if err := writer.Write(row); err != nil {
326+
return fmt.Errorf("failed to write CSV row: %v", err)
327+
}
328+
}
329+
330+
return nil
331+
}

cmd/src/main.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ The commands are:
5353
config manages global, org, and user settings
5454
extensions,ext manages extensions (experimental)
5555
extsvc manages external services
56+
gateway interacts with Cody Gateway
5657
login authenticate to a Sourcegraph instance with your user credentials
5758
lsif manages LSIF data (deprecated: use 'code-intel')
5859
orgs,org manages organizations

0 commit comments

Comments
 (0)