examples/client: add a loadtest command
Add a loadtest client example, to help confirm performance of our streamable transport implementation. For #190
This commit is contained in:
parent
ddaf35ed77
commit
8f11a86898
2 changed files with 125 additions and 3 deletions
|
|
@ -31,10 +31,10 @@ func main() {
|
|||
flag.Parse()
|
||||
args := flag.Args()
|
||||
if len(args) == 0 {
|
||||
fmt.Fprintf(os.Stderr, "Usage: listfeatures <command> [<args>]")
|
||||
fmt.Fprintf(os.Stderr, "List all features for a stdio MCP server")
|
||||
fmt.Fprintln(os.Stderr, "Usage: listfeatures <command> [<args>]")
|
||||
fmt.Fprintln(os.Stderr, "List all features for a stdio MCP server")
|
||||
fmt.Fprintln(os.Stderr)
|
||||
fmt.Fprintf(os.Stderr, "Example: listfeatures npx @modelcontextprotocol/server-everything")
|
||||
fmt.Fprintln(os.Stderr, "Example:\n\tlistfeatures npx @modelcontextprotocol/server-everything")
|
||||
os.Exit(2)
|
||||
}
|
||||
|
||||
|
|
|
|||
122
examples/client/loadtest/main.go
Normal file
122
examples/client/loadtest/main.go
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// The load command load tests a streamable MCP server
|
||||
//
|
||||
// Usage: loadtest <URL>
|
||||
//
|
||||
// For example:
|
||||
//
|
||||
// loadtest -tool=greet -args='{"name": "foo"}' http://localhost:8080
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
var (
|
||||
duration = flag.Duration("duration", 1*time.Minute, "duration of the load test")
|
||||
tool = flag.String("tool", "", "tool to call")
|
||||
jsonArgs = flag.String("args", "", "JSON arguments to pass")
|
||||
workers = flag.Int("workers", 10, "number of concurrent workers")
|
||||
timeout = flag.Duration("timeout", 1*time.Second, "request timeout")
|
||||
qps = flag.Int("qps", 100, "tool calls per second, per worker")
|
||||
v = flag.Bool("v", false, "if set, enable verbose logging of results")
|
||||
)
|
||||
|
||||
func main() {
|
||||
flag.Usage = func() {
|
||||
out := flag.CommandLine.Output()
|
||||
fmt.Fprintf(out, "Usage: loadtest [flags] <URL>")
|
||||
fmt.Fprintf(out, "Load test a streamable HTTP server (CTRL-C to end early)")
|
||||
fmt.Fprintln(out)
|
||||
fmt.Fprintf(out, "Example: loadtest -tool=greet -args='{\"name\": \"foo\"}' http://localhost:8080\n")
|
||||
fmt.Fprintln(out)
|
||||
fmt.Fprintln(out, "Flags:")
|
||||
flag.PrintDefaults()
|
||||
}
|
||||
flag.Parse()
|
||||
args := flag.Args()
|
||||
if len(args) != 1 || *tool == "" {
|
||||
flag.Usage()
|
||||
os.Exit(2)
|
||||
}
|
||||
|
||||
parentCtx, cancel := context.WithTimeout(context.Background(), *duration)
|
||||
defer cancel()
|
||||
parentCtx, restoreSignal := signal.NotifyContext(parentCtx, os.Interrupt)
|
||||
defer restoreSignal()
|
||||
|
||||
var (
|
||||
start = time.Now()
|
||||
success atomic.Int64
|
||||
failure atomic.Int64
|
||||
)
|
||||
|
||||
// Run the test.
|
||||
var wg sync.WaitGroup
|
||||
for range *workers {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
client := mcp.NewClient(&mcp.Implementation{Name: "mcp-client", Version: "v1.0.0"}, nil)
|
||||
cs, err := client.Connect(parentCtx, &mcp.StreamableClientTransport{Endpoint: args[0]}, nil)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
defer cs.Close()
|
||||
|
||||
ticker := time.NewTicker(1 * time.Second / time.Duration(*qps))
|
||||
defer ticker.Stop()
|
||||
|
||||
for range ticker.C {
|
||||
ctx, cancel := context.WithTimeout(parentCtx, *timeout)
|
||||
defer cancel()
|
||||
|
||||
res, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: *tool, Arguments: json.RawMessage(*jsonArgs)})
|
||||
if err != nil {
|
||||
if parentCtx.Err() != nil {
|
||||
return // test ended
|
||||
}
|
||||
failure.Add(1)
|
||||
if *v {
|
||||
log.Printf("FAILURE: %v", err)
|
||||
}
|
||||
} else {
|
||||
success.Add(1)
|
||||
if *v {
|
||||
data, err := json.Marshal(res)
|
||||
if err != nil {
|
||||
log.Fatalf("marshalling result: %v", err)
|
||||
}
|
||||
log.Printf("SUCCESS: %s", string(data))
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
restoreSignal() // call restore signal (redundantly) here to allow ctrl-c to work again
|
||||
|
||||
// Print stats.
|
||||
var (
|
||||
dur = time.Since(start)
|
||||
succ = success.Load()
|
||||
fail = failure.Load()
|
||||
)
|
||||
fmt.Printf("Results (in %s):\n", dur)
|
||||
fmt.Printf("\tsuccess: %d (%g QPS)\n", succ, float64(succ)/dur.Seconds())
|
||||
fmt.Printf("\tfailure: %d (%g QPS)\n", fail, float64(fail)/dur.Seconds())
|
||||
}
|
||||
Loading…
Reference in a new issue