Fixed data race in mock cache for timesweptcache unit tests

util/cache/cacher/lrucache.go

@@ -13,7 +13,7 @@ type KeyPair struct {
 	Value interface{}
 }
 
-// A simple LRU Cache implementation. The least recently used elements
+// A simple LRU Cache implementation. The least recently used element
 // in the cache are removed if the cache's max capacity is hit. The cache's
 // mutex cannot be a RWMutex since Gets alter the cache's underlying data structures.
 // Slower concurrent performance than a Random Removal cache, but more cache hits.
@@ -111,6 +111,7 @@ func (cache *LruCache) Put(key string, value interface{}) {
 	}
 }
 
+// Removes an element from the cache.
 // Does NOT take the cache's lock. Functions calling removeElementNoLock()
 // need to do it themselves
 func (cache *LruCache) removeElementNoLock(key string) {

util/cache/timesweptcache_test.go

@@ -2,43 +2,53 @@ package cache
 
 import (
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/suite"
 )
 
 type MockCache struct {
-	elements map[string]interface{}
+	elements  map[string]interface{}
+	cacheLock sync.RWMutex
 }
 
 func (cache *MockCache) Get(key string) (interface{}, bool) {
+	cache.cacheLock.RLock()
+	defer cache.cacheLock.RUnlock()
 	val, ok := cache.elements[key]
 	return val, ok
 }
 
 func (cache *MockCache) Put(key string, value interface{}) {
+	cache.cacheLock.Lock()
+	defer cache.cacheLock.Unlock()
 	cache.elements[key] = value
 }
 
 func (cache *MockCache) Clear() {
+	cache.cacheLock.Lock()
+	defer cache.cacheLock.Unlock()
 	cache.elements = make(map[string]interface{})
 }
 
 func (cache *MockCache) Remove(key string) {
-	delete(cache.elements, key)
+	cache.cacheLock.Lock()
+	defer cache.cacheLock.Unlock()
+	cache.removeNoLock(key)
 }
 
-func (cache *MockCache) GetIterator() func() (string, interface{}, bool) {
-	return func() (string, interface{}, bool) {
-		return "true", "true", false
-	}
+func (cache *MockCache) removeNoLock(key string) {
+	delete(cache.elements, key)
 }
 
 func (cache *MockCache) ForEach(cleanUp func(string, interface{}) bool) {
+	cache.cacheLock.Lock()
+	defer cache.cacheLock.Unlock()
 	for key, val := range cache.elements {
 		if cleanUp(key, val) {
-			cache.Remove(key)
+			cache.removeNoLock(key)
 		}
 	}
 }