diff --git a/api/server.go b/api/server.go index dabfea0567d..caed31b692f 100644 --- a/api/server.go +++ b/api/server.go @@ -3,6 +3,7 @@ package api import ( "fmt" "net/http" + "time" "github.com/sirupsen/logrus" @@ -19,11 +20,10 @@ func newHandler(logger logrus.FieldLogger) http.Handler { return mux } -// ListenAndServe is analogous to the stdlib one but also takes a core.Engine and logrus.FieldLogger -func ListenAndServe(addr string, engine *core.Engine, logger logrus.FieldLogger) error { - mux := newHandler(logger) - - return http.ListenAndServe(addr, withEngine(engine, newLogger(logger, mux))) +// GetServer returns a http.Server instance that can serve k6's REST API. +func GetServer(addr string, engine *core.Engine, logger logrus.FieldLogger) *http.Server { + mux := withEngine(engine, newLogger(logger, newHandler(logger))) + return &http.Server{Addr: addr, Handler: mux, ReadHeaderTimeout: 10 * time.Second} } type wrappedResponseWriter struct { diff --git a/cmd/run.go b/cmd/run.go index 5e2f55a2478..ab9c5f80d64 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -117,10 +117,20 @@ func (c *cmdRun) run(cmd *cobra.Command, args []string) error { // Spin up the REST API server, if not disabled. if c.gs.flags.address != "" { initBar.Modify(pb.WithConstProgress(0, "Init API server")) + + apiWG := &sync.WaitGroup{} + apiWG.Add(2) + defer apiWG.Wait() + + srvCtx, srvCancel := context.WithCancel(globalCtx) + defer srvCancel() + + // TODO: send the ExecutionState and MetricsEngine instead of the Engine + srv := api.GetServer(c.gs.flags.address, engine, logger) go func() { + defer apiWG.Done() logger.Debugf("Starting the REST API server on %s", c.gs.flags.address) - // TODO: send the ExecutionState and MetricsEngine instead of the Engine - if aerr := api.ListenAndServe(c.gs.flags.address, engine, logger); aerr != nil { + 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 { logger.WithError(aerr).Error("Error from API server") @@ -130,6 +140,13 @@ func (c *cmdRun) run(cmd *cobra.Command, args []string) error { } } }() + go func() { + defer apiWG.Done() + <-srvCtx.Done() + if aerr := srv.Close(); aerr != nil { + logger.WithError(aerr).Debug("REST API server did not shut down correctly") + } + }() } // We do this here so we can get any output URLs below.