diff --git a/util/limiter/limiter.go b/util/limiter/limiter.go index 8896f8604..6345b35bb 100644 --- a/util/limiter/limiter.go +++ b/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, "") + for _, line := range dump { + if onlyLimited && line.Tokens > 0 { + continue + } + kStr := html.EscapeString(fmt.Sprint(line.Key)) + format := "" + if !onlyLimited && line.Tokens <= 0 { + // Make limited entries stand out when showing + // limited+non-limited together + format = "" + } + fmt.Fprintf(w, format, kStr, line.Tokens) + } + io.WriteString(w, "
KeyTokens
%s%d
%s%d
") +} + +// 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 +} diff --git a/util/limiter/limiter_test.go b/util/limiter/limiter_test.go index fdf0d6b7d..098a9ece1 100644 --- a/util/limiter/limiter_test.go +++ b/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{ + "", + "", + "", + "", + "", + "
KeyTokens
qux9
bar6
foo-2
", + }, "") + 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{ + "", + "", + "", + "
KeyTokens
foo-2
", + }, "") + 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{ + "", + "", + "", + "", + "", + "
KeyTokens
qux10
bar9
foo1
", + }, "") + 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++ { diff --git a/util/lru/lru.go b/util/lru/lru.go index db77859a4..b8fe511eb 100644 --- a/util/lru/lru.go +++ b/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, "") + c.ForEach(func(k K, v V) { + kStr := html.EscapeString(fmt.Sprint(k)) + vStr := html.EscapeString(fmt.Sprint(v)) + fmt.Fprintf(w, "", kStr, vStr) + }) + io.WriteString(w, "
KeyValue
%s%v
") +} diff --git a/util/lru/lru_test.go b/util/lru/lru_test.go index cdc169f57..0620688ce 100644 --- a/util/lru/lru_test.go +++ b/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{ + "", + "", + "", + "", + "", + "
KeyValue
4wat
3qux
2bar
", + }, "") + + 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