@@ -20,13 +20,10 @@ import (
2020 "context"
2121 "encoding/json"
2222 "errors"
23- "expvar"
2423 "flag"
2524 "fmt"
2625 "io"
2726 "net"
28- "net/http"
29- "net/http/pprof"
3027 "os"
3128 "path/filepath"
3229 "runtime"
@@ -344,6 +341,8 @@ func run(ctx context.Context, manager Manager, config Config) error {
344341 ttrpcServices = []TTRPCService {}
345342
346343 ttrpcUnaryInterceptors = []ttrpc.UnaryServerInterceptor {}
344+
345+ pprofHandler server
347346 )
348347
349348 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 {
397396 ttrpcUnaryInterceptors = append (ttrpcUnaryInterceptors , src .UnaryServerInterceptor ())
398397 }
399398
399+ if result .Registration .ID == "pprof" {
400+ if src , ok := instance .(server ); ok {
401+ pprofHandler = src
402+ }
400403 }
401404 }
402405
@@ -416,7 +419,7 @@ func run(ctx context.Context, manager Manager, config Config) error {
416419 }
417420 }
418421
419- if err := serve (ctx , server , signals , sd .Shutdown ); err != nil {
422+ if err := serve (ctx , server , signals , sd .Shutdown , pprofHandler ); err != nil {
420423 if ! errors .Is (err , shutdown .ErrShutdown ) {
421424 cleanupSockets (ctx )
422425 return err
@@ -437,7 +440,7 @@ func run(ctx context.Context, manager Manager, config Config) error {
437440
438441// serve serves the ttrpc API over a unix socket in the current working directory
439442// and blocks until the context is canceled
440- func serve (ctx context.Context , server * ttrpc.Server , signals chan os.Signal , shutdown func ()) error {
443+ func serve (ctx context.Context , server * ttrpc.Server , signals chan os.Signal , shutdown func (), pprof server ) error {
441444 dump := make (chan os.Signal , 32 )
442445 setupDumpStacks (dump )
443446
@@ -457,9 +460,9 @@ func serve(ctx context.Context, server *ttrpc.Server, signals chan os.Signal, sh
457460 }
458461 }()
459462
460- if debugFlag {
461- if err := serveDebug (ctx ); err != nil {
462- return err
463+ if debugFlag && pprof != nil {
464+ if err := setupPprof (ctx , pprof ); err != nil {
465+ log . G ( ctx ). WithError ( err ). Warn ( "Could not setup pprof" )
463466 }
464467 }
465468
@@ -478,31 +481,6 @@ func serve(ctx context.Context, server *ttrpc.Server, signals chan os.Signal, sh
478481 return reap (ctx , logger , signals )
479482}
480483
481- func serveDebug (ctx context.Context ) error {
482- l , err := serveListener (debugSocketFlag , 4 )
483- if err != nil {
484- return err
485- }
486- go func () {
487- defer l .Close ()
488- m := http .NewServeMux ()
489- m .Handle ("/debug/vars" , expvar .Handler ())
490- m .Handle ("/debug/pprof/" , http .HandlerFunc (pprof .Index ))
491- m .Handle ("/debug/pprof/cmdline" , http .HandlerFunc (pprof .Cmdline ))
492- m .Handle ("/debug/pprof/profile" , http .HandlerFunc (pprof .Profile ))
493- m .Handle ("/debug/pprof/symbol" , http .HandlerFunc (pprof .Symbol ))
494- m .Handle ("/debug/pprof/trace" , http .HandlerFunc (pprof .Trace ))
495- srv := & http.Server {
496- Handler : m ,
497- ReadHeaderTimeout : 5 * time .Minute ,
498- }
499- if err := srv .Serve (l ); err != nil && ! errors .Is (err , net .ErrClosed ) {
500- log .G (ctx ).WithError (err ).Fatal ("containerd-shim: pprof endpoint failure" )
501- }
502- }()
503- return nil
504- }
505-
506484func dumpStacks (logger * log.Entry ) {
507485 var (
508486 buf []byte
@@ -517,3 +495,22 @@ func dumpStacks(logger *log.Entry) {
517495 buf = buf [:stackSize ]
518496 logger .Infof ("=== BEGIN goroutine stack dump ===\n %s\n === END goroutine stack dump ===" , buf )
519497}
498+
499+ type server interface {
500+ Serve (net.Listener ) error
501+ }
502+
503+ func setupPprof (ctx context.Context , srv server ) error {
504+ l , err := serveListener (debugSocketFlag , 4 )
505+ if err != nil {
506+ return fmt .Errorf ("could not setup pprof listener: %w" , err )
507+ }
508+
509+ go func () {
510+ if err := srv .Serve (l ); err != nil && ! errors .Is (err , net .ErrClosed ) {
511+ log .G (ctx ).WithError (err ).Fatal ("containerd-shim: pprof endpoint failure" )
512+ }
513+ }()
514+
515+ return nil
516+ }
0 commit comments