diff --git a/cns/client/client_test.go b/cns/client/client_test.go index 17d4090d0b..0a93dd8764 100644 --- a/cns/client/client_test.go +++ b/cns/client/client_test.go @@ -170,8 +170,8 @@ func TestMain(m *testing.M) { config := common.ServiceConfig{} httpRestService, err := restserver.NewHTTPRestService(&config, &fakes.WireserverClientFake{}, &fakes.WireserverProxyFake{}, &fakes.NMAgentClientFake{}, nil, nil, nil) - svc = httpRestService.(*restserver.HTTPRestService) - svc.Name = "cns-test-server" + svc = httpRestService + httpRestService.Name = "cns-test-server" fakeNNC := v1alpha.NodeNetworkConfig{ TypeMeta: metav1.TypeMeta{}, ObjectMeta: metav1.ObjectMeta{}, @@ -204,7 +204,7 @@ func TestMain(m *testing.M) { }, }, } - svc.IPAMPoolMonitor = &fakes.MonitorFake{IPsNotInUseCount: 13, NodeNetworkConfig: &fakeNNC} + httpRestService.IPAMPoolMonitor = &fakes.MonitorFake{IPsNotInUseCount: 13, NodeNetworkConfig: &fakeNNC} if err != nil { logger.Errorf("Failed to create CNS object, err:%v.\n", err) diff --git a/cns/restserver/ipam_test.go b/cns/restserver/ipam_test.go index 69455a11bf..d825fe1407 100644 --- a/cns/restserver/ipam_test.go +++ b/cns/restserver/ipam_test.go @@ -64,11 +64,11 @@ type ncState struct { func getTestService() *HTTPRestService { var config common.ServiceConfig httpsvc, _ := NewHTTPRestService(&config, &fakes.WireserverClientFake{}, &fakes.WireserverProxyFake{}, &fakes.NMAgentClientFake{}, store.NewMockStore(""), nil, nil) - svc = httpsvc.(*HTTPRestService) - svc.IPAMPoolMonitor = &fakes.MonitorFake{} + svc = httpsvc + httpsvc.IPAMPoolMonitor = &fakes.MonitorFake{} setOrchestratorTypeInternal(cns.KubernetesCRD) - return svc + return httpsvc } func newSecondaryIPConfig(ipAddress string, ncVersion int) cns.SecondaryIPConfig { diff --git a/cns/restserver/restserver.go b/cns/restserver/restserver.go index 4a938350c7..70676c28e3 100644 --- a/cns/restserver/restserver.go +++ b/cns/restserver/restserver.go @@ -4,6 +4,7 @@ import ( "context" "net" "net/http" + "net/http/pprof" "sync" "time" @@ -151,7 +152,7 @@ type networkInfo struct { // NewHTTPRestService creates a new HTTP Service object. func NewHTTPRestService(config *common.ServiceConfig, wscli interfaceGetter, wsproxy wireserverProxy, nmagentClient nmagentClient, endpointStateStore store.KeyValueStore, gen CNIConflistGenerator, homeAzMonitor *HomeAzMonitor, -) (cns.HTTPService, error) { +) (*HTTPRestService, error) { service, err := cns.NewService(config.Name, config.Version, config.ChannelMode, config.Store) if err != nil { return nil, err @@ -300,6 +301,23 @@ func (service *HTTPRestService) Init(config *common.ServiceConfig) error { return nil } +func (service *HTTPRestService) RegisterPProfEndpoints() { + if service.Listener != nil { + mux := service.Listener.GetMux() + mux.Handle("/debug/pprof/allocs", pprof.Handler("allocs")) + mux.Handle("/debug/pprof/block", pprof.Handler("block")) + mux.Handle("/debug/pprof/goroutine", pprof.Handler("goroutine")) + mux.Handle("/debug/pprof/heap", pprof.Handler("heap")) + mux.Handle("/debug/pprof/mutex", pprof.Handler("mutex")) + mux.Handle("/debug/pprof/threadcreate", pprof.Handler("threadcreate")) + mux.HandleFunc("/debug/pprof/", pprof.Index) + mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline) + mux.HandleFunc("/debug/pprof/profile", pprof.Profile) + mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol) + mux.HandleFunc("/debug/pprof/trace", pprof.Trace) + } +} + // Start starts the CNS listener. func (service *HTTPRestService) Start(config *common.ServiceConfig) error { // Start the listener. diff --git a/cns/service/main.go b/cns/service/main.go index e857ef4927..b064cc24c9 100644 --- a/cns/service/main.go +++ b/cns/service/main.go @@ -10,7 +10,6 @@ import ( "fmt" "io/fs" "net/http" - "net/http/pprof" "os" "os/signal" "runtime" @@ -764,6 +763,10 @@ func main() { logger.Printf("[Azure CNS] Start HTTP listener") if httpRestService != nil { + if cnsconfig.EnablePprof { + httpRestService.RegisterPProfEndpoints() + } + err = httpRestService.Start(&config) if err != nil { logger.Errorf("Failed to start CNS, err:%v.\n", err) @@ -1237,18 +1240,7 @@ func InitializeCRDState(ctx context.Context, httpRestService cns.HTTPService, cn mux := httpRestServiceImplementation.Listener.GetMux() mux.Handle("/readyz", http.StripPrefix("/readyz", &healthz.Handler{})) if cnsconfig.EnablePprof { - // add pprof endpoints - mux.Handle("/debug/pprof/allocs", pprof.Handler("allocs")) - mux.Handle("/debug/pprof/block", pprof.Handler("block")) - mux.Handle("/debug/pprof/goroutine", pprof.Handler("goroutine")) - mux.Handle("/debug/pprof/heap", pprof.Handler("heap")) - mux.Handle("/debug/pprof/mutex", pprof.Handler("mutex")) - mux.Handle("/debug/pprof/threadcreate", pprof.Handler("threadcreate")) - mux.HandleFunc("/debug/pprof/", pprof.Index) - mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline) - mux.HandleFunc("/debug/pprof/profile", pprof.Profile) - mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol) - mux.HandleFunc("/debug/pprof/trace", pprof.Trace) + httpRestServiceImplementation.RegisterPProfEndpoints() } // Start the Manager which starts the reconcile loop.