Make Registry thread-safe

registry.go

@@ -2,6 +2,7 @@ package limiters
 
 import (
 	"container/heap"
+	"sync"
 	"time"
 )
 
@@ -44,8 +45,10 @@ func (pq *gcPq) Pop() interface{} {
 	return item
 }
 
-// Registry is a garbage-collectable registry of values.
+// Registry is a thread-safe garbage-collectable registry of values.
 type Registry struct {
+	// Guards all the fields below it.
+	mx sync.Mutex
 	pq *gcPq
 	m  map[string]*pqItem
 }
@@ -59,6 +62,8 @@ func NewRegistry() *Registry {
 // GetOrCreate gets an existing value by key and updates its expiration time.
 // If the key lookup fails it creates a new value by calling the provided value closure and puts it on the queue.
 func (r *Registry) GetOrCreate(key string, value func() interface{}, ttl time.Duration, now time.Time) interface{} {
+	r.mx.Lock()
+	defer r.mx.Unlock()
 	item, ok := r.m[key]
 	if ok {
 		// Update the expiration time.
@@ -79,6 +84,8 @@ func (r *Registry) GetOrCreate(key string, value func() interface{}, ttl time.Du
 
 // DeleteExpired deletes expired items from the registry and returns the number of deleted items.
 func (r *Registry) DeleteExpired(now time.Time) int {
+	r.mx.Lock()
+	defer r.mx.Unlock()
 	c := 0
 	for {
 		if len(*r.pq) == 0 {
@@ -97,6 +104,8 @@ func (r *Registry) DeleteExpired(now time.Time) int {
 
 // Delete deletes an item from the registry.
 func (r *Registry) Delete(key string) {
+	r.mx.Lock()
+	defer r.mx.Unlock()
 	item, ok := r.m[key]
 	if !ok {
 		return
@@ -107,11 +116,15 @@ func (r *Registry) Delete(key string) {
 
 // Exists returns true if an item with the given key exists in the registry.
 func (r *Registry) Exists(key string) bool {
+	r.mx.Lock()
+	defer r.mx.Unlock()
 	_, ok := r.m[key]
 	return ok
 }
 
 // Len returns the number of items in the registry.
 func (r *Registry) Len() int {
+	r.mx.Lock()
+	defer r.mx.Unlock()
 	return len(*r.pq)
 }

registry_test.go

@@ -3,6 +3,7 @@ package limiters_test
 import (
 	"context"
 	"fmt"
+	"strconv"
 	"testing"
 	"time"
 
@@ -82,3 +83,19 @@ func TestRegistry_Delete(t *testing.T) {
 	registry.Delete("key")
 	assert.False(t, registry.Exists("key"))
 }
+
+// This test is expected to fail when run with the --race flag.
+func TestRegistry_ConcurrentUsage(t *testing.T) {
+	registry := limiters.NewRegistry()
+	clock := newFakeClock()
+	for i := 0; i < 10; i++ {
+		go func(i int) {
+			registry.GetOrCreate(strconv.Itoa(i), func() interface{} { return &struct{}{} }, 0, clock.Now())
+		}(i)
+	}
+	for i := 0; i < 10; i++ {
+		go func(i int) {
+			registry.DeleteExpired(clock.Now())
+		}(i)
+	}
+}