Ensure the REST API server is properly stopped before k6 run finishes

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 {

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.