Skip to content

Commit 6890620

Browse files
committed
feat(pprof): Added pprof
Signed-off-by: Vimal Kumar <[email protected]>
1 parent 2ad09ba commit 6890620

File tree

6 files changed

+205
-7
lines changed

6 files changed

+205
-7
lines changed

cmd/kepler/main.go

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -229,11 +229,18 @@ func createServices(logger *slog.Logger, cfg *config.Config) ([]service.Service,
229229
prometheus.WithCollectors(collectors),
230230
)
231231

232-
return []service.Service{
232+
services := []service.Service{
233233
promExporter,
234234
apiServer,
235235
pm,
236-
}, nil
236+
}
237+
238+
if cfg.EnablePprof {
239+
pprof := server.NewPprof(apiServer)
240+
services = append(services, pprof)
241+
}
242+
243+
return services, nil
237244
}
238245

239246
func createPowerMonitor(logger *slog.Logger, cfg *config.Config) (*monitor.PowerMonitor, error) {

config/config.go

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,19 +33,23 @@ type (
3333
}
3434

3535
Config struct {
36-
Log Log `yaml:"log"`
37-
Host Host `yaml:"host"`
38-
Dev Dev `yaml:"dev"` // WARN: do not expose dev settings as flags
36+
Log Log `yaml:"log"`
37+
Host Host `yaml:"host"`
38+
Dev Dev `yaml:"dev"` // WARN: do not expose dev settings as flags
39+
EnablePprof bool `yaml:"enable-pprof"`
3940
}
4041
)
4142

4243
const (
4344
// Flags
44-
LogLevelFlag = "log.level"
45-
LogFormatFlag = "log.format"
45+
LogLevelFlag = "log.level"
46+
LogFormatFlag = "log.format"
47+
4648
HostSysFSFlag = "host.sysfs"
4749
HostProcFSFlag = "host.procfs"
4850

51+
EnablePprofFlag = "enable.pprof"
52+
4953
// WARN: dev settings shouldn't be exposed as flags as flags are intended for end users
5054
)
5155

@@ -123,6 +127,7 @@ func RegisterFlags(app *kingpin.Application) ConfigUpdaterFn {
123127
logFormat := app.Flag(LogFormatFlag, "Logging format: text or json").Default("text").Enum("text", "json")
124128
hostSysFS := app.Flag(HostSysFSFlag, "Host sysfs path").Default("/sys").ExistingDir()
125129
hostProcFS := app.Flag(HostProcFSFlag, "Host procfs path").Default("/proc").ExistingDir()
130+
enablePprof := app.Flag(EnablePprofFlag, "Enable pprof").Default("false").Bool()
126131
return func(cfg *Config) error {
127132
// Logging settings
128133
if flagsSet[LogLevelFlag] {
@@ -141,6 +146,10 @@ func RegisterFlags(app *kingpin.Application) ConfigUpdaterFn {
141146
cfg.Host.ProcFS = *hostProcFS
142147
}
143148

149+
if flagsSet[EnablePprofFlag] {
150+
cfg.EnablePprof = *enablePprof
151+
}
152+
144153
cfg.sanitize()
145154
return cfg.Validate()
146155
}

config/config_test.go

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,3 +379,36 @@ func TestConfigString(t *testing.T) {
379379
})
380380
}
381381
}
382+
383+
func TestEnablePprof(t *testing.T) {
384+
tt := []struct {
385+
name string
386+
args []string
387+
enabled bool
388+
}{{
389+
name: "enable pprof with flag",
390+
args: []string{"--enable.pprof"},
391+
enabled: true,
392+
}, {
393+
name: "disable pprof no flag",
394+
args: []string{"--log.level=debug"},
395+
enabled: false,
396+
}, {
397+
name: "disable pprof with flag",
398+
args: []string{"--no-enable.pprof"},
399+
enabled: false,
400+
}}
401+
402+
for _, tc := range tt {
403+
t.Run(tc.name, func(t *testing.T) {
404+
app := kingpin.New("test", "Test application")
405+
updateConfig := RegisterFlags(app)
406+
_, parseErr := app.Parse(tc.args)
407+
assert.NoError(t, parseErr, "unexpected flag parsing error")
408+
cfg := DefaultConfig()
409+
err := updateConfig(cfg)
410+
assert.NoError(t, err, "unexpected config update error")
411+
assert.Equal(t, cfg.EnablePprof, tc.enabled, "unexpected flag value")
412+
})
413+
}
414+
}

internal/server/pprof.go

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
// SPDX-FileCopyrightText: 2025 The Kepler Authors
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package server
5+
6+
import (
7+
"context"
8+
"net/http"
9+
"net/http/pprof"
10+
11+
"github.com/sustainable-computing-io/kepler/internal/service"
12+
)
13+
14+
type pp struct {
15+
api APIService
16+
}
17+
18+
var _ service.Service = &pp{}
19+
var _ service.Initializer = &pp{}
20+
21+
func NewPprof(api APIService) *pp {
22+
return &pp{
23+
api: api,
24+
}
25+
}
26+
27+
func (p *pp) Name() string {
28+
return "pprof"
29+
}
30+
31+
func (p *pp) Init(ctx context.Context) error {
32+
return p.api.Register("/debug/pprof/", "pprof", "Profiling Data", handlers())
33+
}
34+
35+
func handlers() http.Handler {
36+
mux := http.NewServeMux()
37+
38+
mux.HandleFunc("/debug/pprof/", pprof.Index)
39+
mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
40+
mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
41+
mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
42+
mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
43+
44+
return mux
45+
}

internal/server/pprof_test.go

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
// SPDX-FileCopyrightText: 2025 The Kepler Authors
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package server
5+
6+
import (
7+
"context"
8+
"net/http"
9+
"net/http/httptest"
10+
"testing"
11+
12+
"github.com/stretchr/testify/assert"
13+
"github.com/stretchr/testify/mock"
14+
"net/http/pprof"
15+
)
16+
17+
// MockAPIService is an implementation of the APIService interface for testing.
18+
type MockAPIService struct {
19+
mock.Mock
20+
}
21+
22+
func (m *MockAPIService) Register(path, name, description string, handler http.Handler) error {
23+
args := m.Called(path, name, description, handler)
24+
return args.Error(0)
25+
}
26+
27+
func (m *MockAPIService) Name() string {
28+
return "mockApiService"
29+
}
30+
31+
// TestNewPprof tests the NewPprof constructor.
32+
func TestNewPprof(t *testing.T) {
33+
api := &MockAPIService{}
34+
p := NewPprof(api)
35+
36+
assert.NotNil(t, p, "NewPprof should return a non-nil pointer")
37+
assert.Equal(t, api, p.api, "NewPprof should set the api field correctly")
38+
}
39+
40+
// TestName tests the Name method.
41+
func TestName(t *testing.T) {
42+
api := &MockAPIService{}
43+
p := NewPprof(api)
44+
45+
name := p.Name()
46+
assert.Equal(t, "pprof", name, "Name should return 'pprof'")
47+
}
48+
49+
// TestInit_Success tests the Init method when the API registration succeeds.
50+
func TestInit_Success(t *testing.T) {
51+
api := &MockAPIService{}
52+
p := NewPprof(api)
53+
54+
// Set up mock expectation
55+
api.On("Register", "/debug/pprof/", "pprof", "Profiling Data", mock.AnythingOfType("*http.ServeMux")).Return(nil)
56+
57+
err := p.Init(context.Background())
58+
assert.NoError(t, err, "Init should not return an error when registration succeeds")
59+
api.AssertExpectations(t)
60+
}
61+
62+
// TestInit_Failure tests the Init method when the API registration fails.
63+
func TestInit_Failure(t *testing.T) {
64+
api := &MockAPIService{}
65+
p := NewPprof(api)
66+
67+
// Set up mock expectation
68+
expectedErr := assert.AnError
69+
api.On("Register", "/debug/pprof/", "pprof", "Profiling Data", mock.AnythingOfType("*http.ServeMux")).Return(expectedErr)
70+
71+
err := p.Init(context.Background())
72+
assert.Error(t, err, "Init should return an error when registration fails")
73+
assert.Equal(t, expectedErr, err, "Init should return the expected error")
74+
api.AssertExpectations(t)
75+
}
76+
77+
// TestPprofHandlers tests the handlers function to ensure it registers the correct pprof endpoints.
78+
func TestPprofHandlers(t *testing.T) {
79+
handler := handlers()
80+
mux, ok := handler.(*http.ServeMux)
81+
assert.True(t, ok, "handlers should return an http.ServeMux")
82+
83+
// Test cases for each pprof endpoint
84+
tests := []struct {
85+
path string
86+
handlerFunc http.HandlerFunc
87+
}{
88+
{"/debug/pprof/", pprof.Index},
89+
{"/debug/pprof/cmdline", pprof.Cmdline},
90+
{"/debug/pprof/profile", pprof.Profile},
91+
{"/debug/pprof/symbol", pprof.Symbol},
92+
{"/debug/pprof/trace", pprof.Trace},
93+
}
94+
95+
for _, tt := range tests {
96+
t.Run(tt.path, func(t *testing.T) {
97+
req := httptest.NewRequest("GET", tt.path, nil)
98+
rr := httptest.NewRecorder()
99+
mux.ServeHTTP(rr, req)
100+
assert.NotEqual(t, http.StatusNotFound, rr.Code, "Handler for %s should be registered", tt.path)
101+
})
102+
}
103+
}

internal/server/server.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,7 @@ func (s *APIServer) Shutdown() error {
159159
}
160160

161161
func (s *APIServer) Register(endpoint, summary, description string, handler http.Handler) error {
162+
s.logger.Debug("Endpoint Registered", "endppoint", endpoint)
162163
s.mux.Handle(endpoint, handler)
163164
s.endpointDescription += fmt.Sprintf("<li> <a href=\"%s\"> %s </a> %s </li>\n", endpoint, summary, description)
164165
return nil

0 commit comments

Comments
 (0)