Add a simple rate limiter with a helper to use it in logging (#1915)

inside.go

@@ -10,6 +10,7 @@ import (
 	"github.com/slackhq/nebula/firewall"
 	"github.com/slackhq/nebula/header"
 	"github.com/slackhq/nebula/iputil"
+	"github.com/slackhq/nebula/logging"
 	"github.com/slackhq/nebula/noiseutil"
 	"github.com/slackhq/nebula/overlay/batch"
 	"github.com/slackhq/nebula/overlay/tio"
@@ -517,7 +518,9 @@ func (f *Interface) SendVia(via *HostInfo, relay *Relay, ad, nb, out []byte, noc
 
 	err = f.writers[q].WriteTo(toSend, via.GetRemote())
 	if err != nil {
-		via.logger(f.l).Info("Failed to WriteTo in sendVia", "error", err)
+		if l, ok := logging.RateLimited(f.l, f.udpRelayTxErrLimiter); ok {
+			via.logger(l).Info("Failed to WriteTo in sendVia", "error", err)
+		}
 	}
 }
 
@@ -582,18 +585,22 @@ func (f *Interface) sendNoMetrics(t header.MessageType, st header.MessageSubType
 	if remote.IsValid() {
 		err = f.writers[q].WriteTo(out, remote)
 		if err != nil {
-			hostinfo.logger(f.l).Error("Failed to write outgoing packet",
-				"error", err,
-				"udpAddr", remote,
-			)
+			if l, ok := logging.RateLimited(f.l, f.udpTxErrLimiter); ok {
+				hostinfo.logger(l).Error("Failed to write outgoing packet",
+					"error", err,
+					"udpAddr", remote,
+				)
+			}
 		}
 	} else if hr := hostinfo.GetRemote(); hr.IsValid() {
 		err = f.writers[q].WriteTo(out, hr)
 		if err != nil {
-			hostinfo.logger(f.l).Error("Failed to write outgoing packet",
-				"error", err,
-				"udpAddr", hr,
-			)
+			if l, ok := logging.RateLimited(f.l, f.udpTxErrLimiter); ok {
+				hostinfo.logger(l).Error("Failed to write outgoing packet",
+					"error", err,
+					"udpAddr", hr,
+				)
+			}
 		}
 	} else {
 		// Try to send via a relay

interface.go

@@ -24,6 +24,7 @@ import (
 	"github.com/slackhq/nebula/overlay"
 	"github.com/slackhq/nebula/overlay/batch"
 	"github.com/slackhq/nebula/overlay/tio"
+	"github.com/slackhq/nebula/ratelimit"
 	"github.com/slackhq/nebula/udp"
 )
 
@@ -124,6 +125,9 @@ type Interface struct {
 	batchers []*batch.MultiCoalescer
 	wg       sync.WaitGroup
 
+	udpTxErrLimiter      *ratelimit.Limiter // udp direct transmit error log limiter
+	udpRelayTxErrLimiter *ratelimit.Limiter // udp relay transmit error log limiter
+
 	// fatalErr holds the first unexpected reader error that caused shutdown.
 	// nil means "no fatal error" (yet)
 	fatalErr atomic.Pointer[error]
@@ -198,6 +202,15 @@ func NewInterface(ctx context.Context, c *InterfaceConfig) (*Interface, error) {
 		return nil, errors.New("no connection manager")
 	}
 
+	udpTxErrLimiter, err := ratelimit.New(1, time.Second)
+	if err != nil {
+		return nil, err
+	}
+	udpRelayTxErrLimiter, err := ratelimit.New(1, time.Second)
+	if err != nil {
+		return nil, err
+	}
+
 	if c.routines <= 1 {
 		c.PinThreads = false //pinning is not useful unless there's more than one tun reader
 	}
@@ -230,6 +243,8 @@ func NewInterface(ctx context.Context, c *InterfaceConfig) (*Interface, error) {
 		conntrackCacheTimeout: c.ConntrackCacheTimeout,
 		cpuAffinity:           c.CpuAffinity,
 		pinThreads:            c.PinThreads,
+		udpTxErrLimiter:       udpTxErrLimiter,
+		udpRelayTxErrLimiter:  udpRelayTxErrLimiter,
 
 		metricHandshakes: metrics.GetOrRegisterHistogram("handshakes", nil, metrics.NewExpDecaySample(1028, 0.015)),
 		metricTxDropped:  metrics.GetOrRegisterCounter("udp.tx.dropped", nil),

logging/ratelimit.go

@@ -0,0 +1,31 @@
+package logging
+
+import "log/slog"
+
+// RateLimiter is what RateLimited needs from a rate limiter, such as a *ratelimit.Limiter.
+// Allow reports whether this call may proceed, and how many calls were denied since the last one that did.
+type RateLimiter interface {
+	Allow() (denied uint64, ok bool)
+}
+
+// RateLimited checks lim before anything is built.
+// When allowed it returns l, with a "suppressed" count attached if lim denied calls since its last allowed one.
+// A nil lim allows everything.
+//
+//	if l, ok := logging.RateLimited(f.l, f.udpTxErrLimiter); ok {
+//		hostinfo.logger(l).Error("Failed to write outgoing packet", "error", err)
+//	}
+func RateLimited(l *slog.Logger, lim RateLimiter) (*slog.Logger, bool) {
+	if lim == nil {
+		return l, true
+	}
+
+	denied, ok := lim.Allow()
+	if !ok {
+		return nil, false
+	}
+	if denied > 0 {
+		return l.With("suppressed", denied), true
+	}
+	return l, true
+}

logging/ratelimit_test.go

@@ -0,0 +1,106 @@
+package logging
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"log/slog"
+	"net/netip"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/slackhq/nebula/ratelimit"
+)
+
+type fakeLimiter struct {
+	denied uint64
+	ok     bool
+}
+
+func (f fakeLimiter) Allow() (uint64, bool) { return f.denied, f.ok }
+
+func TestRateLimited(t *testing.T) {
+	var buf bytes.Buffer
+	base := NewLogger(&buf)
+
+	if l, ok := RateLimited(base, fakeLimiter{ok: false}); ok || l != nil {
+		t.Fatalf("denied call: ok=%v logger=%v", ok, l)
+	}
+
+	if l, ok := RateLimited(base, fakeLimiter{ok: true}); !ok || l != base {
+		t.Fatal("nothing denied should hand back the same logger")
+	}
+
+	if l, ok := RateLimited(base, nil); !ok || l != base {
+		t.Fatal("a nil limiter should allow with the same logger")
+	}
+
+	l, ok := RateLimited(base, fakeLimiter{denied: 7, ok: true})
+	if !ok {
+		t.Fatal("allowed call denied")
+	}
+	l.Info("hello", "k", "v")
+	if !strings.Contains(buf.String(), "suppressed=7") {
+		t.Fatalf("want suppressed=7, got %q", buf.String())
+	}
+}
+
+// The benchmarks compare a log line through slog directly against the same line behind RateLimited: logged with
+// nothing denied, denied, and logged after denials.
+
+var (
+	benchErr  = errors.New("sendto: no buffer space available")
+	benchAddr = netip.MustParseAddrPort("192.0.2.1:4242")
+)
+
+func benchLogger() *slog.Logger {
+	return NewLogger(io.Discard).With("localIndex", 1234, "remoteIndex", 5678)
+}
+
+func BenchmarkRateLimitedBaseline(b *testing.B) {
+	lg := benchLogger()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		lg.Error("Failed to write outgoing packet", "error", benchErr, "udpAddr", benchAddr)
+	}
+}
+
+func BenchmarkRateLimitedAllowed(b *testing.B) {
+	lg, lim := benchLogger(), mustLimiter(b, 1, 0)
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		if l, ok := RateLimited(lg, lim); ok {
+			l.Error("Failed to write outgoing packet", "error", benchErr, "udpAddr", benchAddr)
+		}
+	}
+}
+
+func BenchmarkRateLimitedDenied(b *testing.B) {
+	lg, lim := benchLogger(), mustLimiter(b, 1, time.Hour)
+	lim.Allow()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		if l, ok := RateLimited(lg, lim); ok {
+			l.Error("Failed to write outgoing packet", "error", benchErr, "udpAddr", benchAddr)
+		}
+	}
+}
+
+func BenchmarkRateLimitedAfterDenied(b *testing.B) {
+	lg, lim := benchLogger(), fakeLimiter{denied: 5, ok: true}
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		if l, ok := RateLimited(lg, lim); ok {
+			l.Error("Failed to write outgoing packet", "error", benchErr, "udpAddr", benchAddr)
+		}
+	}
+}
+
+func mustLimiter(b *testing.B, burst int, every time.Duration) *ratelimit.Limiter {
+	l, err := ratelimit.New(burst, every)
+	if err != nil {
+		b.Fatal(err)
+	}
+	return l
+}

ratelimit/limiter.go

@@ -0,0 +1,91 @@
+package ratelimit
+
+import (
+	"fmt"
+	"math"
+	"sync/atomic"
+	"time"
+)
+
+// Limiter is a token bucket: `burst` calls are allowed at once, then one more per `every`.
+// It counts the calls it denies and hands that count to the next allowed call.
+//
+// The budget belongs to the Limiter instance.
+// Every Allow that returns true spends from it, whatever the caller then does.
+//
+// Time follows Go's monotonic clock.
+//
+// Create one Limiter for each thing you want to limit,
+// sharing one between unrelated events lets the noisy one starve the quiet one.
+//
+// Check Allow before doing the work being limited, so a denied call stays cheap:
+//
+//	if denied, ok := probeLimiter.Allow(); ok {
+//		sendProbe()
+//		probesDenied.Inc(int64(denied))
+//	}
+//
+// For log lines use logging.RateLimited, which attaches the denied count to the logger as "suppressed".
+type Limiter struct {
+	every  int64
+	slack  int64
+	tat    atomic.Int64
+	denied atomic.Uint64
+}
+
+// maxSpan bounds `burst` * `every` to half of int64, leaving the other half for process uptime so the bucket
+// math can't overflow.
+const maxSpan = math.MaxInt64 / 2
+
+// New returns a Limiter allowing `burst` calls at once and one more per `every`.
+// A `burst` under 1 is 1, and an `every` of 0 or less never limits.
+// It returns an error if `burst` * `every` is over about 146 years.
+func New(burst int, every time.Duration) (*Limiter, error) {
+	l := &Limiter{}
+	if every > 0 {
+		burst = max(burst, 1)
+		if int64(burst) > maxSpan/int64(every) {
+			return nil, fmt.Errorf("ratelimit: burst %d * every %v is too long", burst, every)
+		}
+		l.every = int64(every)
+		l.slack = int64(burst-1) * l.every
+	}
+	return l, nil
+}
+
+var monoStart = time.Now()
+
+func monoNow() int64 { return int64(time.Since(monoStart)) }
+
+// Allow reports whether this call is allowed.
+// If a token is available it spends it and returns the calls denied since the last allowed one, and true.
+// Otherwise it returns 0 and false, and counts the denial.
+// A nil Limiter allows everything.
+func (l *Limiter) Allow() (uint64, bool) {
+	if l == nil {
+		return 0, true
+	}
+	return l.allow(monoNow())
+}
+
+// allow is Allow at a given monotonic time.
+func (l *Limiter) allow(now int64) (uint64, bool) {
+	// Concurrent Allow calls read the clock before racing for the bucket, so one can see a time another already spent
+	// past and be denied. An unlimited Limiter must never deny, so it skips the bucket.
+	if l.every == 0 {
+		return 0, true
+	}
+
+	for {
+		tat := l.tat.Load()
+		next := max(tat, now)
+		if next-now > l.slack {
+			l.denied.Add(1)
+			return 0, false
+		}
+
+		if l.tat.CompareAndSwap(tat, next+l.every) {
+			return l.denied.Swap(0), true
+		}
+	}
+}

ratelimit/limiter_test.go

@@ -0,0 +1,264 @@
+package ratelimit
+
+import (
+	"math"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func mustLimiter(t testing.TB, burst int, every time.Duration) *Limiter {
+	t.Helper()
+	l, err := New(burst, every)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return l
+}
+
+func TestLimiterBurstAndRefill(t *testing.T) {
+	l, now := mustLimiter(t, 3, time.Second), int64(0)
+
+	for i := 0; i < 3; i++ {
+		if n, ok := l.allow(now); !ok || n != 0 {
+			t.Fatalf("burst line %d: ok=%v denied=%d", i, ok, n)
+		}
+	}
+	for i := 0; i < 5; i++ {
+		if _, ok := l.allow(now); ok {
+			t.Fatalf("line %d past the burst logged", i)
+		}
+	}
+
+	now += int64(time.Second)
+	if n, ok := l.allow(now); !ok || n != 5 {
+		t.Fatalf("after one refill: ok=%v denied=%d, want true 5", ok, n)
+	}
+	if _, ok := l.allow(now); ok {
+		t.Fatal("one refill earned more than one line")
+	}
+
+	// A long quiet spell refills to burst, not past it, and the count reported above does not carry over.
+	now += int64(time.Hour)
+	for i, want := range []uint64{1, 0, 0} {
+		if n, ok := l.allow(now); !ok || n != want {
+			t.Fatalf("refilled line %d: ok=%v denied=%d, want true %d", i, ok, n, want)
+		}
+	}
+	if _, ok := l.allow(now); ok {
+		t.Fatal("bucket refilled past burst")
+	}
+}
+
+func TestLimiterBounds(t *testing.T) {
+	for _, every := range []time.Duration{0, -time.Second} {
+		l := mustLimiter(t, 1, every)
+		for i := 0; i < 100; i++ {
+			if _, ok := l.allow(0); !ok {
+				t.Fatalf("every=%v limited line %d", every, i)
+			}
+		}
+	}
+
+	l := mustLimiter(t, 0, time.Second)
+	if _, ok := l.allow(0); !ok {
+		t.Fatal("burst 0 should allow one line")
+	}
+	if _, ok := l.allow(0); ok {
+		t.Fatal("burst 0 should allow only one line")
+	}
+}
+
+func TestLimiterOverflow(t *testing.T) {
+	for _, c := range []struct {
+		burst int
+		every time.Duration
+	}{
+		{math.MaxInt, time.Second},
+		{1, math.MaxInt64},
+		{2, maxSpan/2 + 1},
+	} {
+		if _, err := New(c.burst, c.every); err == nil {
+			t.Fatalf("burst %d every %d should not fit", c.burst, c.every)
+		}
+	}
+
+	l := mustLimiter(t, 2, maxSpan/2)
+	for i := 0; i < 2; i++ {
+		if _, ok := l.allow(0); !ok {
+			t.Fatalf("the largest span should allow its burst, denied call %d", i)
+		}
+	}
+	if _, ok := l.allow(0); ok {
+		t.Fatal("the largest span allowed past its burst")
+	}
+}
+
+// Concurrent callers read the clock at slightly different times, which must not limit a Limiter that never limits.
+func TestLimiterUnlimitedConcurrent(t *testing.T) {
+	l := mustLimiter(t, 1, 0)
+
+	var denied atomic.Int64
+	var wg sync.WaitGroup
+	for i := 0; i < 16; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < 10000; j++ {
+				if _, ok := l.Allow(); !ok {
+					denied.Add(1)
+				}
+			}
+		}()
+	}
+	wg.Wait()
+
+	if got := denied.Load(); got != 0 {
+		t.Fatalf("denied %d lines", got)
+	}
+}
+
+func TestLimiterNil(t *testing.T) {
+	var l *Limiter
+	for i := 0; i < 10; i++ {
+		if _, ok := l.Allow(); !ok {
+			t.Fatal("nil limiter limited")
+		}
+	}
+}
+
+func TestLimiterConcurrent(t *testing.T) {
+	l := mustLimiter(t, 10, time.Hour)
+
+	var allowed atomic.Int64
+	var reported atomic.Uint64
+	var wg sync.WaitGroup
+	for i := 0; i < 16; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < 1000; j++ {
+				if n, ok := l.allow(0); ok {
+					allowed.Add(1)
+					reported.Add(n)
+				}
+			}
+		}()
+	}
+	wg.Wait()
+
+	if got := allowed.Load(); got != 10 {
+		t.Fatalf("allowed %d lines, want the burst of 10", got)
+	}
+	// An allowed call that lands after some denials reports them, so the rest are still pending.
+	if got := reported.Load() + l.denied.Load(); got != 16*1000-10 {
+		t.Fatalf("denied %d, want %d", got, 16*1000-10)
+	}
+}
+
+func BenchmarkAllowAllowed(b *testing.B) {
+	l := mustLimiter(b, 1, time.Nanosecond)
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		l.tat.Store(0)
+		l.Allow()
+	}
+}
+
+func BenchmarkAllowDenied(b *testing.B) {
+	l := mustLimiter(b, 1, time.Hour)
+	l.Allow()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		l.Allow()
+	}
+}
+
+func BenchmarkAllowDeniedParallel(b *testing.B) {
+	l := mustLimiter(b, 1, time.Hour)
+	l.Allow()
+	b.ReportAllocs()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			l.Allow()
+		}
+	})
+}
+
+// FuzzLimiter checks the bucket against its exact definition over arbitrary call timings.
+// A call is denied if and only if some allowed call a would leave too many calls for the time since a:
+// (allowed calls since a, plus this one) * every > burst * every + (now - a).
+// Every denied call must also be reported exactly once, by a later allowed call or still pending.
+func FuzzLimiter(f *testing.F) {
+	f.Add(uint8(3), uint16(1000), []byte{0, 0, 0, 0, 16, 0, 0, 200, 1, 1, 1, 15, 17, 0})
+	f.Add(uint8(1), uint16(1), []byte{0, 16, 15, 16, 0, 255, 0, 0})
+	f.Fuzz(func(t *testing.T, burst uint8, everyUs uint16, steps []byte) {
+		if len(steps) > 512 {
+			steps = steps[:512]
+		}
+		b := int64(burst%16) + 1
+		every := int64(everyUs%1000+1) * int64(time.Microsecond)
+		l := mustLimiter(t, int(b), time.Duration(every))
+
+		var now int64
+		var allowedAt []int64
+		var denied, reported uint64
+		for _, s := range steps {
+			// Each step advances by a sixteenth of every, so timings land on and around the refill.
+			now += int64(s) * every / 16
+
+			over := false
+			for i, a := range allowedAt {
+				if int64(len(allowedAt)-i+1)*every > b*every+now-a {
+					over = true
+					break
+				}
+			}
+
+			n, ok := l.allow(now)
+			if ok == over {
+				t.Fatalf("at %d with %d allowed before: ok=%v, the definition says %v", now, len(allowedAt), ok, !over)
+			}
+			if !ok {
+				denied++
+				continue
+			}
+			reported += n
+			allowedAt = append(allowedAt, now)
+		}
+
+		if got := reported + l.denied.Load(); got != denied {
+			t.Fatalf("reported %d denied calls, made %d", got, denied)
+		}
+	})
+}
+
+// Concurrent callers on the real clock must account for every denied call exactly once.
+func TestLimiterDeniedConservedConcurrent(t *testing.T) {
+	l := mustLimiter(t, 4, time.Microsecond)
+
+	var denied, reported atomic.Uint64
+	var wg sync.WaitGroup
+	for i := 0; i < 16; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < 20000; j++ {
+				if n, ok := l.Allow(); ok {
+					reported.Add(n)
+				} else {
+					denied.Add(1)
+				}
+			}
+		}()
+	}
+	wg.Wait()
+
+	if denied.Load() == 0 {
+		t.Fatal("nothing was denied, the test proves nothing")
+	}
+	if got := reported.Load() + l.denied.Load(); got != denied.Load() {
+		t.Fatalf("reported %d denied calls, made %d", got, denied.Load())
+	}
+}