prometheus: Remove interface indirection, and output strings not bytes.

This is part of a series of changes to add metric charts in performance
benchmarks.

This change is meant to do three things:

- Remove the interface indirection from the Prometheus library, which is
  performance-critical due to its use in writing out profiling metrics
  (although the runsc metric server also benefits from this too).
- Use a `StringWriter`-like writer contract, to avoid needless casting
  between strings and bytes within the Prometheus library. The library only
  ever needs to deal with strings, so it is up to callers to do the
  conversion to bytes if they need to (which the runsc metric-server does).
- Avoid buffer allocations in the metric server when each snapshot is larger
  than the buffer size. Instead, buffers are saved and reused.

PiperOrigin-RevId: 630500004

pkg/prometheus/BUILD

@@ -11,6 +11,7 @@ go_library(
         "prometheus.go",
         "prometheus_verify.go",
     ],
+    stateify = False,
     visibility = ["//:sandbox"],
     deps = ["//pkg/metric:metric_go_proto"],
 )

pkg/prometheus/prometheus_test.go

@@ -1170,9 +1170,9 @@ func TestVerifier(t *testing.T) {
 	}
 }
 
-// shortWriter implements io.Writer but fails after a given number of bytes.
+// shortWriter implements io.StringWriter but fails after a given number of bytes.
 type shortWriter struct {
-	buf     bytes.Buffer
+	buf     strings.Builder
 	size    int
 	maxSize int
 }
@@ -1189,9 +1189,9 @@ func (s *shortWriter) String() string {
 	return s.buf.String()
 }
 
-// Write implements io.Writer.Write.
-func (s *shortWriter) Write(b []byte) (n int, err error) {
-	toWrite := len(b)
+// Write implements io.StringWriter.WriteString.
+func (s *shortWriter) WriteString(x string) (n int, err error) {
+	toWrite := len(x)
 	leftToWrite := s.maxSize - s.size
 	if leftToWrite < toWrite {
 		toWrite = leftToWrite
@@ -1199,9 +1199,9 @@ func (s *shortWriter) Write(b []byte) (n int, err error) {
 	if toWrite == 0 {
 		return 0, errors.New("writer out of capacity")
 	}
-	written, err := s.buf.Write(b[:toWrite])
+	written, err := s.buf.WriteString(x[:toWrite])
 	s.size += written
-	if written == len(b) {
+	if written == len(x) {
 		return written, err
 	}
 	return written, errors.New("short write")
@@ -1696,7 +1696,7 @@ func TestGroupSameNameMetrics(t *testing.T) {
 	// Make sure the data written does parse.
 	// We don't use this result here because the Prometheus library is more permissive than this test.
 	if _, err := (&expfmt.TextParser{}).TextToMetricFamilies(&buf); err != nil {
-		t.Fatalf("cannot parse data written from snapshots: %v", err)
+		t.Fatalf("cannot parse data written from snapshots: %v\nraw data:\n%s\n(end of raw data)", err, rawData)
 	}
 
 	// Verify that we see all metrics, and that each time we see a new one, it's one we haven't seen

runsc/metricserver/metricserver.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"io/ioutil"
 	"math/rand"
 	"net"
@@ -262,6 +261,10 @@ type metricServer struct {
 	// Used to efficiently reallocate a map of the right size during the next export.
 	lastMetricsWrittenSize atomicbitops.Uint32
 
+	// Pool of `prometheus.ReusableWriter`s. Used to avoid large buffer allocations for
+	// successive snapshots.
+	promWriterPool sync.Pool
+
 	// mu protects the fields below.
 	mu sync.Mutex
 
@@ -613,7 +616,7 @@ func queryMultiSandboxMetrics(ctx context.Context, loadedSandboxes []sandboxLoad
 }
 
 // serveMetrics serves metrics requests.
-func (m *metricServer) serveMetrics(w http.ResponseWriter, req *http.Request) httpResult {
+func (m *metricServer) serveMetrics(w *httpResponseWriter, req *http.Request) httpResult {
 	ctx, ctxCancel := context.WithTimeout(req.Context(), metricsExportTimeout)
 	defer ctxCancel()
 
@@ -738,10 +741,12 @@ func (m *metricServer) serveMetrics(w http.ResponseWriter, req *http.Request) ht
 	if metricsFilter != "" {
 		commentHeader = fmt.Sprintf("%s (filtered using regular expression: %q)", commentHeader, metricsFilter)
 	}
-	written, err := prometheus.Write(w, prometheus.ExportOptions{
+	promWriter := m.promWriterPool.Get().(*prometheus.ReusableWriter[*httpResponseWriter])
+	written, err := promWriter.Write(w, prometheus.ExportOptions{
 		CommentHeader:  commentHeader,
 		MetricsWritten: metricsWritten,
 	}, snapshotsToOptions)
+	m.promWriterPool.Put(promWriter)
 	if err != nil {
 		if written == 0 {
 			return httpResult{http.StatusServiceUnavailable, err}
@@ -761,7 +766,7 @@ func (m *metricServer) serveMetrics(w http.ResponseWriter, req *http.Request) ht
 // Returns a response prefixed by "runsc-metrics:OK" on success.
 // Clients can use this to assert that they are talking to the metrics server, as opposed to some
 // other random HTTP server.
-func (m *metricServer) serveHealthCheck(w http.ResponseWriter, req *http.Request) httpResult {
+func (m *metricServer) serveHealthCheck(w *httpResponseWriter, req *http.Request) httpResult {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	if m.shuttingDown {
@@ -775,18 +780,18 @@ func (m *metricServer) serveHealthCheck(w http.ResponseWriter, req *http.Request
 		return httpResult{http.StatusBadRequest, fmt.Errorf("this metric server is configured to serve root directory: %s", m.rootDir)}
 	}
 	w.WriteHeader(http.StatusOK)
-	io.WriteString(w, "runsc-metrics:OK")
+	w.WriteString("runsc-metrics:OK")
 	return httpOK
 }
 
 // servePID serves the PID of the metric server process.
-func (m *metricServer) servePID(w http.ResponseWriter, req *http.Request) httpResult {
+func (m *metricServer) servePID(w *httpResponseWriter, req *http.Request) httpResult {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	if m.shuttingDown {
 		return httpResult{http.StatusServiceUnavailable, errors.New("server is shutting down")}
 	}
-	io.WriteString(w, strconv.Itoa(m.pid))
+	w.WriteString(strconv.Itoa(m.pid))
 	return httpOK
 }
 
@@ -823,6 +828,11 @@ func (s *Server) Run(ctx context.Context) error {
 		pidFile:                s.PIDFile,
 		exposeProfileEndpoints: s.ExposeProfileEndpoints,
 		allowUnknownRoot:       s.AllowUnknownRoot,
+		promWriterPool: sync.Pool{
+			New: func() any {
+				return &prometheus.ReusableWriter[*httpResponseWriter]{}
+			},
+		},
 	}
 	conf := s.Config
 	if conf.MetricServer == "" {

runsc/metricserver/metricserver_http.go

@@ -16,7 +16,6 @@ package metricserver
 
 import (
 	"errors"
-	"fmt"
 	"net/http"
 	"runtime"
 	"strings"
@@ -38,7 +37,7 @@ type httpResult struct {
 var httpOK = httpResult{code: http.StatusOK}
 
 // serveIndex serves the index page.
-func (m *metricServer) serveIndex(w http.ResponseWriter, req *http.Request) httpResult {
+func (m *metricServer) serveIndex(w *httpResponseWriter, req *http.Request) httpResult {
 	if req.URL.Path != "/" {
 		if strings.HasPrefix(req.URL.Path, "/metrics?") {
 			// Prometheus's scrape_config.metrics_path takes in a query path and automatically encodes
@@ -52,15 +51,40 @@ func (m *metricServer) serveIndex(w http.ResponseWriter, req *http.Request) http
 		}
 		return httpResult{http.StatusNotFound, errors.New("path not found")}
 	}
-	fmt.Fprintf(w, "<html><head><title>runsc metrics</title></head><body>")
-	fmt.Fprintf(w, "<p>You have reached the runsc metrics server page!</p>")
-	fmt.Fprintf(w, `<p>To see actual metric data, head over to <a href="/metrics">/metrics</a>.</p>`)
-	fmt.Fprintf(w, "</body></html>")
+	w.WriteString("<html><head><title>runsc metrics</title></head><body>")
+	w.WriteString("<p>You have reached the runsc metrics server page!</p>")
+	w.WriteString(`<p>To see actual metric data, head over to <a href="/metrics">/metrics</a>.</p>`)
+	w.WriteString("</body></html>")
 	return httpOK
 }
 
+// httpResponseWriter is a ResponseWriter that also implements io.StringWriter.
+type httpResponseWriter struct {
+	resp http.ResponseWriter
+}
+
+// Header implements http.ResponseWriter.Header.
+func (w *httpResponseWriter) Header() http.Header {
+	return w.resp.Header()
+}
+
+// Write implements http.ResponseWriter.Write.
+func (w *httpResponseWriter) Write(b []byte) (int, error) {
+	return w.resp.Write(b)
+}
+
+// WriteHeader implements http.ResponseWriter.WriteHeader.
+func (w *httpResponseWriter) WriteHeader(code int) {
+	w.resp.WriteHeader(code)
+}
+
+// WriteString implements io.StringWriter.WriteString.
+func (w *httpResponseWriter) WriteString(s string) (int, error) {
+	return w.resp.Write([]byte(s))
+}
+
 // logRequest wraps an HTTP handler and adds logging to it.
-func logRequest(f func(w http.ResponseWriter, req *http.Request) httpResult) func(w http.ResponseWriter, req *http.Request) {
+func logRequest(f func(w *httpResponseWriter, req *http.Request) httpResult) func(w http.ResponseWriter, req *http.Request) {
 	return func(w http.ResponseWriter, req *http.Request) {
 		log.Infof("Request: %s %s", req.Method, req.URL.Path)
 		defer func() {
@@ -68,7 +92,7 @@ func logRequest(f func(w http.ResponseWriter, req *http.Request) httpResult) fun
 				log.Warningf("Request: %s %s: Panic:\n%v", req.Method, req.URL.Path, r)
 			}
 		}()
-		result := f(w, req)
+		result := f(&httpResponseWriter{resp: w}, req)
 		if result.err != nil {
 			http.Error(w, result.err.Error(), result.code)
 			log.Warningf("Request: %s %s: Failed with HTTP code %d: %v", req.Method, req.URL.Path, result.code, result.err)

runsc/metricserver/metricserver_profile.go

@@ -26,7 +26,7 @@ import (
 )
 
 // profileCPU returns a CPU profile over HTTP.
-func (m *metricServer) profileCPU(w http.ResponseWriter, req *http.Request) httpResult {
+func (m *metricServer) profileCPU(w *httpResponseWriter, req *http.Request) httpResult {
 	// Time to finish up profiling and flush out the results to the client.
 	const finishProfilingBuffer = 250 * time.Millisecond
 
@@ -66,7 +66,7 @@ func (m *metricServer) profileCPU(w http.ResponseWriter, req *http.Request) http
 }
 
 // profileHeap returns a heap profile over HTTP.
-func (m *metricServer) profileHeap(w http.ResponseWriter, req *http.Request) httpResult {
+func (m *metricServer) profileHeap(w *httpResponseWriter, req *http.Request) httpResult {
 	m.mu.Lock()
 	if m.shuttingDown {
 		m.mu.Unlock()