diff --git a/cmd/kepler/main.go b/cmd/kepler/main.go index caea8f0e35..5837e5098b 100644 --- a/cmd/kepler/main.go +++ b/cmd/kepler/main.go @@ -229,11 +229,18 @@ func createServices(logger *slog.Logger, cfg *config.Config) ([]service.Service, prometheus.WithCollectors(collectors), ) - return []service.Service{ + services := []service.Service{ promExporter, apiServer, pm, - }, nil + } + + if cfg.EnablePprof { + pprof := server.NewPprof(apiServer) + services = append(services, pprof) + } + + return services, nil } func createPowerMonitor(logger *slog.Logger, cfg *config.Config) (*monitor.PowerMonitor, error) { diff --git a/config/config.go b/config/config.go index 0ff47637ef..4ddf88f02a 100644 --- a/config/config.go +++ b/config/config.go @@ -33,19 +33,23 @@ type ( } Config struct { - Log Log `yaml:"log"` - Host Host `yaml:"host"` - Dev Dev `yaml:"dev"` // WARN: do not expose dev settings as flags + Log Log `yaml:"log"` + Host Host `yaml:"host"` + Dev Dev `yaml:"dev"` // WARN: do not expose dev settings as flags + EnablePprof bool `yaml:"enable-pprof"` } ) const ( // Flags - LogLevelFlag = "log.level" - LogFormatFlag = "log.format" + LogLevelFlag = "log.level" + LogFormatFlag = "log.format" + HostSysFSFlag = "host.sysfs" HostProcFSFlag = "host.procfs" + EnablePprofFlag = "enable.pprof" + // WARN: dev settings shouldn't be exposed as flags as flags are intended for end users ) @@ -123,6 +127,7 @@ func RegisterFlags(app *kingpin.Application) ConfigUpdaterFn { logFormat := app.Flag(LogFormatFlag, "Logging format: text or json").Default("text").Enum("text", "json") hostSysFS := app.Flag(HostSysFSFlag, "Host sysfs path").Default("/sys").ExistingDir() hostProcFS := app.Flag(HostProcFSFlag, "Host procfs path").Default("/proc").ExistingDir() + enablePprof := app.Flag(EnablePprofFlag, "Enable pprof").Default("false").Bool() return func(cfg *Config) error { // Logging settings if flagsSet[LogLevelFlag] { @@ -141,6 +146,10 @@ func RegisterFlags(app *kingpin.Application) ConfigUpdaterFn { cfg.Host.ProcFS = *hostProcFS } + if flagsSet[EnablePprofFlag] { + cfg.EnablePprof = *enablePprof + } + cfg.sanitize() return cfg.Validate() } diff --git a/config/config_test.go b/config/config_test.go index 417742c698..088bd274a3 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -379,3 +379,36 @@ func TestConfigString(t *testing.T) { }) } } + +func TestEnablePprof(t *testing.T) { + tt := []struct { + name string + args []string + enabled bool + }{{ + name: "enable pprof with flag", + args: []string{"--enable.pprof"}, + enabled: true, + }, { + name: "disable pprof no flag", + args: []string{"--log.level=debug"}, + enabled: false, + }, { + name: "disable pprof with flag", + args: []string{"--no-enable.pprof"}, + enabled: false, + }} + + for _, tc := range tt { + t.Run(tc.name, func(t *testing.T) { + app := kingpin.New("test", "Test application") + updateConfig := RegisterFlags(app) + _, parseErr := app.Parse(tc.args) + assert.NoError(t, parseErr, "unexpected flag parsing error") + cfg := DefaultConfig() + err := updateConfig(cfg) + assert.NoError(t, err, "unexpected config update error") + assert.Equal(t, cfg.EnablePprof, tc.enabled, "unexpected flag value") + }) + } +} diff --git a/internal/server/pprof.go b/internal/server/pprof.go new file mode 100644 index 0000000000..13307864ea --- /dev/null +++ b/internal/server/pprof.go @@ -0,0 +1,45 @@ +// SPDX-FileCopyrightText: 2025 The Kepler Authors +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "net/http" + "net/http/pprof" + + "github.com/sustainable-computing-io/kepler/internal/service" +) + +type pp struct { + api APIService +} + +var _ service.Service = (*pp)(nil) +var _ service.Initializer = (*pp)(nil) + +func NewPprof(api APIService) *pp { + return &pp{ + api: api, + } +} + +func (p *pp) Name() string { + return "pprof" +} + +func (p *pp) Init(ctx context.Context) error { + return p.api.Register("/debug/pprof/", "pprof", "Profiling Data", handlers()) +} + +func handlers() http.Handler { + mux := http.NewServeMux() + + 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) + + return mux +} diff --git a/internal/server/pprof_test.go b/internal/server/pprof_test.go new file mode 100644 index 0000000000..00cf93d76c --- /dev/null +++ b/internal/server/pprof_test.go @@ -0,0 +1,103 @@ +// SPDX-FileCopyrightText: 2025 The Kepler Authors +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "net/http/pprof" +) + +// MockAPIService is an implementation of the APIService interface for testing. +type MockAPIService struct { + mock.Mock +} + +func (m *MockAPIService) Register(path, name, description string, handler http.Handler) error { + args := m.Called(path, name, description, handler) + return args.Error(0) +} + +func (m *MockAPIService) Name() string { + return "mockApiService" +} + +// TestNewPprof tests the NewPprof constructor. +func TestNewPprof(t *testing.T) { + api := &MockAPIService{} + p := NewPprof(api) + + assert.NotNil(t, p, "NewPprof should return a non-nil pointer") + assert.Equal(t, api, p.api, "NewPprof should set the api field correctly") +} + +// TestPprofName tests the Name method. +func TestPprofName(t *testing.T) { + api := &MockAPIService{} + p := NewPprof(api) + + name := p.Name() + assert.Equal(t, "pprof", name, "Name should return 'pprof'") +} + +// TestPprofInit_Success tests the Init method when the API registration succeeds. +func TestPprofInit_Success(t *testing.T) { + api := &MockAPIService{} + p := NewPprof(api) + + // Set up mock expectation + api.On("Register", "/debug/pprof/", "pprof", "Profiling Data", mock.AnythingOfType("*http.ServeMux")).Return(nil) + + err := p.Init(context.Background()) + assert.NoError(t, err, "Init should not return an error when registration succeeds") + api.AssertExpectations(t) +} + +// TestPprofInit_Failure tests the Init method when the API registration fails. +func TestPprofInit_Failure(t *testing.T) { + api := &MockAPIService{} + p := NewPprof(api) + + // Set up mock expectation + expectedErr := assert.AnError + api.On("Register", "/debug/pprof/", "pprof", "Profiling Data", mock.AnythingOfType("*http.ServeMux")).Return(expectedErr) + + err := p.Init(context.Background()) + assert.Error(t, err, "Init should return an error when registration fails") + assert.Equal(t, expectedErr, err, "Init should return the expected error") + api.AssertExpectations(t) +} + +// TestPprofHandlers tests the handlers function to ensure it registers the correct pprof endpoints. +func TestPprofHandlers(t *testing.T) { + handler := handlers() + mux, ok := handler.(*http.ServeMux) + assert.True(t, ok, "handlers should return an http.ServeMux") + + // Test cases for each pprof endpoint + tests := []struct { + path string + handlerFunc http.HandlerFunc + }{ + {"/debug/pprof/", pprof.Index}, + {"/debug/pprof/cmdline", pprof.Cmdline}, + {"/debug/pprof/profile", pprof.Profile}, + {"/debug/pprof/symbol", pprof.Symbol}, + {"/debug/pprof/trace", pprof.Trace}, + } + + for _, tt := range tests { + t.Run(tt.path, func(t *testing.T) { + req := httptest.NewRequest("GET", tt.path, nil) + rr := httptest.NewRecorder() + mux.ServeHTTP(rr, req) + assert.NotEqual(t, http.StatusNotFound, rr.Code, "Handler for %s should be registered", tt.path) + }) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 05e975cc92..f2dea5b714 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -159,6 +159,7 @@ func (s *APIServer) Shutdown() error { } func (s *APIServer) Register(endpoint, summary, description string, handler http.Handler) error { + s.logger.Debug("Endpoint Registered", "endpoint", endpoint) s.mux.Handle(endpoint, handler) s.endpointDescription += fmt.Sprintf("
  • %s %s
    • \n", endpoint, summary, description) return nil