Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Expose profiling endpoints #3370

Merged
merged 1 commit into from
Nov 8, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 30 additions & 4 deletions api/server.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
// Package api contains the REST API implementation for k6.
// It also registers the services endpoints like pprof
package api

import (
"context"
"fmt"
"net/http"
_ "net/http/pprof" //nolint:gosec // Register pprof handlers
"time"

"github.com/sirupsen/logrus"
Expand All @@ -15,18 +18,41 @@ import (
"go.k6.io/k6/metrics/engine"
)

func newHandler(cs *v1.ControlSurface) http.Handler {
func newHandler(cs *v1.ControlSurface, profilingEnabled bool) http.Handler {
mux := http.NewServeMux()
mux.Handle("/v1/", v1.NewHandler(cs))
mux.Handle("/ping", handlePing(cs.RunState.Logger))
mux.Handle("/", handlePing(cs.RunState.Logger))

injectProfilerHandler(mux, profilingEnabled)

return mux
}

func injectProfilerHandler(mux *http.ServeMux, profilingEnabled bool) {
var handler http.Handler

handler = http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
rw.Header().Add("Content-Type", "text/plain; charset=utf-8")
_, _ = rw.Write([]byte("To enable profiling, please run k6 with the --profiling-enabled flag"))
})

if profilingEnabled {
handler = http.DefaultServeMux
}

mux.Handle("/debug/pprof/", handler)
}

// GetServer returns a http.Server instance that can serve k6's REST API.
func GetServer(
runCtx context.Context, addr string, runState *lib.TestRunState,
samples chan metrics.SampleContainer, me *engine.MetricsEngine, es *execution.Scheduler,
runCtx context.Context,
addr string,
profilingEnabled bool,
runState *lib.TestRunState,
samples chan metrics.SampleContainer,
me *engine.MetricsEngine,
es *execution.Scheduler,
) *http.Server {
// TODO: reduce the control surface as much as possible? For example, if
// we refactor the Runner API, we won't need to send the Samples channel.
Expand All @@ -38,7 +64,7 @@ func GetServer(
RunState: runState,
}

mux := withLoggingHandler(runState.Logger, newHandler(cs))
mux := withLoggingHandler(runState.Logger, newHandler(cs, profilingEnabled))
return &http.Server{Addr: addr, Handler: mux, ReadHeaderTimeout: 10 * time.Second}
}

Expand Down
6 changes: 6 additions & 0 deletions cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,12 @@ func rootCmdPersistentFlagSet(gs *state.GlobalState) *pflag.FlagSet {
flags.BoolVarP(&gs.Flags.Verbose, "verbose", "v", gs.DefaultFlags.Verbose, "enable verbose logging")
flags.BoolVarP(&gs.Flags.Quiet, "quiet", "q", gs.DefaultFlags.Quiet, "disable progress updates")
flags.StringVarP(&gs.Flags.Address, "address", "a", gs.DefaultFlags.Address, "address for the REST API server")
flags.BoolVar(
&gs.Flags.ProfilingEnabled,
"profiling-enabled",
gs.DefaultFlags.ProfilingEnabled,
"enable profiling (pprof) endpoints, k6's REST API should be enabled as well",
)

return flags
}
Expand Down
12 changes: 11 additions & 1 deletion cmd/run.go
Original file line number Diff line number Diff line change
Expand Up @@ -270,10 +270,20 @@ func (c *cmdRun) run(cmd *cobra.Command, args []string) (err error) {
srvCtx, srvCancel := context.WithCancel(globalCtx)
defer srvCancel()

srv := api.GetServer(runCtx, c.gs.Flags.Address, testRunState, samples, metricsEngine, execScheduler)
srv := api.GetServer(
runCtx,
c.gs.Flags.Address, c.gs.Flags.ProfilingEnabled,
testRunState,
samples,
metricsEngine,
execScheduler,
)
go func() {
defer apiWG.Done()
logger.Debugf("Starting the REST API server on %s", c.gs.Flags.Address)
if c.gs.Flags.ProfilingEnabled {
logger.Debugf("Profiling exposed on http://%s/debug/pprof/", c.gs.Flags.Address)
}
if aerr := srv.ListenAndServe(); aerr != nil && !errors.Is(aerr, http.ErrServerClosed) {
// Only exit k6 if the user has explicitly set the REST API address
if cmd.Flags().Lookup("address").Changed {
Expand Down
22 changes: 12 additions & 10 deletions cmd/state/state.go
Original file line number Diff line number Diff line change
Expand Up @@ -134,21 +134,23 @@ func NewGlobalState(ctx context.Context) *GlobalState {

// GlobalFlags contains global config values that apply for all k6 sub-commands.
type GlobalFlags struct {
ConfigFilePath string
Quiet bool
NoColor bool
Address string
LogOutput string
LogFormat string
Verbose bool
ConfigFilePath string
Quiet bool
NoColor bool
Address string
ProfilingEnabled bool
LogOutput string
LogFormat string
Verbose bool
}

// GetDefaultFlags returns the default global flags.
func GetDefaultFlags(homeDir string) GlobalFlags {
return GlobalFlags{
Address: "localhost:6565",
ConfigFilePath: filepath.Join(homeDir, "loadimpact", "k6", defaultConfigFileName),
LogOutput: "stderr",
Address: "localhost:6565",
ProfilingEnabled: false,
ConfigFilePath: filepath.Join(homeDir, "loadimpact", "k6", defaultConfigFileName),
LogOutput: "stderr",
}
}

Expand Down
4 changes: 4 additions & 0 deletions cmd/ui.go
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,10 @@ func printExecutionDescription(
}

fmt.Fprintf(buf, " output: %s\n", valueColor.Sprint(strings.Join(outputDescriptions, ", ")))
if gs.Flags.ProfilingEnabled && gs.Flags.Address != "" {
fmt.Fprintf(buf, " profiling: %s\n", valueColor.Sprintf("http://%s/debug/pprof/", gs.Flags.Address))
}

fmt.Fprintf(buf, "\n")

maxDuration, _ := lib.GetEndOffset(execPlan)
Expand Down