examples/server/distributed: a stateless distributed server

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 (

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))
+}

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 {