From 07b65d76a7fdb710e757f26e2c03418609bc1b4b Mon Sep 17 00:00:00 2001 From: Robert Findley Date: Tue, 2 Sep 2025 14:32:12 -0400 Subject: [PATCH] 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 --- examples/client/loadtest/main.go | 12 +- examples/server/distributed/main.go | 166 ++++++++++++++++++++++++++++ examples/server/everything/main.go | 28 ++++- 3 files changed, 198 insertions(+), 8 deletions(-) create mode 100644 examples/server/distributed/main.go diff --git a/examples/client/loadtest/main.go b/examples/client/loadtest/main.go index 2c6a5c0..d5a04c2 100644 --- a/examples/client/loadtest/main.go +++ b/examples/client/loadtest/main.go @@ -33,7 +33,7 @@ var ( 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") + verbose = flag.Bool("v", false, "if set, enable verbose logging") ) func main() { @@ -56,8 +56,8 @@ func main() { parentCtx, cancel := context.WithTimeout(context.Background(), *duration) defer cancel() - parentCtx, restoreSignal := signal.NotifyContext(parentCtx, os.Interrupt) - defer restoreSignal() + parentCtx, stop := signal.NotifyContext(parentCtx, os.Interrupt) + defer stop() var ( start = time.Now() @@ -91,12 +91,12 @@ func main() { return // test ended } failure.Add(1) - if *v { + if *verbose { log.Printf("FAILURE: %v", err) } } else { success.Add(1) - if *v { + if *verbose { data, err := json.Marshal(res) if err != nil { log.Fatalf("marshalling result: %v", err) @@ -108,7 +108,7 @@ func main() { }() } wg.Wait() - restoreSignal() // call restore signal (redundantly) here to allow ctrl-c to work again + stop() // restore the interrupt signal // Print stats. var ( diff --git a/examples/server/distributed/main.go b/examples/server/distributed/main.go new file mode 100644 index 0000000..43d959e --- /dev/null +++ b/examples/server/distributed/main.go @@ -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)) +} diff --git a/examples/server/everything/main.go b/examples/server/everything/main.go index d2b7b33..0b81919 100644 --- a/examples/server/everything/main.go +++ b/examples/server/everything/main.go @@ -11,19 +11,40 @@ import ( "fmt" "log" "net/http" + _ "net/http/pprof" "net/url" "os" + "runtime" "strings" "github.com/google/jsonschema-go/jsonschema" "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() { 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{ Instructions: "Use this server!", CompletionHandler: complete, // support completions by setting this handler @@ -56,7 +77,10 @@ func main() { return server }, nil) 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 { t := &mcp.LoggingTransport{Transport: &mcp.StdioTransport{}, Writer: os.Stderr} if err := server.Run(context.Background(), t); err != nil {