util/lru, util/limiter: add debug helper to dump state as HTML

For use in tsweb debug handlers, so that we can easily inspect cache
and limiter state when troubleshooting.

Updates tailscale/corp#3601

Signed-off-by: David Anderson <[email protected]>

util/limiter/limiter.go

@@ -4,6 +4,9 @@
 package limiter
 
 import (
+	"fmt"
+	"html"
+	"io"
 	"sync"
 	"time"
 
@@ -147,3 +150,53 @@ func (l *Limiter[K]) tokensForTest(key K) (int64, bool) {
 	}
 	return 0, false
 }
+
+// DumpHTML writes the state of the limiter to the given writer,
+// formatted as an HTML table. If onlyLimited is true, the output only
+// lists keys that are currently being limited.
+//
+// DumpHTML blocks other callers of the limiter while it collects the
+// state for dumping. It should not be called on large limiters
+// involved in hot codepaths.
+func (l *Limiter[K]) DumpHTML(w io.Writer, onlyLimited bool) {
+	l.dumpHTML(w, onlyLimited, time.Now())
+}
+
+func (l *Limiter[K]) dumpHTML(w io.Writer, onlyLimited bool, now time.Time) {
+	dump := l.collectDump(now)
+	io.WriteString(w, "<table><tr><th>Key</th><th>Tokens</th></tr>")
+	for _, line := range dump {
+		if onlyLimited && line.Tokens > 0 {
+			continue
+		}
+		kStr := html.EscapeString(fmt.Sprint(line.Key))
+		format := "<tr><td>%s</td><td>%d</td></tr>"
+		if !onlyLimited && line.Tokens <= 0 {
+			// Make limited entries stand out when showing
+			// limited+non-limited together
+			format = "<tr><td>%s</td><td><b>%d</b></td></tr>"
+		}
+		fmt.Fprintf(w, format, kStr, line.Tokens)
+	}
+	io.WriteString(w, "</table>")
+}
+
+// collectDump grabs a copy of the limiter state needed by DumpHTML.
+func (l *Limiter[K]) collectDump(now time.Time) []dumpEntry[K] {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	ret := make([]dumpEntry[K], 0, l.cache.Len())
+	l.cache.ForEach(func(k K, v *bucket) {
+		l.updateBucketLocked(v, now) // so stats are accurate
+		ret = append(ret, dumpEntry[K]{k, v.cur})
+	})
+	return ret
+}
+
+// dumpEntry is the per-key information that DumpHTML needs to print
+// limiter state.
+type dumpEntry[K comparable] struct {
+	Key    K
+	Tokens int64
+}

util/limiter/limiter_test.go

@@ -4,8 +4,12 @@
 package limiter
 
 import (
+	"bytes"
+	"strings"
 	"testing"
 	"time"
+
+	"github.com/google/go-cmp/cmp"
 )
 
 const testRefillInterval = time.Second
@@ -113,6 +117,64 @@ func TestLimiterOverdraft(t *testing.T) {
 	hasTokens(t, l, "foo", -1)
 }
 
+func TestDumpHTML(t *testing.T) {
+	l := &Limiter[string]{
+		Size:           3,
+		Max:            10,
+		Overdraft:      10,
+		RefillInterval: testRefillInterval,
+	}
+
+	now := time.Now().Truncate(testRefillInterval).Add(time.Millisecond)
+	allowed(t, l, "foo", 10, now)
+	denied(t, l, "foo", 2, now)
+	allowed(t, l, "bar", 4, now)
+	allowed(t, l, "qux", 1, now)
+
+	var out bytes.Buffer
+	l.DumpHTML(&out, false)
+	want := strings.Join([]string{
+		"<table>",
+		"<tr><th>Key</th><th>Tokens</th></tr>",
+		"<tr><td>qux</td><td>9</td></tr>",
+		"<tr><td>bar</td><td>6</td></tr>",
+		"<tr><td>foo</td><td><b>-2</b></td></tr>",
+		"</table>",
+	}, "")
+	if diff := cmp.Diff(out.String(), want); diff != "" {
+		t.Fatalf("wrong DumpHTML output (-got+want):\n%s", diff)
+	}
+
+	out.Reset()
+	l.DumpHTML(&out, true)
+	want = strings.Join([]string{
+		"<table>",
+		"<tr><th>Key</th><th>Tokens</th></tr>",
+		"<tr><td>foo</td><td>-2</td></tr>",
+		"</table>",
+	}, "")
+	if diff := cmp.Diff(out.String(), want); diff != "" {
+		t.Fatalf("wrong DumpHTML output (-got+want):\n%s", diff)
+	}
+
+	// Check that DumpHTML updates tokens even if the key wasn't hit
+	// organically.
+	now = now.Add(3 * time.Second)
+	out.Reset()
+	l.dumpHTML(&out, false, now)
+	want = strings.Join([]string{
+		"<table>",
+		"<tr><th>Key</th><th>Tokens</th></tr>",
+		"<tr><td>qux</td><td>10</td></tr>",
+		"<tr><td>bar</td><td>9</td></tr>",
+		"<tr><td>foo</td><td>1</td></tr>",
+		"</table>",
+	}, "")
+	if diff := cmp.Diff(out.String(), want); diff != "" {
+		t.Fatalf("wrong DumpHTML output (-got+want):\n%s", diff)
+	}
+}
+
 func allowed(t *testing.T, l *Limiter[string], key string, count int, now time.Time) {
 	t.Helper()
 	for i := 0; i < count; i++ {

util/lru/lru.go

@@ -4,6 +4,12 @@
 // Package lru contains a typed Least-Recently-Used cache.
 package lru
 
+import (
+	"fmt"
+	"html"
+	"io"
+)
+
 // Cache is container type keyed by K, storing V, optionally evicting the least
 // recently used items if a maximum size is exceeded.
 //
@@ -171,3 +177,31 @@ func (c *Cache[K, V]) deleteElement(ent *entry[K, V]) {
 	}
 	delete(c.lookup, ent.key)
 }
+
+// ForEach calls fn for each entry in the cache, from most recently
+// used to least recently used.
+func (c *Cache[K, V]) ForEach(fn func(K, V)) {
+	if c.head == nil {
+		return
+	}
+	cur := c.head
+	for {
+		fn(cur.key, cur.value)
+		cur = cur.next
+		if cur == c.head {
+			return
+		}
+	}
+}
+
+// DumpHTML writes the state of the cache to the given writer,
+// formatted as an HTML table.
+func (c *Cache[K, V]) DumpHTML(w io.Writer) {
+	io.WriteString(w, "<table><tr><th>Key</th><th>Value</th></tr>")
+	c.ForEach(func(k K, v V) {
+		kStr := html.EscapeString(fmt.Sprint(k))
+		vStr := html.EscapeString(fmt.Sprint(v))
+		fmt.Fprintf(w, "<tr><td>%s</td><td>%v</td></tr>", kStr, vStr)
+	})
+	io.WriteString(w, "</table>")
+}

util/lru/lru_test.go

@@ -4,8 +4,12 @@
 package lru
 
 import (
+	"bytes"
 	"math/rand"
+	"strings"
 	"testing"
+
+	"github.com/google/go-cmp/cmp"
 )
 
 func TestLRU(t *testing.T) {
@@ -44,6 +48,31 @@ func TestLRU(t *testing.T) {
 	}
 }
 
+func TestDumpHTML(t *testing.T) {
+	c := Cache[int, string]{MaxEntries: 3}
+
+	c.Set(1, "foo")
+	c.Set(2, "bar")
+	c.Set(3, "qux")
+	c.Set(4, "wat")
+
+	var out bytes.Buffer
+	c.DumpHTML(&out)
+
+	want := strings.Join([]string{
+		"<table>",
+		"<tr><th>Key</th><th>Value</th></tr>",
+		"<tr><td>4</td><td>wat</td></tr>",
+		"<tr><td>3</td><td>qux</td></tr>",
+		"<tr><td>2</td><td>bar</td></tr>",
+		"</table>",
+	}, "")
+
+	if diff := cmp.Diff(out.String(), want); diff != "" {
+		t.Fatalf("wrong DumpHTML output (-got+want):\n%s", diff)
+	}
+}
+
 func BenchmarkLRU(b *testing.B) {
 	const lruSize = 10
 	const maxval = 15 // 33% more keys than the LRU can hold