util/expvarx: add a time and concurrency limiting expvar.Func wrapper

expvarx.SafeFunc wraps an expvar.Func with a time limit. On reaching the
time limit, calls to Value return nil, and no new concurrent calls to
the underlying expvar.Func will be started until the call completes.

Updates tailscale/corp#16999
Signed-off-by: James Tucker <james@tailscale.com>

util/expvarx/expvarx.go

@@ -0,0 +1,89 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+// Package expvarx provides some extensions to the [expvar] package.
+package expvarx
+
+import (
+	"encoding/json"
+	"expvar"
+	"sync"
+	"time"
+
+	"tailscale.com/types/lazy"
+)
+
+// SafeFunc is a wrapper around [expvar.Func] that guards against unbounded call
+// time and ensures that only a single call is in progress at any given time.
+type SafeFunc struct {
+	f      expvar.Func
+	limit  time.Duration
+	onSlow func(time.Duration, any)
+
+	mu       sync.Mutex
+	inflight *lazy.SyncValue[any]
+}
+
+// NewSafeFunc returns a new SafeFunc that wraps f.
+// If f takes longer than limit to execute then Value calls return nil.
+// If onSlow is non-nil, it is called when f takes longer than limit to execute.
+// onSlow is called with the duration of the slow call and the final computed
+// value.
+func NewSafeFunc(f expvar.Func, limit time.Duration, onSlow func(time.Duration, any)) *SafeFunc {
+	return &SafeFunc{f: f, limit: limit, onSlow: onSlow}
+}
+
+// Value acts similarly to [expvar.Func.Value], but if the underlying function
+// takes longer than the configured limit, all callers will receive nil until
+// the underlying operation completes. On completion of the underlying
+// operation, the onSlow callback is called if set.
+func (s *SafeFunc) Value() any {
+	s.mu.Lock()
+
+	if s.inflight == nil {
+		s.inflight = new(lazy.SyncValue[any])
+	}
+	var inflight = s.inflight
+	s.mu.Unlock()
+
+	// inflight ensures that only a single work routine is spawned at any given
+	// time, but if the routine takes too long inflight is populated with a nil
+	// result. The long running computed value is lost forever.
+	return inflight.Get(func() any {
+		start := time.Now()
+		result := make(chan any, 1)
+
+		// work is spawned in routine so that the caller can timeout.
+		go func() {
+			// Allow new work to be started after this work completes
+			defer func() {
+				s.mu.Lock()
+				s.inflight = nil
+				s.mu.Unlock()
+
+			}()
+
+			v := s.f.Value()
+			result <- v
+		}()
+
+		select {
+		case v := <-result:
+			return v
+		case <-time.After(s.limit):
+			if s.onSlow != nil {
+				go func() {
+					s.onSlow(time.Since(start), <-result)
+				}()
+			}
+			return nil
+		}
+	})
+}
+
+// String implements stringer in the same pattern as [expvar.Func], calling
+// Value and serializing the result as JSON, ignoring errors.
+func (s *SafeFunc) String() string {
+	v, _ := json.Marshal(s.Value())
+	return string(v)
+}