diff --git a/main.go b/main.go index 5f3ba855a..618ff645e 100644 --- a/main.go +++ b/main.go @@ -46,6 +46,7 @@ import ( "github.com/mongodb/mongodb-kubernetes/mongodb-community-operator/pkg/util/envvar" "github.com/mongodb/mongodb-kubernetes/pkg/images" "github.com/mongodb/mongodb-kubernetes/pkg/multicluster" + "github.com/mongodb/mongodb-kubernetes/pkg/pprof" "github.com/mongodb/mongodb-kubernetes/pkg/telemetry" "github.com/mongodb/mongodb-kubernetes/pkg/util" "github.com/mongodb/mongodb-kubernetes/pkg/util/architectures" @@ -277,6 +278,16 @@ func main() { log.Info("Not running telemetry component!") } + pprofEnabledString := env.ReadOrDefault(util.OperatorPprofEnabledEnv, "") + if pprofEnabled, err := pprof.IsPprofEnabled(pprofEnabledString, getOperatorEnv()); err != nil { + log.Errorf("Unable to check if pprof is enabled: %s", err) + } else if pprofEnabled { + port := env.ReadIntOrDefault(util.OperatorPprofPortEnv, util.OperatorPprofDefaultPort) + if err := mgr.Add(pprof.NewRunnable(port, log)); err != nil { + log.Errorf("Unable to start pprof server: %s", err) + } + } + log.Info("Starting the Cmd.") // Start the Manager diff --git a/pkg/pprof/pprof.go b/pkg/pprof/pprof.go new file mode 100644 index 000000000..c66f4e13e --- /dev/null +++ b/pkg/pprof/pprof.go @@ -0,0 +1,82 @@ +package pprof + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/http/pprof" + "strconv" + "time" + + "go.uber.org/zap" + + "github.com/mongodb/mongodb-kubernetes/pkg/util" +) + +type Runnable struct { + port int + log *zap.SugaredLogger +} + +func NewRunnable(port int, log *zap.SugaredLogger) *Runnable { + return &Runnable{ + port: port, + log: log, + } +} + +func (p *Runnable) Start(ctx context.Context) error { + pprofAddress := fmt.Sprintf("localhost:%d", p.port) + + handler := http.NewServeMux() + handler.HandleFunc("GET /debug/pprof/", pprof.Index) + handler.HandleFunc("GET /debug/pprof/cmdline/", pprof.Cmdline) + handler.HandleFunc("GET /debug/pprof/profile/", pprof.Profile) + handler.HandleFunc("GET /debug/pprof/symbol/", pprof.Symbol) + handler.HandleFunc("GET /debug/pprof/trace/", pprof.Trace) + + server := &http.Server{ + Addr: pprofAddress, + ReadHeaderTimeout: 10 * time.Second, + Handler: handler, + } + + go func() { + p.log.Infof("Starting pprof server at %s", pprofAddress) + if err := server.ListenAndServe(); err != nil { + if !errors.Is(err, http.ErrServerClosed) { + p.log.Errorf("unable to start pprof server: %s", err.Error()) + } + } + p.log.Info("pprof server stopped") + }() + + go func() { + <-ctx.Done() + p.log.Info("Stopping pprof server") + if err := server.Shutdown(context.Background()); err != nil { + p.log.Errorf("unable to shutdown pprof server: %s", err.Error()) + } + }() + + return nil +} + +// IsPprofEnabled checks if pprof is enabled based on the MDB_OPERATOR_PPROF_ENABLED +// and OPERATOR_ENV environment variables. It returns true if: +// - MDB_OPERATOR_PPROF_ENABLED is set to true +// - OPERATOR_ENV is set to dev or local and MDB_OPERATOR_PPROF_ENABLED is not set +// Otherwise, it returns false. +func IsPprofEnabled(pprofEnabledString string, operatorEnv util.OperatorEnvironment) (bool, error) { + if pprofEnabledString != "" { + pprofEnabled, err := strconv.ParseBool(pprofEnabledString) + if err != nil { + return false, fmt.Errorf("unable to parse %s environment variable: %w", util.OperatorPprofEnabledEnv, err) + } + + return pprofEnabled, nil + } + + return operatorEnv == util.OperatorEnvironmentDev || operatorEnv == util.OperatorEnvironmentLocal, nil +} diff --git a/pkg/pprof/pprof_test.go b/pkg/pprof/pprof_test.go new file mode 100644 index 000000000..6d8c4c3b5 --- /dev/null +++ b/pkg/pprof/pprof_test.go @@ -0,0 +1,80 @@ +package pprof + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/mongodb/mongodb-kubernetes/pkg/util" +) + +func TestIsPprofEnabled(t *testing.T) { + tests := map[string]struct { + pprofEnabledString string + operatorEnv util.OperatorEnvironment + expected bool + expectedErrMsg string + }{ + "pprof enabled by default in dev": { + operatorEnv: util.OperatorEnvironmentDev, + expected: true, + }, + "pprof enabled by default in local": { + operatorEnv: util.OperatorEnvironmentLocal, + expected: true, + }, + "pprof disabled by default in prod": { + operatorEnv: util.OperatorEnvironmentProd, + expected: false, + }, + "pprof enabled in prod": { + pprofEnabledString: "true", + operatorEnv: util.OperatorEnvironmentProd, + expected: true, + }, + "pprof explicitly enabled in dev": { + pprofEnabledString: "true", + operatorEnv: util.OperatorEnvironmentDev, + expected: true, + }, + "pprof explicitly enabled in local": { + pprofEnabledString: "true", + operatorEnv: util.OperatorEnvironmentLocal, + expected: true, + }, + "pprof disabled in dev": { + pprofEnabledString: "false", + operatorEnv: util.OperatorEnvironmentDev, + expected: false, + }, + "pprof disabled in local": { + pprofEnabledString: "false", + operatorEnv: util.OperatorEnvironmentLocal, + expected: false, + }, + "pprof disabled explicitly in prod": { + pprofEnabledString: "false", + operatorEnv: util.OperatorEnvironmentProd, + expected: false, + }, + "pprof misconfigured": { + pprofEnabledString: "false11", + operatorEnv: util.OperatorEnvironmentProd, + expected: false, + expectedErrMsg: "unable to parse MDB_OPERATOR_PPROF_ENABLED environment variable: strconv.ParseBool: parsing \"false11\": invalid syntax", + }, + } + + for name, test := range tests { + t.Run(name, func(t *testing.T) { + result, err := IsPprofEnabled(test.pprofEnabledString, test.operatorEnv) + if test.expectedErrMsg != "" { + require.Error(t, err) + assert.Equal(t, test.expectedErrMsg, err.Error()) + } + + assert.Equal(t, test.expected, result) + }) + } +} diff --git a/pkg/util/constants.go b/pkg/util/constants.go index 86219ee05..a359838c0 100644 --- a/pkg/util/constants.go +++ b/pkg/util/constants.go @@ -154,6 +154,11 @@ const ( LDAP = "LDAP" MinimumScramSha256MdbVersion = "4.0.0" + // pprof variables + OperatorPprofEnabledEnv = "MDB_OPERATOR_PPROF_ENABLED" + OperatorPprofPortEnv = "MDB_OPERATOR_PPROF_PORT" + OperatorPprofDefaultPort = 10081 + // these were historically used and constituted a security issue—if set they should be changed InvalidKeyFileContents = "DUMMYFILE" InvalidAutomationAgentPassword = "D9XK2SfdR2obIevI9aKsYlVH" //nolint //Part of the algorithm