add more tests and benchmark for cache

refactor cache to use atomic uint
2021-01-19 23:23:41 +01:00
parent 806d000f22
commit af870e005e
4 changed files with 109 additions and 62 deletions
@@ -178,7 +178,7 @@ func (s Service) ListAccounts(ctx context.Context, in *proto.ListAccountsRequest
 			vh.Write([]byte(a.PasswordProfile.Password))
 			v := vh.Sum([]byte(password))

-			e := passwordValidCache.Get(k)
+			e := passwordValidCache.Load(k)

 			if e == nil {
 				suspicious = !isPasswordValid(s.log, a.PasswordProfile.Password, password)
@@ -187,12 +187,12 @@ func (s Service) ListAccounts(ctx context.Context, in *proto.ListAccountsRequest
 			}

 			if suspicious {
-				passwordValidCache.Unset(k)
+				passwordValidCache.Delete(k)
 				return merrors.Unauthorized(s.id, "account not found or invalid credentials")
 			}

 			if e == nil {
-				passwordValidCache.Set(k, v, time.Now().Add(passwordValidCacheExpiration))
+				passwordValidCache.Store(k, v, time.Now().Add(passwordValidCacheExpiration))
 			}
 		}

@@ -2,6 +2,7 @@ package sync

 import (
 	"sync"
+	"sync/atomic"
 	"time"
 )

@@ -9,8 +10,8 @@ import (
 type Cache struct {
 	entries  sync.Map
 	pool     sync.Pool
-	capacity int
-	length   int
+	capacity uint64
+	length   uint64
 }

 // CacheEntry represents an entry on the cache. You can type assert on V.
@@ -22,15 +23,15 @@ type CacheEntry struct {
 // NewCache returns a new instance of Cache.
 func NewCache(capacity int) Cache {
 	return Cache{
-		capacity: capacity,
+		capacity: uint64(capacity),
 		pool: sync.Pool{New: func() interface{} {
 			return new(CacheEntry)
 		}},
 	}
 }

-// Get gets an entry by given key
-func (c *Cache) Get(key string) *CacheEntry {
+// Load loads an entry by given key
+func (c *Cache) Load(key string) *CacheEntry {
 	if mapEntry, ok := c.entries.Load(key); ok {
 		entry := mapEntry.(*CacheEntry)
 		if c.expired(entry) {
@@ -42,9 +43,9 @@ func (c *Cache) Get(key string) *CacheEntry {
 	return nil
 }

-// Set adds an entry for given key and value
-func (c *Cache) Set(key string, val interface{}, expiration time.Time) {
-	if !c.fits() {
+// Store adds an entry for given key and value
+func (c *Cache) Store(key string, val interface{}, expiration time.Time) {
+	if c.length > c.capacity {
 		c.evict()
 	}

@@ -60,18 +61,19 @@ func (c *Cache) Set(key string, val interface{}, expiration time.Time) {
 		entry.V = val
 		entry.expiration = expiration

-		c.length++
+		atomic.AddUint64(&c.length, 1)
 	}
 }

-// Unset removes an entry by given key
-func (c *Cache) Unset(key string) bool {
-	if _, loaded := c.entries.LoadAndDelete(key); !loaded {
-		return false
+// Delete removes an entry by given key
+func (c *Cache) Delete(key string) bool {
+	_, loaded := c.entries.LoadAndDelete(key)
+
+	if loaded {
+		atomic.AddUint64(&c.length, ^uint64(0))
 	}

-	c.length--
-	return true
+	return loaded
 }

 // evict frees memory from the cache by removing entries that exceeded the cache TTL.
@@ -79,8 +81,7 @@ func (c *Cache) evict() {
 	c.entries.Range(func(key, mapEntry interface{}) bool {
 		entry := mapEntry.(*CacheEntry)
 		if c.expired(entry) {
-			c.entries.Delete(key)
-			c.length--
+			c.Delete(key.(string))
 		}
 		return true
 	})
@@ -90,8 +91,3 @@ func (c *Cache) evict() {
 func (c *Cache) expired(e *CacheEntry) bool {
 	return e.expiration.Before(time.Now())
 }
-
-// fits returns whether the cache fits more entries.
-func (c *Cache) fits() bool {
-	return c.capacity > c.length
-}
@@ -3,52 +3,103 @@ package sync
 import (
 	"github.com/stretchr/testify/assert"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 )

-func TestCache_Get(t *testing.T) {
-	size := 1024
+func cacheRunner(size int) (*Cache, func(f func(i int))) {
 	c := NewCache(size)
-
-	for i := 0; i < size; i++ {
-		c.Set(strconv.Itoa(i), i, time.Now().Add(10*time.Second))
+	run := func(f func(i int)) {
+		wg := sync.WaitGroup{}
+		for i := 0; i < size; i++ {
+			wg.Add(1)
+			go func(i int) {
+				f(i)
+				wg.Done()
+			}(i)
+		}
+		wg.Wait()
 	}

-	for i := 0; i < size; i++ {
-		assert.Equal(t, i, c.Get(strconv.Itoa(i)).V, "entry value is the same")
-	}
-
-	assert.Nil(t, c.Get("unknown"), "entry is nil if unknown")
-
-	wait := 10 * time.Millisecond
-	c.Set("expired", size, time.Now().Add(wait))
-	time.Sleep(wait + 1)
-	assert.Nil(t, c.Get(strconv.Itoa(size)), "entry is nil if it's expired")
+	return &c, run
 }

-func TestCache_Set(t *testing.T) {
-	c := NewCache(1)
+func BenchmarkCache(b *testing.B) {
+	b.ReportAllocs()
+	size := 1024
+	c, cr := cacheRunner(size)

-	c.Set("new", "new", time.Now().Add(10*time.Millisecond))
-	assert.Equal(t, "new", c.Get("new").V, "new entries can be added")
-	assert.Equal(t, 1, c.length, "adding new entries will increase the cache size")
-
-	replacedExpiration := time.Now().Add(10 * time.Millisecond)
-	c.Set("new", "updated", replacedExpiration)
-	assert.Equal(t, "updated", c.Get("new").V, "entry values can be updated")
-	assert.Equal(t, replacedExpiration, c.Get("new").expiration, "entry expiration can be updated")
-
-	time.Sleep(11 * time.Millisecond)
-	c.Set("eviction", "eviction", time.Now())
-	assert.Equal(t, 1, c.length, "expired entries get removed")
+	cr(func(i int) { c.Store(strconv.Itoa(i), i, time.Now().Add(10*time.Millisecond)) })
+	cr(func(i int) { c.Delete(strconv.Itoa(i)) })
 }

-func TestCache_Unset(t *testing.T) {
-	c := NewCache(1)
+func TestCache(t *testing.T) {
+	size := 1024
+	c, cr := cacheRunner(size)

-	c.Set("new", "new", time.Now().Add(10*time.Millisecond))
-	c.Unset("new")
-	assert.Nil(t, c.Get("new"), "entries can be removed")
-	assert.Equal(t, 0, c.length, "removing a entry decreases the cache size")
+	cr(func(i int) { c.Store(strconv.Itoa(i), i, time.Now().Add(10*time.Millisecond)) })
+	assert.Equal(t, size, int(c.length), "length is atomic")
+
+	cr(func(i int) { c.Delete(strconv.Itoa(i)) })
+	assert.Equal(t, 0, int(c.length), "delete is atomic")
+
+	cr(func(i int) {
+		time.Sleep(11 * time.Millisecond)
+		c.evict()
+	})
+	assert.Equal(t, 0, int(c.length), "evict is atomic")
+}
+
+func TestCache_Load(t *testing.T) {
+	size := 1024
+	c, cr := cacheRunner(size)
+
+	cr(func(i int) {
+		c.Store(strconv.Itoa(i), i, time.Now().Add(10*time.Second))
+	})
+
+	cr(func(i int) {
+		assert.Equal(t, i, c.Load(strconv.Itoa(i)).V, "entry value is the same")
+	})
+
+	cr(func(i int) {
+		assert.Nil(t, c.Load(strconv.Itoa(i+size)), "entry is nil if unknown")
+	})
+
+	cr(func(i int) {
+		wait := 10 * time.Millisecond
+		c.Store(strconv.Itoa(i), i, time.Now().Add(wait))
+		time.Sleep(wait + 1)
+		assert.Nil(t, c.Load(strconv.Itoa(i)), "entry is nil if it's expired")
+	})
+}
+
+func TestCache_Store(t *testing.T) {
+	c, cr := cacheRunner(1024)
+
+	cr(func(i int) {
+		c.Store(strconv.Itoa(i), i, time.Now().Add(10*time.Millisecond))
+		assert.Equal(t, i, c.Load(strconv.Itoa(i)).V, "new entries can be added")
+	})
+
+	cr(func(i int) {
+		replacedExpiration := time.Now().Add(10 * time.Millisecond)
+		c.Store(strconv.Itoa(i), "old", time.Now().Add(10*time.Minute))
+		c.Store(strconv.Itoa(i), "updated", replacedExpiration)
+		assert.Equal(t, "updated", c.Load(strconv.Itoa(i)).V, "entry values can be updated")
+		assert.Equal(t, replacedExpiration, c.Load(strconv.Itoa(i)).expiration, "entry expiration can be updated")
+	})
+}
+
+func TestCache_Delete(t *testing.T) {
+	c, cr := cacheRunner(1024)
+
+	cr(func(i int) {
+		c.Store(strconv.Itoa(i), i, time.Now().Add(10*time.Millisecond))
+		c.Delete(strconv.Itoa(i))
+		assert.Nil(t, c.Load(strconv.Itoa(i)), "entries can be deleted")
+	})
+
+	assert.Equal(t, 0, int(c.length), "removing a entry decreases the cache size")
 }
@@ -79,7 +79,7 @@ type oidcAuth struct {
 }

 func (m oidcAuth) getClaims(token string, req *http.Request) (claims oidc.StandardClaims, status int) {
-	hit := m.tokenCache.Get(token)
+	hit := m.tokenCache.Load(token)
 	if hit == nil {
 		// TODO cache userinfo for access token if we can determine the expiry (which works in case it is a jwt based access token)
 		oauth2Token := &oauth2.Token{
@@ -106,7 +106,7 @@ func (m oidcAuth) getClaims(token string, req *http.Request) (claims oidc.Standa
 		claims.Iss = m.oidcIss

 		expiration := m.extractExpiration(token)
-		m.tokenCache.Set(token, claims, expiration)
+		m.tokenCache.Store(token, claims, expiration)

 		m.logger.Debug().Interface("claims", claims).Interface("userInfo", userInfo).Time("expiration", expiration.UTC()).Msg("unmarshalled and cached userinfo")
 		return