servenv: Move away from using default HTTP muxer

Up until now we've been using the implicit default HTTP mux. The problem
of using this mux is that it works implicitly also for profiling
endpoints. So the moment that `net/http/pprof` is imported, those
endpoints get added.

We want to be able to make enabling the profiling endpoints optional,
but there's no way to do that with the default mux except by recompiling
and removing the import.

By moving everything away from the default mux to an explicit mux for
the servenv, we can make this configurable in the near future. It still
gets added to the default mux, but we don't use that anymore so it
doesn't get exposed.

The changes here now explicitly add the profiling endpoints so that we
can make that optional in a follow up.

Signed-off-by: Dirkjan Bussink <d.bussink@gmail.com>

go/cmd/vtgate/index.go

@@ -18,14 +18,16 @@ package main
 
 import (
 	"net/http"
+
+	"vitess.io/vitess/go/vt/servenv"
 )
 
 // This is a separate file so it can be selectively included/excluded from
 // builds to opt in/out of the redirect.
 
 func init() {
 	// Anything unrecognized gets redirected to the status page.
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		http.Redirect(w, r, "/debug/status", http.StatusFound)
 	})
 }

go/cmd/vtorc/index.go

@@ -18,11 +18,13 @@ package main
 
 import (
 	"net/http"
+
+	"vitess.io/vitess/go/vt/servenv"
 )
 
 func init() {
 	// Anything unrecognized gets redirected to the status page.
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		http.Redirect(w, r, "/debug/status", http.StatusFound)
 	})
 }

go/cmd/vttablet/index.go

@@ -18,14 +18,16 @@ package main
 
 import (
 	"net/http"
+
+	"vitess.io/vitess/go/vt/servenv"
 )
 
 // This is a separate file so it can be selectively included/excluded from
 // builds to opt in/out of the redirect.
 
 func init() {
 	// Anything unrecognized gets redirected to the status page.
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		http.Redirect(w, r, "/debug/status", http.StatusFound)
 	})
 }

go/stats/opentsdb/opentsdb.go

@@ -113,7 +113,7 @@ func InitWithoutServenv(prefix string) {
 
 	stats.RegisterPushBackend("opentsdb", backend)
 
-	http.HandleFunc("/debug/opentsdb", func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc("/debug/opentsdb", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json; charset=utf-8")
 		dataPoints := (*backend).getDataPoints()
 		sort.Sort(byMetric(dataPoints))

go/stats/prometheusbackend/prometheusbackend.go

@@ -18,14 +18,14 @@ package prometheusbackend
 
 import (
 	"expvar"
-	"net/http"
 	"strings"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 
 	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/servenv"
 )
 
 // PromBackend implements PullBackend using Prometheus as the backing metrics storage.
@@ -39,7 +39,7 @@ var (
 
 // Init initializes the Prometheus be with the given namespace.
 func Init(namespace string) {
-	http.Handle("/metrics", promhttp.Handler())
+	servenv.HTTPHandle("/metrics", promhttp.Handler())
 	be.namespace = namespace
 	stats.Register(be.publishPrometheusMetric)
 }

go/streamlog/streamlog.go

@@ -183,7 +183,7 @@ func (logger *StreamLogger[T]) Name() string {
 // ServeLogs registers the URL on which messages will be broadcast.
 // It is safe to register multiple URLs for the same StreamLogger.
 func (logger *StreamLogger[T]) ServeLogs(url string, logf LogFormatter) {
-	http.HandleFunc(url, func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc(url, func(w http.ResponseWriter, r *http.Request) {
 		if err := acl.CheckAccessHTTP(r, acl.DEBUGGING); err != nil {
 			acl.SendError(w, err)
 			return

go/streamlog/streamlog_flaky_test.go

@@ -28,6 +28,8 @@ import (
 	"syscall"
 	"testing"
 	"time"
+
+	"vitess.io/vitess/go/vt/servenv"
 )
 
 type logMessage struct {
@@ -51,7 +53,12 @@ func TestHTTP(t *testing.T) {
 	defer l.Close()
 	addr := l.Addr().String()
 
-	go http.Serve(l, nil)
+	go func() {
+		err := servenv.HTTPServe(l)
+		if err != nil {
+			t.Errorf("http serve returned unexpected error: %v", err)
+		}
+	}()
 
 	logger := New[*logMessage]("logger", 1)
 	logger.ServeLogs("/log", testLogf)

go/vt/discovery/healthcheck.go

@@ -348,7 +348,7 @@ func NewHealthCheck(ctx context.Context, retryDelay, healthCheckTimeout time.Dur
 
 	hc.topoWatchers = topoWatchers
 	healthcheckOnce.Do(func() {
-		http.Handle("/debug/gateway", hc)
+		servenv.HTTPHandle("/debug/gateway", hc)
 	})
 
 	// start the topo watches here

go/vt/servenv/exporter.go

@@ -105,6 +105,10 @@ type Exporter struct {
 	mu          sync.Mutex
 }
 
+func init() {
+	HTTPHandle("/debug/vars", expvar.Handler())
+}
+
 // NewExporter creates a new Exporter with name as namespace.
 // label is the name of the additional dimension for the stats vars.
 func NewExporter(name, label string) *Exporter {
@@ -153,12 +157,12 @@ func (e *Exporter) URLPrefix() string {
 
 // HandleFunc sets or overwrites the handler for url. If Exporter has a name,
 // url remapped from /path to /name/path. If name is empty, the request
-// is passed through to http.HandleFunc.
+// is passed through to HTTPHandleFunc.
 func (e *Exporter) HandleFunc(url string, f func(w http.ResponseWriter, r *http.Request)) {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 	if e.name == "" {
-		http.HandleFunc(url, f)
+		HTTPHandleFunc(url, f)
 		return
 	}
 
@@ -169,7 +173,7 @@ func (e *Exporter) HandleFunc(url string, f func(w http.ResponseWriter, r *http.
 	hf := &handleFunc{f: f}
 	e.handleFuncs[url] = hf
 
-	http.HandleFunc(e.URLPrefix()+url, func(w http.ResponseWriter, r *http.Request) {
+	HTTPHandleFunc(e.URLPrefix()+url, func(w http.ResponseWriter, r *http.Request) {
 		if f := hf.Get(); f != nil {
 			f(w, r)
 		}

go/vt/servenv/exporter_test.go

@@ -43,7 +43,12 @@ func TestHandleFunc(t *testing.T) {
 	}
 	defer listener.Close()
 	port := listener.Addr().(*net.TCPAddr).Port
-	go http.Serve(listener, nil)
+	go func() {
+		err := HTTPServe(listener)
+		if err != nil {
+			t.Errorf("HTTPServe returned: %v", err)
+		}
+	}()
 
 	ebd := NewExporter("", "")
 	ebd.HandleFunc("/path", func(w http.ResponseWriter, r *http.Request) {

go/vt/servenv/flushlogs.go

@@ -25,7 +25,7 @@ import (
 
 func init() {
 	OnInit(func() {
-		http.HandleFunc("/debug/flushlogs", func(w http.ResponseWriter, r *http.Request) {
+		HTTPHandleFunc("/debug/flushlogs", func(w http.ResponseWriter, r *http.Request) {
 			logutil.Flush()
 			fmt.Fprint(w, "flushed")
 		})

go/vt/servenv/http.go

@@ -0,0 +1,39 @@
+package servenv
+
+import (
+	"errors"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/http/pprof"
+)
+
+var mux = http.NewServeMux()
+
+func HTTPHandle(pattern string, handler http.Handler) {
+	mux.Handle(pattern, handler)
+}
+
+func HTTPHandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
+	mux.HandleFunc(pattern, handler)
+}
+
+func HTTPServe(l net.Listener) error {
+	err := http.Serve(l, mux)
+	if errors.Is(err, http.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
+		return nil
+	}
+	return err
+}
+
+func HTTPTestServer() *httptest.Server {
+	return httptest.NewServer(mux)
+}
+
+func HTTPRegisterProfile() {
+	HTTPHandleFunc("/debug/pprof/", pprof.Index)
+	HTTPHandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	HTTPHandleFunc("/debug/pprof/profile", pprof.Profile)
+	HTTPHandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	HTTPHandleFunc("/debug/pprof/trace", pprof.Trace)
+}

go/vt/servenv/liveness.go

@@ -29,7 +29,7 @@ import (
 // further behind on its backlog.
 
 func init() {
-	http.HandleFunc("/debug/liveness", func(rw http.ResponseWriter, r *http.Request) {
+	HTTPHandleFunc("/debug/liveness", func(rw http.ResponseWriter, r *http.Request) {
 		// Do nothing. Return success immediately.
 	})
 }

go/vt/servenv/liveness_test.go

@@ -19,12 +19,11 @@ package servenv
 import (
 	"io"
 	"net/http"
-	"net/http/httptest"
 	"testing"
 )
 
 func TestLivenessHandler(t *testing.T) {
-	server := httptest.NewServer(nil)
+	server := HTTPTestServer()
 	defer server.Close()
 
 	resp, err := http.Get(server.URL + "/debug/liveness")

go/vt/servenv/pprof.go

@@ -342,4 +342,5 @@ func init() {
 		fs.StringSliceVar(&pprofFlag, "pprof", pprofFlag, "enable profiling")
 	})
 	OnInit(pprofInit)
+	OnInit(HTTPRegisterProfile)
 }

go/vt/servenv/run.go

@@ -19,7 +19,6 @@ package servenv
 import (
 	"fmt"
 	"net"
-	"net/http"
 	"net/url"
 	"os"
 	"os/signal"
@@ -49,7 +48,12 @@ func Run(port int) {
 	if err != nil {
 		log.Exit(err)
 	}
-	go http.Serve(l, nil)
+	go func() {
+		err := HTTPServe(l)
+		if err != nil {
+			log.Errorf("http serve returned unexpected error: %v", err)
+		}
+	}()
 
 	ExitChan = make(chan os.Signal, 1)
 	signal.Notify(ExitChan, syscall.SIGTERM, syscall.SIGINT)

go/vt/servenv/servenv.go

@@ -29,8 +29,6 @@ limitations under the License.
 package servenv
 
 import (
-	// register the HTTP handlers for profiling
-	_ "net/http/pprof"
 	"net/url"
 	"os"
 	"os/signal"

go/vt/servenv/status.go

@@ -167,12 +167,12 @@ func newStatusPage(name string) *statusPage {
 	}
 	sp.tmpl = template.Must(sp.reparse(nil))
 	if name == "" {
-		http.HandleFunc(StatusURLPath(), sp.statusHandler)
+		HTTPHandleFunc(StatusURLPath(), sp.statusHandler)
 		// Debug profiles are only supported for the top level status page.
 		registerDebugBlockProfileRate()
 		registerDebugMutexProfileFraction()
 	} else {
-		http.HandleFunc("/"+name+StatusURLPath(), sp.statusHandler)
+		HTTPHandleFunc("/"+name+StatusURLPath(), sp.statusHandler)
 	}
 	return sp
 }
@@ -277,7 +277,7 @@ func (sp *statusPage) reparse(sections []section) (*template.Template, error) {
 
 // Toggle the block profile rate to/from 100%, unless specific rate is passed in
 func registerDebugBlockProfileRate() {
-	http.HandleFunc("/debug/blockprofilerate", func(w http.ResponseWriter, r *http.Request) {
+	HTTPHandleFunc("/debug/blockprofilerate", func(w http.ResponseWriter, r *http.Request) {
 		if err := acl.CheckAccessHTTP(r, acl.DEBUGGING); err != nil {
 			acl.SendError(w, err)
 			return
@@ -307,7 +307,7 @@ func registerDebugBlockProfileRate() {
 
 // Toggle the mutex profiling fraction to/from 100%, unless specific fraction is passed in
 func registerDebugMutexProfileFraction() {
-	http.HandleFunc("/debug/mutexprofilefraction", func(w http.ResponseWriter, r *http.Request) {
+	HTTPHandleFunc("/debug/mutexprofilefraction", func(w http.ResponseWriter, r *http.Request) {
 		if err := acl.CheckAccessHTTP(r, acl.DEBUGGING); err != nil {
 			acl.SendError(w, err)
 			return

go/vt/servenv/status_test.go

@@ -19,7 +19,6 @@ package servenv
 import (
 	"io"
 	"net/http"
-	"net/http/httptest"
 	"regexp"
 	"strings"
 	"testing"
@@ -43,7 +42,7 @@ func init() {
 }
 
 func TestStatus(t *testing.T) {
-	server := httptest.NewServer(nil)
+	server := HTTPTestServer()
 	defer server.Close()
 
 	resp, err := http.Get(server.URL + StatusURLPath())
@@ -68,7 +67,7 @@ func TestStatus(t *testing.T) {
 }
 
 func TestNamedStatus(t *testing.T) {
-	server := httptest.NewServer(nil)
+	server := HTTPTestServer()
 	defer server.Close()
 
 	name := "test"

go/vt/throttler/demo/throttler_demo.go

@@ -302,7 +302,7 @@ func main() {
 	servenv.ParseFlags(flagSetName)
 
 	go servenv.RunDefault()
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		http.Redirect(w, r, "/throttlerz", http.StatusTemporaryRedirect)
 	})
 

go/vt/throttler/throttlerlogz.go

@@ -27,6 +27,7 @@ import (
 	"golang.org/x/exp/slices"
 
 	"vitess.io/vitess/go/vt/logz"
+	"vitess.io/vitess/go/vt/servenv"
 )
 
 const logHeaderHTML = `
@@ -101,7 +102,7 @@ var (
 )
 
 func init() {
-	http.HandleFunc("/throttlerlogz/", func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc("/throttlerlogz/", func(w http.ResponseWriter, r *http.Request) {
 		throttlerlogzHandler(w, r, GlobalManager)
 	})
 }

go/vt/throttler/throttlerz.go

@@ -24,6 +24,7 @@ import (
 	"golang.org/x/exp/slices"
 
 	"vitess.io/vitess/go/vt/log"
+	"vitess.io/vitess/go/vt/servenv"
 )
 
 const listHTML = `<!DOCTYPE html>
@@ -50,7 +51,7 @@ var (
 )
 
 func init() {
-	http.HandleFunc("/throttlerz/", func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc("/throttlerz/", func(w http.ResponseWriter, r *http.Request) {
 		throttlerzHandler(w, r, GlobalManager)
 	})
 }

go/vt/vtctld/api.go

@@ -122,7 +122,7 @@ func httpErrorf(w http.ResponseWriter, r *http.Request, format string, args ...a
 }
 
 func handleAPI(apiPath string, handlerFunc func(w http.ResponseWriter, r *http.Request) error) {
-	http.HandleFunc(apiPrefix+apiPath, func(w http.ResponseWriter, r *http.Request) {
+	servenv.HTTPHandleFunc(apiPrefix+apiPath, func(w http.ResponseWriter, r *http.Request) {
 		defer func() {
 			if x := recover(); x != nil {
 				httpErrorf(w, r, "uncaught panic: %v", x)