add Clock interface and LeakyBucket

go.mod

@@ -3,6 +3,7 @@ module github.com/livekit/protocol
 go 1.21
 
 require (
+	github.com/benbjohnson/clock v1.3.5
 	github.com/eapache/channels v1.1.0
 	github.com/frostbyte73/core v0.0.10
 	github.com/fsnotify/fsnotify v1.7.0

go.sum

@@ -1,3 +1,5 @@
+github.com/benbjohnson/clock v1.3.5 h1:VvXlSJBzZpA/zum6Sj74hxwYI2DIxRWuNIoXAzHZz5o=
+github.com/benbjohnson/clock v1.3.5/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=

utils/clock.go

@@ -0,0 +1,20 @@
+package utils
+
+import "time"
+
+type Clock interface {
+	Now() time.Time
+	Sleep(time.Duration)
+}
+
+type SystemClock struct{}
+
+var _ Clock = &SystemClock{}
+
+func (clock *SystemClock) Now() time.Time {
+	return time.Now()
+}
+
+func (clock *SystemClock) Sleep(d time.Duration) {
+	time.Sleep(d)
+}

utils/rate.go

@@ -0,0 +1,92 @@
+// Copyright (c) 2022 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+// EDIT: slight modification to allow setting rate limit on the fly
+// SCOPE: LeakyBucket
+package utils
+
+import (
+	"sync/atomic"
+	"time"
+)
+
+type LeakyBucket struct {
+	//lint:ignore U1000 Padding is unused but it is crucial to maintain performance
+	// of this rate limiter in case of collocation with other frequently accessed memory.
+	prepadding [64]byte // cache line size = 64; created to avoid false sharing.
+	state      int64    // unix nanoseconds of the next permissions issue.
+	//lint:ignore U1000 like prepadding.
+	postpadding [56]byte // cache line size - state size = 64 - 8; created to avoid false sharing.
+
+	perRequest time.Duration
+	maxSlack   time.Duration
+	clock      Clock
+}
+
+func NewLeakyBucket(rateLimit int, slack time.Duration, clock Clock) *LeakyBucket {
+	var lb LeakyBucket
+	lb.SetRateLimit(rateLimit)
+	lb.maxSlack = slack * lb.perRequest
+	lb.clock = clock
+	atomic.StoreInt64(&lb.state, 0)
+	return &lb
+}
+
+func (lb *LeakyBucket) SetRateLimit(rateLimit int) {
+	lb.perRequest = time.Second / time.Duration(rateLimit)
+}
+
+// Take blocks to ensure that the time spent between multiple
+// Take calls is on average time.Second/rate.
+func (t *LeakyBucket) Take() time.Time {
+	var (
+		newTimeOfNextPermissionIssue int64
+		now                          int64
+	)
+	for {
+		now = t.clock.Now().UnixNano()
+		timeOfNextPermissionIssue := atomic.LoadInt64(&t.state)
+
+		switch {
+		case timeOfNextPermissionIssue == 0 || (t.maxSlack == 0 && now-timeOfNextPermissionIssue > int64(t.perRequest)):
+			// if this is our first call or t.maxSlack == 0 we need to shrink issue time to now
+			newTimeOfNextPermissionIssue = now
+		case t.maxSlack > 0 && now-timeOfNextPermissionIssue > int64(t.maxSlack)+int64(t.perRequest):
+			// a lot of nanoseconds passed since the last Take call
+			// we will limit max accumulated time to maxSlack
+			newTimeOfNextPermissionIssue = now - int64(t.maxSlack)
+		default:
+			// calculate the time at which our permission was issued
+			newTimeOfNextPermissionIssue = timeOfNextPermissionIssue + int64(t.perRequest)
+		}
+
+		if atomic.CompareAndSwapInt64(&t.state, timeOfNextPermissionIssue, newTimeOfNextPermissionIssue) {
+			break
+		}
+	}
+
+	sleepDuration := time.Duration(newTimeOfNextPermissionIssue - now)
+	if sleepDuration > 0 {
+		t.clock.Sleep(sleepDuration)
+		return time.Unix(0, newTimeOfNextPermissionIssue)
+	}
+	// return now if we don't sleep as atomicLimiter does
+	return time.Unix(0, now)
+}