tailscale/util/expvarx/expvarx.go
James Tucker 0f3b2e7b86 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>
2024-01-24 17:47:16 -08:00

90 lines
2.4 KiB
Go

// 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)
}