Skip to content

Commit

Permalink
Expose & display profiling endpoints
Browse files Browse the repository at this point in the history
  • Loading branch information
olegbespalov committed Nov 8, 2023
1 parent e44a08f commit a43edef
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 15 deletions.
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

0 comments on commit a43edef

Please sign in to comment.