servenv: Move away from using default HTTP muxer

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,54 @@
+/*
+Copyright 2023 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 servenv
+
+import (
+	"errors"
+	"net"
+	"net/http"
+	"net/http/pprof"
+
+	"vitess.io/vitess/go/vt/servenv/internal/mux"
+)
+
+// HTTPHandle registers the given handler for the internal servenv mux.
+func HTTPHandle(pattern string, handler http.Handler) {
+	mux.Mux.Handle(pattern, handler)
+}
+
+// HTTPHandleFunc registers the given handler func for the internal servenv mux.
+func HTTPHandleFunc(pattern string, handler func(http.ResponseWriter, *http.Request)) {
+	mux.Mux.HandleFunc(pattern, handler)
+}
+
+// HTTPServe starts the HTTP server for the internal servenv mux on the listener.
+func HTTPServe(l net.Listener) error {
+	err := http.Serve(l, mux.Mux)
+	if errors.Is(err, http.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
+		return nil
+	}
+	return err
+}
+
+// HTTPRegisterProfile registers the default pprof HTTP endpoints with the internal servenv 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/internal/mux/mux.go

@@ -0,0 +1,21 @@
+/*
+Copyright 2023 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 mux
+
+import "net/http"
+
+var Mux = http.NewServeMux()

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,13 @@ package servenv
 import (
 	"io"
 	"net/http"
-	"net/http/httptest"
 	"testing"
+
+	"vitess.io/vitess/go/vt/servenv/testutils"
 )
 
 func TestLivenessHandler(t *testing.T) {
-	server := httptest.NewServer(nil)
+	server := testutils.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,13 +19,14 @@ package servenv
 import (
 	"io"
 	"net/http"
-	"net/http/httptest"
 	"regexp"
 	"strings"
 	"testing"
 
 	"github.com/google/safehtml/template"
 	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/vt/servenv/testutils"
 )
 
 func init() {
@@ -43,7 +44,7 @@ func init() {
 }
 
 func TestStatus(t *testing.T) {
-	server := httptest.NewServer(nil)
+	server := testutils.HTTPTestServer()
 	defer server.Close()
 
 	resp, err := http.Get(server.URL + StatusURLPath())
@@ -68,7 +69,7 @@ func TestStatus(t *testing.T) {
 }
 
 func TestNamedStatus(t *testing.T) {
-	server := httptest.NewServer(nil)
+	server := testutils.HTTPTestServer()
 	defer server.Close()
 
 	name := "test"