diff --git a/.vscode/launch.json b/.vscode/launch.json index 08dba90..74a4190 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -96,6 +96,7 @@ "SCAFFOLD_DOCKER_IMAGE": "cortex-axon-agent:local", "PORT" : "7399", "BUILTIN_PLUGIN_DIR": "${workspaceFolder}/agent/server/snykbroker/plugins", + "ENABLE_PPROF": "true" } }, diff --git a/agent/config/config.go b/agent/config/config.go index c3e7516..3edc9de 100644 --- a/agent/config/config.go +++ b/agent/config/config.go @@ -72,6 +72,7 @@ type AgentConfig struct { WebhookServerPort int SnykBrokerPort int EnableApiProxy bool + EnablePprof bool FailWaitTime time.Duration AutoRegisterFrequency time.Duration VerboseOutput bool @@ -109,6 +110,9 @@ func (ac AgentConfig) Print() { } else { fmt.Println("\tAPI Proxy: Disabled") } + if ac.EnablePprof { + fmt.Println("\tpprof: Enabled at /pprof") + } fmt.Printf("\tFast fail time: %v\n", ac.FailWaitTime) } @@ -180,6 +184,11 @@ func NewAgentEnvConfig() AgentConfig { dryRun = dryRunEnv == "true" || dryRunEnv == "1" } + enablePprof := false + if pprofEnv := os.Getenv("ENABLE_PPROF"); pprofEnv != "" { + enablePprof = pprofEnv == "true" || pprofEnv == "1" + } + dequeueWaitTime := 1 * time.Second if dequeueWaitTimeEnv := os.Getenv("DEQUEUE_WAIT_TIME"); dequeueWaitTimeEnv != "" { dwt, err := time.ParseDuration(dequeueWaitTimeEnv) @@ -243,6 +252,7 @@ func NewAgentEnvConfig() AgentConfig { WebhookServerPort: WebhookServerPort, SnykBrokerPort: snykBrokerPort, EnableApiProxy: true, + EnablePprof: enablePprof, FailWaitTime: time.Second * 2, PluginDirs: []string{"./plugins"}, AutoRegisterFrequency: reregisterFrequency, diff --git a/agent/server/http/pprof_handler.go b/agent/server/http/pprof_handler.go new file mode 100644 index 0000000..8cfbcdb --- /dev/null +++ b/agent/server/http/pprof_handler.go @@ -0,0 +1,51 @@ +package http + +import ( + "io" + "net/http" + "net/http/pprof" + + "github.com/cortexapps/axon/config" + "github.com/gorilla/mux" + "go.uber.org/zap" +) + +const PprofPathRoot = "/pprof" + +type pprofHandler struct { + io.Closer + config config.AgentConfig + logger *zap.Logger +} + +func NewPprofHandler(config config.AgentConfig, logger *zap.Logger) RegisterableHandler { + return &pprofHandler{ + config: config, + logger: logger, + } +} + +func (h *pprofHandler) RegisterRoutes(m *mux.Router) error { + sub := m.PathPrefix(PprofPathRoot).Subrouter() + + // The index page generates relative links to the named profiles, so they + // resolve under /pprof/. Individual profiles must be registered explicitly + // since pprof.Index only auto-dispatches under the hardcoded /debug/pprof/. + sub.HandleFunc("/", pprof.Index) + sub.HandleFunc("/cmdline", pprof.Cmdline) + sub.HandleFunc("/profile", pprof.Profile) + sub.HandleFunc("/symbol", pprof.Symbol) + sub.HandleFunc("/trace", pprof.Trace) + + for _, name := range []string{"heap", "goroutine", "allocs", "block", "mutex", "threadcreate"} { + sub.Handle("/"+name, pprof.Handler(name)) + } + + h.logger.Info("pprof profiling endpoint enabled", zap.String("path", PprofPathRoot)) + + return nil +} + +func (h *pprofHandler) ServeHTTP(_ http.ResponseWriter, _ *http.Request) { + panic("ServeHTTP should not be called directly") +} diff --git a/agent/server/http/pprof_handler_test.go b/agent/server/http/pprof_handler_test.go new file mode 100644 index 0000000..9a968bd --- /dev/null +++ b/agent/server/http/pprof_handler_test.go @@ -0,0 +1,49 @@ +package http + +import ( + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/cortexapps/axon/config" + "github.com/gorilla/mux" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +func newPprofTestServer(t *testing.T) *httptest.Server { + t.Helper() + handler := NewPprofHandler(config.AgentConfig{}, zap.NewNop()) + router := mux.NewRouter() + require.NoError(t, handler.RegisterRoutes(router)) + ts := httptest.NewServer(router) + t.Cleanup(ts.Close) + return ts +} + +func TestPprofIndex(t *testing.T) { + ts := newPprofTestServer(t) + + resp, err := http.Get(ts.URL + "/pprof/") + require.NoError(t, err) + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Contains(t, string(body), "goroutine") +} + +func TestPprofHeapProfile(t *testing.T) { + ts := newPprofTestServer(t) + + resp, err := http.Get(ts.URL + "/pprof/heap") + require.NoError(t, err) + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.NotEmpty(t, body) +} diff --git a/agent/server/main_http_server.go b/agent/server/main_http_server.go index 12a44b6..8abc4a1 100644 --- a/agent/server/main_http_server.go +++ b/agent/server/main_http_server.go @@ -54,5 +54,10 @@ func NewMainHttpServer(p MainHttpServerParams) cortexHttp.Server { httpServer.RegisterHandler(metricsHandler) } + if config.EnablePprof { + pprofHandler := cortexHttp.NewPprofHandler(config, p.Logger) + httpServer.RegisterHandler(pprofHandler) + } + return httpServer } diff --git a/docker/Dockerfile b/docker/Dockerfile index 0831538..5f15029 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -27,7 +27,7 @@ WORKDIR /agent # when the scheduled Trivy scan flags OS-package CVEs whose fixes are already # in the archive — the cache is just serving a stale layer. Leaves the rest of # the build (Go, npm, snyk-broker clone) hitting cache as normal. -ARG APT_CACHE_BUST=2026-05-27 +ARG APT_CACHE_BUST=2026-05-29 RUN echo "apt cache bust: $APT_CACHE_BUST" \ && apt-get update && apt-get upgrade -y && apt-get install -y \ protobuf-compiler git python3 python3-venv wget build-essential openssl jq