commands: Improve server startup to make tests less flaky

Do this by announcing/listen on the local address before we start the server.

1 files changed, 39 insertions, 31 deletions

diff --git a/commands/server.go b/commands/server.go
index 4dd9116c7..bb6a4e15d 100644
--- a/commands/server.go
+++ b/commands/server.go
@@ -48,7 +48,7 @@ import (
 
 type serverCmd struct {
 	// Can be used to stop the server. Useful in tests
-	stop <-chan bool
+	stop chan bool
 
 	disableLiveReload bool
 	navigateToChanged bool
@@ -70,7 +70,7 @@ func (b *commandsBuilder) newServerCmd() *serverCmd {
 	return b.newServerCmdSignaled(nil)
 }
 
-func (b *commandsBuilder) newServerCmdSignaled(stop <-chan bool) *serverCmd {
+func (b *commandsBuilder) newServerCmdSignaled(stop chan bool) *serverCmd {
 	cc := &serverCmd{stop: stop}
 
 	cc.baseBuilderCmd = b.newBuilderCmd(&cobra.Command{
@@ -89,7 +89,13 @@ By default hugo will also watch your files for any changes you make and
 automatically rebuild the site. It will then live reload any open browser pages
 and push the latest content to them. As most Hugo sites are built in a fraction
 of a second, you will be able to save and see your changes nearly instantly.`,
-		RunE: cc.server,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			err := cc.server(cmd, args)
+			if err != nil && cc.stop != nil {
+				cc.stop <- true
+			}
+			return err
+		},
 	})
 
 	cc.cmd.Flags().IntVarP(&cc.serverPort, "port", "p", 1313, "port on which the server will listen")
@@ -130,8 +136,6 @@ func (f noDirFile) Readdir(count int) ([]os.FileInfo, error) {
 	return nil, nil
 }
 
-var serverPorts []int
-
 func (sc *serverCmd) server(cmd *cobra.Command, args []string) error {
 	// If a Destination is provided via flag write to disk
 	destination, _ := cmd.Flags().GetString("destination")
@@ -166,22 +170,21 @@ func (sc *serverCmd) server(cmd *cobra.Command, args []string) error {
 
 		// We can only do this once.
 		serverCfgInit.Do(func() {
-			serverPorts = make([]int, 1)
+			c.serverPorts = make([]serverPortListener, 1)
 
 			if c.languages.IsMultihost() {
 				if !sc.serverAppend {
 					rerr = newSystemError("--appendPort=false not supported when in multihost mode")
 				}
-				serverPorts = make([]int, len(c.languages))
+				c.serverPorts = make([]serverPortListener, len(c.languages))
 			}
 
 			currentServerPort := sc.serverPort
 
-			for i := 0; i < len(serverPorts); i++ {
+			for i := 0; i < len(c.serverPorts); i++ {
 				l, err := net.Listen("tcp", net.JoinHostPort(sc.serverInterface, strconv.Itoa(currentServerPort)))
 				if err == nil {
-					l.Close()
-					serverPorts[i] = currentServerPort
+					c.serverPorts[i] = serverPortListener{ln: l, p: currentServerPort}
 				} else {
 					if i == 0 && sc.cmd.Flags().Changed("port") {
 						// port set explicitly by user -- he/she probably meant it!
@@ -189,15 +192,15 @@ func (sc *serverCmd) server(cmd *cobra.Command, args []string) error {
 						return
 					}
 					c.logger.Println("port", sc.serverPort, "already in use, attempting to use an available port")
-					sp, err := helpers.FindAvailablePort()
+					l, sp, err := helpers.TCPListen()
 					if err != nil {
 						rerr = newSystemError("Unable to find alternative port to use:", err)
 						return
 					}
-					serverPorts[i] = sp.Port
+					c.serverPorts[i] = serverPortListener{ln: l, p: sp.Port}
 				}
 
-				currentServerPort = serverPorts[i] + 1
+				currentServerPort = c.serverPorts[i].p + 1
 			}
 		})
 
@@ -205,22 +208,20 @@ func (sc *serverCmd) server(cmd *cobra.Command, args []string) error {
 			return
 		}
 
-		c.serverPorts = serverPorts
-
 		c.Set("port", sc.serverPort)
 		if sc.liveReloadPort != -1 {
 			c.Set("liveReloadPort", sc.liveReloadPort)
 		} else {
-			c.Set("liveReloadPort", serverPorts[0])
+			c.Set("liveReloadPort", c.serverPorts[0].p)
 		}
 
 		isMultiHost := c.languages.IsMultihost()
 		for i, language := range c.languages {
 			var serverPort int
 			if isMultiHost {
-				serverPort = serverPorts[i]
+				serverPort = c.serverPorts[i].p
 			} else {
-				serverPort = serverPorts[0]
+				serverPort = c.serverPorts[0].p
 			}
 
 			baseURL, err := sc.fixURL(language, sc.baseURL, serverPort)
@@ -320,10 +321,11 @@ func (f *fileServer) rewriteRequest(r *http.Request, toPath string) *http.Reques
 	return r2
 }
 
-func (f *fileServer) createEndpoint(i int) (*http.ServeMux, string, string, error) {
+func (f *fileServer) createEndpoint(i int) (*http.ServeMux, net.Listener, string, string, error) {
 	baseURL := f.baseURLs[i]
 	root := f.roots[i]
-	port := f.c.serverPorts[i]
+	port := f.c.serverPorts[i].p
+	listener := f.c.serverPorts[i].ln
 
 	publishDir := f.c.Cfg.GetString("publishDir")
 
@@ -353,7 +355,7 @@ func (f *fileServer) createEndpoint(i int) (*http.ServeMux, string, string, erro
 	// We're only interested in the path
 	u, err := url.Parse(baseURL)
 	if err != nil {
-		return nil, "", "", errors.Wrap(err, "Invalid baseURL")
+		return nil, nil, "", "", errors.Wrap(err, "Invalid baseURL")
 	}
 
 	decorate := func(h http.Handler) http.Handler {
@@ -459,7 +461,7 @@ func (f *fileServer) createEndpoint(i int) (*http.ServeMux, string, string, erro
 
 	endpoint := net.JoinHostPort(f.s.serverInterface, strconv.Itoa(port))
 
-	return mu, u.String(), endpoint, nil
+	return mu, listener, u.String(), endpoint, nil
 }
 
 var logErrorRe = regexp.MustCompile(`(?s)ERROR \d{4}/\d{2}/\d{2} \d{2}:\d{2}:\d{2} `)
@@ -514,8 +516,10 @@ func (c *commandeer) serve(s *serverCmd) error {
 	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
 	var servers []*http.Server
 
+	wg1, ctx := errgroup.WithContext(context.Background())
+
 	for i := range baseURLs {
-		mu, serverURL, endpoint, err := srv.createEndpoint(i)
+		mu, listener, serverURL, endpoint, err := srv.createEndpoint(i)
 		srv := &http.Server{
 			Addr:    endpoint,
 			Handler: mu,
@@ -532,13 +536,13 @@ func (c *commandeer) serve(s *serverCmd) error {
 			mu.HandleFunc(u.Path+"/livereload", livereload.Handler)
 		}
 		jww.FEEDBACK.Printf("Web Server is available at %s (bind address %s)\n", serverURL, s.serverInterface)
-		go func() {
-			err = srv.ListenAndServe()
+		wg1.Go(func() error {
+			err = srv.Serve(listener)
 			if err != nil && err != http.ErrServerClosed {
-				c.logger.Errorf("Error: %s\n", err.Error())
-				os.Exit(1)
+				return err
 			}
-		}()
+			return nil
+		})
 	}
 
 	jww.FEEDBACK.Println("Press Ctrl+C to stop")
@@ -556,15 +560,19 @@ func (c *commandeer) serve(s *serverCmd) error {
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-	wg, ctx := errgroup.WithContext(ctx)
+	wg2, ctx := errgroup.WithContext(ctx)
 	for _, srv := range servers {
 		srv := srv
-		wg.Go(func() error {
+		wg2.Go(func() error {
 			return srv.Shutdown(ctx)
 		})
 	}
 
-	return wg.Wait()
+	err1, err2 := wg1.Wait(), wg2.Wait()
+	if err1 != nil {
+		return err1
+	}
+	return err2
 }
 
 // fixURL massages the baseURL into a form needed for serving