Add pprof to runc-shim

Signed-off-by: Henry Wang <henwang@amazon.com>

pkg/shim/shim.go

@@ -20,10 +20,13 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"expvar"
 	"flag"
 	"fmt"
 	"io"
 	"net"
+	"net/http"
+	"net/http/pprof"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -121,6 +124,7 @@ var (
 	id                   string
 	namespaceFlag        string
 	socketFlag           string
+	debugSocketFlag      string
 	bundlePath           string
 	addressFlag          string
 	containerdBinaryFlag string
@@ -143,6 +147,7 @@ func parseFlags() {
 	flag.StringVar(&namespaceFlag, "namespace", "", "namespace that owns the shim")
 	flag.StringVar(&id, "id", "", "id of the task")
 	flag.StringVar(&socketFlag, "socket", "", "socket path to serve")
+	flag.StringVar(&debugSocketFlag, "debug-socket", "", "debug socket path to serve")
 	flag.StringVar(&bundlePath, "bundle", "", "path to the bundle if not workdir")
 
 	flag.StringVar(&addressFlag, "address", "", "grpc address back to main containerd")
@@ -435,7 +440,7 @@ func serve(ctx context.Context, server *ttrpc.Server, signals chan os.Signal, sh
 		return err
 	}
 
-	l, err := serveListener(socketFlag)
+	l, err := serveListener(socketFlag, 3)
 	if err != nil {
 		return err
 	}
@@ -445,6 +450,13 @@ func serve(ctx context.Context, server *ttrpc.Server, signals chan os.Signal, sh
 			log.G(ctx).WithError(err).Fatal("containerd-shim: ttrpc server failure")
 		}
 	}()
+
+	if debugFlag {
+		if err := serveDebug(ctx); err != nil {
+			return err
+		}
+	}
+
 	logger := log.G(ctx).WithFields(log.Fields{
 		"pid":       os.Getpid(),
 		"path":      path,
@@ -460,6 +472,31 @@ 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

pkg/shim/shim_unix.go

@@ -49,13 +49,13 @@ func setupDumpStacks(dump chan<- os.Signal) {
 	signal.Notify(dump, syscall.SIGUSR1)
 }
 
-func serveListener(path string) (net.Listener, error) {
+func serveListener(path string, fd uintptr) (net.Listener, error) {
 	var (
 		l   net.Listener
 		err error
 	)
 	if path == "" {
-		l, err = net.FileListener(os.NewFile(3, "socket"))
+		l, err = net.FileListener(os.NewFile(fd, "socket"))
 		path = "[inherited from parent]"
 	} else {
 		if len(path) > socketPathLimit {

pkg/shim/shim_windows.go

@@ -42,7 +42,7 @@ func subreaper() error {
 func setupDumpStacks(dump chan<- os.Signal) {
 }
 
-func serveListener(path string) (net.Listener, error) {
+func serveListener(path string, fd uintptr) (net.Listener, error) {
 	return nil, errdefs.ErrNotImplemented
 }
 

pkg/shim/util_unix.go

@@ -76,12 +76,16 @@ func AdjustOOMScore(pid int) error {
 const socketRoot = defaults.DefaultStateDir
 
 // SocketAddress returns a socket address
-func SocketAddress(ctx context.Context, socketPath, id string) (string, error) {
+func SocketAddress(ctx context.Context, socketPath, id string, debug bool) (string, error) {
 	ns, err := namespaces.NamespaceRequired(ctx)
 	if err != nil {
 		return "", err
 	}
-	d := sha256.Sum256([]byte(filepath.Join(socketPath, ns, id)))
+	path := filepath.Join(socketPath, ns, id)
+	if debug {
+		path = filepath.Join(path, "debug")
+	}
+	d := sha256.Sum256([]byte(path))
 	return fmt.Sprintf("unix://%s/%x", filepath.Join(socketRoot, "s"), d), nil
 }
 
@@ -286,7 +290,12 @@ func cleanupSockets(ctx context.Context) {
 	}
 	if len(socketFlag) > 0 {
 		_ = RemoveSocket("unix://" + socketFlag)
-	} else if address, err := SocketAddress(ctx, addressFlag, id); err == nil {
+	} else if address, err := SocketAddress(ctx, addressFlag, id, false); err == nil {
+		_ = RemoveSocket(address)
+	}
+	if len(debugSocketFlag) > 0 {
+		_ = RemoveSocket("unix://" + debugSocketFlag)
+	} else if address, err := SocketAddress(ctx, addressFlag, id, true); err == nil {
 		_ = RemoveSocket(address)
 	}
 }