From 36efc508176b79e035332b95d414048f6ed46fa7 Mon Sep 17 00:00:00 2001
From: Tom DNetto <tom@tailscale.com>
Date: Tue, 6 Feb 2024 16:18:33 -0800
Subject: [PATCH] tsweb: implementing bucketed statistics for started/finished
 counts

Signed-off-by: Tom DNetto <tom@tailscale.com>
Updates: corp#17075
---
 tsweb/tsweb.go      | 63 +++++++++++++++++++++++++++++++++++++++++++++
 tsweb/tsweb_test.go | 28 ++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/tsweb/tsweb.go b/tsweb/tsweb.go
index c9d1c3a43..5ce45226a 100644
--- a/tsweb/tsweb.go
+++ b/tsweb/tsweb.go
@@ -18,6 +18,7 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strconv"
 	"strings"
 	"sync"
@@ -175,6 +176,52 @@ type ReturnHandler interface {
 	ServeHTTPReturn(http.ResponseWriter, *http.Request) error
 }
 
+// BucketedStatsOptions describes tsweb handler options surrounding
+// the generation of metrics, grouped into buckets.
+type BucketedStatsOptions struct {
+	// Bucket returns which bucket the given request is in.
+	// If nil, [NormalizedPath] is used to compute the bucket.
+	Bucket func(req *http.Request) string
+
+	// If non-nil, Started maintains a counter of all requests which
+	// have begun processing.
+	Started *expvar.Map
+
+	// If non-nil, Finished maintains a counter of all requests which
+	// have finished processing (that is, the HTTP handler has returned).
+	Finished *expvar.Map
+}
+
+var (
+	hexSequenceRegex = regexp.MustCompile("[a-fA-F0-9]{9,}")
+)
+
+// NormalizedPath returns the given path with any query parameters
+// removed, and any hex strings of 9 or more characters replaced
+// with an ellipsis.
+func NormalizedPath(p string) string {
+	// Fastpath: No hex sequences in there we might have to trim.
+	// Avoids allocating.
+	if hexSequenceRegex.FindStringIndex(p) == nil {
+		b, _, _ := strings.Cut(p, "?")
+		return b
+	}
+
+	// If we got here, there's at least one hex sequences we need to
+	// replace with an ellipsis.
+	replaced := hexSequenceRegex.ReplaceAllString(p, "…")
+	b, _, _ := strings.Cut(replaced, "?")
+	return b
+}
+
+func (o *BucketedStatsOptions) bucketForRequest(r *http.Request) string {
+	if o.Bucket != nil {
+		return o.Bucket(r)
+	}
+
+	return NormalizedPath(r.URL.Path)
+}
+
 type HandlerOptions struct {
 	QuietLoggingIfSuccessful bool // if set, do not log successfully handled HTTP requests (200 and 304 status codes)
 	Logf                     logger.Logf
@@ -189,6 +236,10 @@ type HandlerOptions struct {
 	// The keys are HTTP numeric response codes e.g. 200, 404, ...
 	StatusCodeCountersFull *expvar.Map
 
+	// If non-nil, BucketedStats computes and exposes statistics
+	// for each bucket based on the contained parameters.
+	BucketedStats *BucketedStatsOptions
+
 	// OnError is called if the handler returned a HTTPError. This
 	// is intended to be used to present pretty error pages if
 	// the user agent is determined to be a browser.
@@ -250,6 +301,14 @@ func (h retHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		RequestID:  RequestIDFromContext(r.Context()),
 	}
 
+	var bucket string
+	if bs := h.opts.BucketedStats; bs != nil {
+		bucket = bs.bucketForRequest(r)
+		if bs.Started != nil {
+			bs.Started.Add(bucket, 1)
+		}
+	}
+
 	lw := &loggingResponseWriter{ResponseWriter: w, logf: h.opts.Logf}
 	err := h.rh.ServeHTTPReturn(lw, r)
 
@@ -332,6 +391,10 @@ func (h retHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
+	if bs := h.opts.BucketedStats; bs != nil && bs.Finished != nil {
+		bs.Finished.Add(bucket, 1)
+	}
+
 	if !h.opts.QuietLoggingIfSuccessful || (msg.Code != http.StatusOK && msg.Code != http.StatusNotModified) {
 		h.opts.Logf("%s", msg)
 	}
diff --git a/tsweb/tsweb_test.go b/tsweb/tsweb_test.go
index 9e7005cb0..3a97191fb 100644
--- a/tsweb/tsweb_test.go
+++ b/tsweb/tsweb_test.go
@@ -11,12 +11,14 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"tailscale.com/tstest"
+	"tailscale.com/util/must"
 	"tailscale.com/util/vizerror"
 )
 
@@ -668,3 +670,29 @@ func TestCleanRedirectURL(t *testing.T) {
 		}
 	}
 }
+
+func TestBucket(t *testing.T) {
+	tcs := []struct {
+		path string
+		want string
+	}{
+		{"/map", "/map"},
+		{"/key?v=63", "/key"},
+		{"/map/a87e865a9d1c7", "/map/…"},
+		{"/machine/37fc1acb57f256b69b0d76749d814d91c68b241057c6b127fee3df37e4af111e", "/machine/…"},
+		{"/machine/37fc1acb57f256b69b0d76749d814d91c68b241057c6b127fee3df37e4af111e/map", "/machine/…/map"},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.path, func(t *testing.T) {
+			o := BucketedStatsOptions{}
+			bucket := (&o).bucketForRequest(&http.Request{
+				URL: must.Get(url.Parse(tc.path)),
+			})
+
+			if bucket != tc.want {
+				t.Errorf("bucket for %q was %q, want %q", tc.path, bucket, tc.want)
+			}
+		})
+	}
+}