examples/server/distributed: a stateless distributed server
Add an example of a server that proxies HTTP requests to multiple stateless streamable MCP backends. For #148
This commit is contained in:
parent
24fc13eb88
commit
07b65d76a7
3 changed files with 198 additions and 8 deletions
|
|
@ -33,7 +33,7 @@ var (
|
||||||
workers = flag.Int("workers", 10, "number of concurrent workers")
|
workers = flag.Int("workers", 10, "number of concurrent workers")
|
||||||
timeout = flag.Duration("timeout", 1*time.Second, "request timeout")
|
timeout = flag.Duration("timeout", 1*time.Second, "request timeout")
|
||||||
qps = flag.Int("qps", 100, "tool calls per second, per worker")
|
qps = flag.Int("qps", 100, "tool calls per second, per worker")
|
||||||
v = flag.Bool("v", false, "if set, enable verbose logging of results")
|
verbose = flag.Bool("v", false, "if set, enable verbose logging")
|
||||||
)
|
)
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
|
|
@ -56,8 +56,8 @@ func main() {
|
||||||
|
|
||||||
parentCtx, cancel := context.WithTimeout(context.Background(), *duration)
|
parentCtx, cancel := context.WithTimeout(context.Background(), *duration)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
parentCtx, restoreSignal := signal.NotifyContext(parentCtx, os.Interrupt)
|
parentCtx, stop := signal.NotifyContext(parentCtx, os.Interrupt)
|
||||||
defer restoreSignal()
|
defer stop()
|
||||||
|
|
||||||
var (
|
var (
|
||||||
start = time.Now()
|
start = time.Now()
|
||||||
|
|
@ -91,12 +91,12 @@ func main() {
|
||||||
return // test ended
|
return // test ended
|
||||||
}
|
}
|
||||||
failure.Add(1)
|
failure.Add(1)
|
||||||
if *v {
|
if *verbose {
|
||||||
log.Printf("FAILURE: %v", err)
|
log.Printf("FAILURE: %v", err)
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
success.Add(1)
|
success.Add(1)
|
||||||
if *v {
|
if *verbose {
|
||||||
data, err := json.Marshal(res)
|
data, err := json.Marshal(res)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Fatalf("marshalling result: %v", err)
|
log.Fatalf("marshalling result: %v", err)
|
||||||
|
|
@ -108,7 +108,7 @@ func main() {
|
||||||
}()
|
}()
|
||||||
}
|
}
|
||||||
wg.Wait()
|
wg.Wait()
|
||||||
restoreSignal() // call restore signal (redundantly) here to allow ctrl-c to work again
|
stop() // restore the interrupt signal
|
||||||
|
|
||||||
// Print stats.
|
// Print stats.
|
||||||
var (
|
var (
|
||||||
|
|
|
||||||
166
examples/server/distributed/main.go
Normal file
166
examples/server/distributed/main.go
Normal file
|
|
@ -0,0 +1,166 @@
|
||||||
|
// 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 distributed command is an example of a distributed MCP server.
|
||||||
|
//
|
||||||
|
// It forks multiple child processes (according to the -child_ports flag), each
|
||||||
|
// of which is a streamable HTTP MCP server with the 'inc' tool, and proxies
|
||||||
|
// incoming http requests to them.
|
||||||
|
//
|
||||||
|
// Distributed MCP servers must be stateless, because there's no guarantee that
|
||||||
|
// subsequent requests for a session land on the same backend. However, they
|
||||||
|
// may still have logical session IDs, as can be seen with verbose logging
|
||||||
|
// (-v).
|
||||||
|
//
|
||||||
|
// Example:
|
||||||
|
//
|
||||||
|
// ./distributed -http=localhost:8080 -child_ports=8081,8082
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"flag"
|
||||||
|
"fmt"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httputil"
|
||||||
|
"net/url"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"os/signal"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
const childPortVar = "MCP_CHILD_PORT"
|
||||||
|
|
||||||
|
var (
|
||||||
|
httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout")
|
||||||
|
childPorts = flag.String("child_ports", "", "comma-separated child ports to distribute to")
|
||||||
|
verbose = flag.Bool("v", false, "if set, enable verbose logging")
|
||||||
|
)
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
// This command runs as either a parent or a child, depending on whether
|
||||||
|
// childPortVar is set (a.k.a. the fork-and-exec trick).
|
||||||
|
//
|
||||||
|
// Each child is a streamable HTTP server, and the parent is a reverse proxy.
|
||||||
|
flag.Parse()
|
||||||
|
if v := os.Getenv(childPortVar); v != "" {
|
||||||
|
child(v)
|
||||||
|
} else {
|
||||||
|
parent()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func parent() {
|
||||||
|
exe, err := os.Executable()
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if *httpAddr == "" {
|
||||||
|
log.Fatal("must provide -http")
|
||||||
|
}
|
||||||
|
if *childPorts == "" {
|
||||||
|
log.Fatal("must provide -child_ports")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ensure that children are cleaned up on CTRL-C
|
||||||
|
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
|
||||||
|
defer stop()
|
||||||
|
|
||||||
|
// Start the child processes.
|
||||||
|
ports := strings.Split(*childPorts, ",")
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
childURLs := make([]*url.URL, len(ports))
|
||||||
|
for i, port := range ports {
|
||||||
|
wg.Add(1)
|
||||||
|
childURL := fmt.Sprintf("http://localhost:%s", port)
|
||||||
|
childURLs[i], err = url.Parse(childURL)
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
cmd := exec.CommandContext(ctx, exe, os.Args[1:]...)
|
||||||
|
cmd.Env = append(os.Environ(), fmt.Sprintf("%s=%s", childPortVar, port))
|
||||||
|
cmd.Stderr = os.Stderr
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
log.Printf("starting child %d at %s", i, childURL)
|
||||||
|
if err := cmd.Run(); err != nil && ctx.Err() == nil {
|
||||||
|
log.Printf("child %d failed: %v", i, err)
|
||||||
|
} else {
|
||||||
|
log.Printf("child %d exited gracefully", i)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start a reverse proxy that round-robin's requests to each backend.
|
||||||
|
var nextBackend atomic.Int64
|
||||||
|
server := http.Server{
|
||||||
|
Addr: *httpAddr,
|
||||||
|
Handler: &httputil.ReverseProxy{
|
||||||
|
Rewrite: func(r *httputil.ProxyRequest) {
|
||||||
|
child := int(nextBackend.Add(1)) % len(childURLs)
|
||||||
|
if *verbose {
|
||||||
|
log.Printf("dispatching to localhost:%s", ports[child])
|
||||||
|
}
|
||||||
|
r.SetURL(childURLs[child])
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
if err := server.ListenAndServe(); err != nil && ctx.Err() == nil {
|
||||||
|
log.Printf("Server failed: %v", err)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
log.Printf("Serving at %s (CTRL-C to cancel)", *httpAddr)
|
||||||
|
|
||||||
|
<-ctx.Done()
|
||||||
|
stop() // restore the interrupt signal
|
||||||
|
|
||||||
|
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
// Attempt the graceful shutdown.
|
||||||
|
if err := server.Shutdown(shutdownCtx); err != nil {
|
||||||
|
log.Fatalf("Server shutdown failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for the subprocesses and http server to stop.
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
log.Println("Server shutdown gracefully.")
|
||||||
|
}
|
||||||
|
|
||||||
|
func child(port string) {
|
||||||
|
// Create a server with a single tool that increments a counter.
|
||||||
|
server := mcp.NewServer(&mcp.Implementation{Name: "counter"}, nil)
|
||||||
|
|
||||||
|
var count atomic.Int64
|
||||||
|
inc := func(ctx context.Context, req *mcp.CallToolRequest, args struct{}) (*mcp.CallToolResult, struct{ Count int64 }, error) {
|
||||||
|
n := count.Add(1)
|
||||||
|
if *verbose {
|
||||||
|
log.Printf("request %d (session %s)", n, req.Session.ID())
|
||||||
|
}
|
||||||
|
return nil, struct{ Count int64 }{n}, nil
|
||||||
|
}
|
||||||
|
mcp.AddTool(server, &mcp.Tool{Name: "inc"}, inc)
|
||||||
|
|
||||||
|
handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {
|
||||||
|
return server
|
||||||
|
}, &mcp.StreamableHTTPOptions{
|
||||||
|
Stateless: true,
|
||||||
|
})
|
||||||
|
log.Printf("child listening on localhost:%s", port)
|
||||||
|
log.Fatal(http.ListenAndServe(fmt.Sprintf("localhost:%s", port), handler))
|
||||||
|
}
|
||||||
|
|
@ -11,19 +11,40 @@ import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"log"
|
"log"
|
||||||
"net/http"
|
"net/http"
|
||||||
|
_ "net/http/pprof"
|
||||||
"net/url"
|
"net/url"
|
||||||
"os"
|
"os"
|
||||||
|
"runtime"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/google/jsonschema-go/jsonschema"
|
"github.com/google/jsonschema-go/jsonschema"
|
||||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
)
|
)
|
||||||
|
|
||||||
var httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout")
|
var (
|
||||||
|
httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout")
|
||||||
|
pprofAddr = flag.String("pprof", "", "if set, host the pprof debugging server at this address")
|
||||||
|
)
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
flag.Parse()
|
flag.Parse()
|
||||||
|
|
||||||
|
if *pprofAddr != "" {
|
||||||
|
// For debugging memory leaks, add an endpoint to trigger a few garbage
|
||||||
|
// collection cycles and ensure the /debug/pprof/heap endpoint only reports
|
||||||
|
// reachable memory.
|
||||||
|
http.DefaultServeMux.Handle("/gc", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||||||
|
for range 3 {
|
||||||
|
runtime.GC()
|
||||||
|
}
|
||||||
|
fmt.Fprintln(w, "GC'ed")
|
||||||
|
}))
|
||||||
|
go func() {
|
||||||
|
// DefaultServeMux was mutated by the /debug/pprof import.
|
||||||
|
http.ListenAndServe(*pprofAddr, http.DefaultServeMux)
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
opts := &mcp.ServerOptions{
|
opts := &mcp.ServerOptions{
|
||||||
Instructions: "Use this server!",
|
Instructions: "Use this server!",
|
||||||
CompletionHandler: complete, // support completions by setting this handler
|
CompletionHandler: complete, // support completions by setting this handler
|
||||||
|
|
@ -56,7 +77,10 @@ func main() {
|
||||||
return server
|
return server
|
||||||
}, nil)
|
}, nil)
|
||||||
log.Printf("MCP handler listening at %s", *httpAddr)
|
log.Printf("MCP handler listening at %s", *httpAddr)
|
||||||
http.ListenAndServe(*httpAddr, handler)
|
if *pprofAddr != "" {
|
||||||
|
log.Printf("pprof listening at http://%s/debug/pprof", *pprofAddr)
|
||||||
|
}
|
||||||
|
log.Fatal(http.ListenAndServe(*httpAddr, handler))
|
||||||
} else {
|
} else {
|
||||||
t := &mcp.LoggingTransport{Transport: &mcp.StdioTransport{}, Writer: os.Stderr}
|
t := &mcp.LoggingTransport{Transport: &mcp.StdioTransport{}, Writer: os.Stderr}
|
||||||
if err := server.Run(context.Background(), t); err != nil {
|
if err := server.Run(context.Background(), t); err != nil {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue