Adding Pprof register method (#1885)

* adding pprof register method that can be invoked in multiple paths * fixing lint issues
Azure · Apr 5, 2023 · 6048ba2 · 6048ba2
1 parent 88f70b1
commit 6048ba2
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 20 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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.