shim: Move pprof server to plugin

Makes the pprof server a plugin and also gates by the `shim_tracing`
build tag (like otel is).
With this change, `net/http` is no longer a dependency in the shim.

Signed-off-by: Brian Goff <cpuguy83@gmail.com>

pkg/shim/shim.go

@@ -20,13 +20,10 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
-	"expvar"
 	"flag"
 	"fmt"
 	"io"
 	"net"
-	"net/http"
-	"net/http/pprof"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -344,6 +341,8 @@ func run(ctx context.Context, manager Manager, config Config) error {
 		ttrpcServices = []TTRPCService{}
 
 		ttrpcUnaryInterceptors = []ttrpc.UnaryServerInterceptor{}
+
+		pprofHandler server
 	)
 
 	for _, p := range registry.Graph(func(*plugin.Registration) bool { return false }) {
@@ -397,6 +396,10 @@ func run(ctx context.Context, manager Manager, config Config) error {
 			ttrpcUnaryInterceptors = append(ttrpcUnaryInterceptors, src.UnaryServerInterceptor())
 		}
 
+		if result.Registration.ID == "pprof" {
+			if src, ok := instance.(server); ok {
+				pprofHandler = src
+			}
 		}
 	}
 
@@ -416,7 +419,7 @@ func run(ctx context.Context, manager Manager, config Config) error {
 		}
 	}
 
-	if err := serve(ctx, server, signals, sd.Shutdown); err != nil {
+	if err := serve(ctx, server, signals, sd.Shutdown, pprofHandler); err != nil {
 		if !errors.Is(err, shutdown.ErrShutdown) {
 			cleanupSockets(ctx)
 			return err
@@ -437,7 +440,7 @@ func run(ctx context.Context, manager Manager, config Config) error {
 
 // serve serves the ttrpc API over a unix socket in the current working directory
 // and blocks until the context is canceled
-func serve(ctx context.Context, server *ttrpc.Server, signals chan os.Signal, shutdown func()) error {
+func serve(ctx context.Context, server *ttrpc.Server, signals chan os.Signal, shutdown func(), pprof server) error {
 	dump := make(chan os.Signal, 32)
 	setupDumpStacks(dump)
 
@@ -457,9 +460,9 @@ func serve(ctx context.Context, server *ttrpc.Server, signals chan os.Signal, sh
 		}
 	}()
 
-	if debugFlag {
-		if err := serveDebug(ctx); err != nil {
-			return err
+	if debugFlag && pprof != nil {
+		if err := setupPprof(ctx, pprof); err != nil {
+			log.G(ctx).WithError(err).Warn("Could not setup pprof")
 		}
 	}
 
@@ -478,31 +481,6 @@ func serve(ctx context.Context, server *ttrpc.Server, signals chan os.Signal, sh
 	return reap(ctx, logger, signals)
 }
 
-func serveDebug(ctx context.Context) error {
-	l, err := serveListener(debugSocketFlag, 4)
-	if err != nil {
-		return err
-	}
-	go func() {
-		defer l.Close()
-		m := http.NewServeMux()
-		m.Handle("/debug/vars", expvar.Handler())
-		m.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
-		m.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
-		m.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
-		m.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
-		m.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
-		srv := &http.Server{
-			Handler:           m,
-			ReadHeaderTimeout: 5 * time.Minute,
-		}
-		if err := srv.Serve(l); err != nil && !errors.Is(err, net.ErrClosed) {
-			log.G(ctx).WithError(err).Fatal("containerd-shim: pprof endpoint failure")
-		}
-	}()
-	return nil
-}
-
 func dumpStacks(logger *log.Entry) {
 	var (
 		buf       []byte
@@ -517,3 +495,22 @@ func dumpStacks(logger *log.Entry) {
 	buf = buf[:stackSize]
 	logger.Infof("=== BEGIN goroutine stack dump ===\n%s\n=== END goroutine stack dump ===", buf)
 }
+
+type server interface {
+	Serve(net.Listener) error
+}
+
+func setupPprof(ctx context.Context, srv server) error {
+	l, err := serveListener(debugSocketFlag, 4)
+	if err != nil {
+		return fmt.Errorf("could not setup pprof listener: %w", err)
+	}
+
+	go func() {
+		if err := srv.Serve(l); err != nil && !errors.Is(err, net.ErrClosed) {
+			log.G(ctx).WithError(err).Fatal("containerd-shim: pprof endpoint failure")
+		}
+	}()
+
+	return nil
+}