Moved rate.Limiter into its own package

socketserver/server/rate/ratelimit.go

@@ -0,0 +1,77 @@
+package rate
+
+import (
+	"io"
+	"time"
+)
+
+// A Limiter supports a constant number of Performed() calls every
+// time a certain amount of time passes.
+//
+// Calls to Performed() when no "action tokens" are available will block
+// until one is available.
+type Limiter interface {
+	// Run begins emitting tokens for the ratelimiter.
+	// A call to Run must be followed by a call to Close.
+	Run()
+	// Performed consumes one token from the rate limiter.
+	// If no tokens are available, the call will block until one is.
+	Performed()
+	// Close stops the rate limiter. Any future calls to Performed() will block forever.
+	// Close never returns an error.
+	io.Closer
+}
+
+type timeRateLimit struct {
+	count  int
+	period time.Duration
+	ch     chan struct{}
+	done   chan struct{}
+}
+
+// Construct a new Limiter with the given count and duration.
+func NewRateLimit(count int, period time.Duration) Limiter {
+	return &timeRateLimit{
+		count:  count,
+		period: period,
+		ch:     make(chan struct{}),
+		done:   make(chan struct{}),
+	}
+}
+
+func (r *timeRateLimit) Run() {
+	for {
+		waiter := time.After(r.period)
+		for i := 0; i < r.count; i++ {
+			select {
+			case r.ch <- struct{}{}:
+				// ok
+			case <-r.done:
+				return
+			}
+		}
+		<-waiter
+	}
+}
+
+func (r *timeRateLimit) Performed() {
+	<-r.ch
+}
+
+func (r *timeRateLimit) Close() error {
+	close(r.done)
+	return nil
+}
+
+type unlimited struct{}
+
+var unlimitedInstance unlimited
+
+// Unlimited returns a Limiter that never blocks. The Run() and Close() calls are no-ops.
+func Unlimited() Limiter {
+	return unlimitedInstance
+}
+
+func (r unlimited) Run()         {}
+func (r unlimited) Performed()   {}
+func (r unlimited) Close() error { return nil }

socketserver/server/rate/ratelimit_test.go

@@ -0,0 +1,40 @@
+package rate
+
+import (
+	"testing"
+	"time"
+)
+
+var exampleData = []string{}
+
+func ExampleNewRateLimit() {
+	rl := NewRateLimit(100, 1*time.Minute)
+	go rl.Run()
+	defer rl.Close()
+
+	for _, v := range exampleData {
+		rl.Performed()
+		// do something with v
+		_ = v
+	}
+}
+
+func TestRateLimit(t *testing.T) {
+	rl := NewRateLimit(3, 100*time.Millisecond)
+	start := time.Now()
+	go rl.Run()
+	for i := 0; i < 4; i++ {
+		rl.Performed()
+	}
+	end := time.Now()
+	if end.Sub(start) < 100*time.Millisecond {
+		t.Error("ratelimiter did not wait for period to expire")
+	}
+	rl.Performed()
+	rl.Performed()
+	end2 := time.Now()
+	if end2.Sub(end) > 10*time.Millisecond {
+		t.Error("ratelimiter improperly waited when tokens were available")
+	}
+	rl.Close()
+}