Implement a central state for the web server

Signed-off-by: Florent Poinsard <[email protected]>

go/cmd/root.go

@@ -24,32 +24,26 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/vitessio/vt/go/web"
+	"github.com/vitessio/vt/go/web/state"
 )
 
+//nolint:gochecknoglobals // the state is protected using mutexes
+var wstate *state.State
+
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
 	// rootCmd represents the base command when called without any subcommands
-	var port int64
-	webserverStarted := false
-	ch := make(chan int, 2)
 	root := &cobra.Command{
 		Use:   "vt",
 		Short: "Utils tools for testing, running and benchmarking Vitess.",
-		RunE: func(_ *cobra.Command, _ []string) error {
-			if port > 0 {
-				if webserverStarted {
-					return nil
-				}
-				webserverStarted = true
-				go startWebServer(port, ch)
-				<-ch
-			}
-			return nil
+		CompletionOptions: cobra.CompletionOptions{
+			DisableDefaultCmd: true,
 		},
 	}
+
+	var port int64
 	root.PersistentFlags().Int64VarP(&port, "port", "p", 8080, "Port to run the web server on")
-	root.CompletionOptions.HiddenDefaultCmd = true
 
 	root.AddCommand(summarizeCmd(&port))
 	root.AddCommand(testerCmd())
@@ -63,16 +57,20 @@ func Execute() {
 		panic(err)
 	}
 
-	if !webserverStarted && port > 0 {
-		webserverStarted = true
-		go startWebServer(port, ch)
+	// Start the web server for all commands no matter what
+	wstate = state.NewState(port)
+	ch := make(chan int, 1)
+	if port > 0 {
+		wstate.SetStarted(true)
+		go startWebServer(ch)
+		if !wstate.WaitUntilAvailable(10 * time.Second) {
+			fmt.Println("Timed out waiting for server to start")
+			os.Exit(1)
+		}
 	} else {
 		ch <- 1
 	}
 
-	// FIXME: add sync b/w webserver and root command, for now just add a wait to make sure webserver is running
-	time.Sleep(2 * time.Second)
-
 	err := root.Execute()
 	if err != nil {
 		fmt.Printf("Error: %v\n", err)
@@ -81,13 +79,15 @@ func Execute() {
 	<-ch
 }
 
-func startWebServer(port int64, ch chan int) {
-	if port > 0 && port != 8080 {
-		panic("(FIXME: make port configurable) Port is not 8080")
+func startWebServer(ch chan int) {
+	err := web.Run(wstate)
+	if err != nil {
+		panic(err)
 	}
-	web.Run(port)
-	if os.WriteFile("/dev/stderr", []byte("Web server is running, use Ctrl-C to exit\n"), 0o600) != nil {
-		panic("Failed to write to /dev/stderr")
+
+	_, err = fmt.Fprint(os.Stderr, "Web server is running, use Ctrl-C to exit\n")
+	if err != nil {
+		panic(err)
 	}
 	ch <- 1
 }

go/cmd/summarize.go

@@ -22,8 +22,10 @@ import (
 	"github.com/vitessio/vt/go/summarize"
 )
 
-func summarizeCmd(_ *int64) *cobra.Command {
-	var cfg summarize.Config
+func summarizeCmd(port *int64) *cobra.Command {
+	cfg := summarize.Config{
+		WState: wstate,
+	}
 
 	cmd := &cobra.Command{
 		Use:     "summarize old_file.json [new_file.json]",
@@ -33,14 +35,14 @@ func summarizeCmd(_ *int64) *cobra.Command {
 		Args:    cobra.RangeArgs(1, 2),
 		Run: func(_ *cobra.Command, args []string) {
 			cfg.Files = args
+			cfg.Port = *port
 			summarize.Run(&cfg)
 		},
 	}
 
 	cmd.Flags().StringVar(&cfg.HotMetric, "hot-metric", "total-time", "Metric to determine hot queries (options: usage-count, total-rows-examined, avg-rows-examined, avg-time, total-time)")
 	cmd.Flags().BoolVar(&cfg.ShowGraph, "graph", false, "Show the query graph in the browser")
 	cmd.Flags().StringVar(&cfg.OutputFormat, "format", "html", "Output format (options: html, markdown)")
-	cmd.Flags().IntVar(&cfg.Port, "port", 8000, "Port to listen on for the webserver")
 
 	return cmd
 }

go/summarize/summarize.go

@@ -32,6 +32,7 @@ import (
 	"golang.org/x/term"
 
 	"github.com/vitessio/vt/go/data"
+	"github.com/vitessio/vt/go/web/state"
 )
 
 type (
@@ -41,9 +42,11 @@ type (
 
 		OutputFormat string
 
-		Port int
+		Port int64
 
 		ShowGraph bool
+
+		WState *state.State
 	}
 
 	traceSummary struct {
@@ -118,7 +121,7 @@ func exitIfError(err error) {
 	os.Exit(1)
 }
 
-func printSummary(hotMetric string, workers []summaryWorker, outputFormat string, port int) (*Summary, error) {
+func printSummary(hotMetric string, workers []summaryWorker, outputFormat string, port int64) (*Summary, error) {
 	s, err := NewSummary(hotMetric)
 	if err != nil {
 		return nil, err

go/web/state/state.go

@@ -0,0 +1,72 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package state
+
+import (
+	"fmt"
+	"net/http"
+	"sync"
+	"time"
+)
+
+type State struct {
+	mu sync.Mutex
+
+	port    int64
+	started bool
+}
+
+func NewState(port int64) *State {
+	return &State{
+		port: port,
+	}
+}
+
+func (s *State) SetStarted(v bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.started = v
+}
+
+func (s *State) Started() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.started
+}
+
+func (s *State) GetPort() int64 {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.port
+}
+
+func (s *State) WaitUntilAvailable(timeout time.Duration) bool {
+	for {
+		select {
+		case <-time.After(timeout):
+			return false
+		case <-time.After(100 * time.Millisecond):
+			r, err := http.Get(fmt.Sprintf("http://localhost:%d/", s.port))
+			if err != nil {
+				continue
+			}
+			if r.StatusCode == http.StatusOK {
+				return true
+			}
+		}
+	}
+}

go/web/web.go

@@ -6,14 +6,14 @@ import (
 	"fmt"
 	"html/template"
 	"io"
-	"log"
 	"net/http"
 	"os"
 	"time"
 
 	"github.com/gin-gonic/gin"
 
 	"github.com/vitessio/vt/go/summarize"
+	"github.com/vitessio/vt/go/web/state"
 )
 
 func RenderFileToGin(fileName string, data any, c *gin.Context) {
@@ -84,7 +84,7 @@ func addFuncMap(r *gin.Engine) {
 	r.SetFuncMap(getFuncMap())
 }
 
-func Run(port int64) {
+func Run(s *state.State) error {
 	gin.SetMode(gin.ReleaseMode)
 	gin.DefaultWriter = io.Discard // Disable logging
 	r := gin.Default()
@@ -125,10 +125,12 @@ func Run(port int64) {
 		RenderFileToGin("summarize.html", &summarizeOutput, c)
 	})
 
-	if os.WriteFile("/dev/stderr", []byte(fmt.Sprintf("Starting web server on http://localhost:%d\n", port)), 0o600) != nil {
-		panic("Failed to write to /dev/stderr")
+	if _, err := fmt.Fprintf(os.Stderr, "Starting web server on http://localhost:%d\n", s.GetPort()); err != nil {
+		return err
 	}
-	if err := r.Run(fmt.Sprintf(":%d", port)); err != nil {
-		log.Fatalf("Failed to start server: %v", err)
+
+	if err := r.Run(fmt.Sprintf(":%d", s.GetPort())); err != nil {
+		return err
 	}
+	return nil
 }