diff --git a/core/pkg/policyhandler/cache.go b/core/pkg/policyhandler/cache.go index d1c0690b..12136939 100644 --- a/core/pkg/policyhandler/cache.go +++ b/core/pkg/policyhandler/cache.go @@ -14,14 +14,16 @@ type TimedCache[T any] struct { value T isSet bool ttl time.Duration - expiration int64 + expiration time.Time mutex sync.RWMutex + stopChan chan struct{} // to stop the invalidateTask goroutine } func NewTimedCache[T any](ttl time.Duration) *TimedCache[T] { cache := &TimedCache[T]{ - ttl: ttl, - isSet: false, + ttl: ttl, + isSet: false, + stopChan: make(chan struct{}), } // start the invalidate task only when the ttl is greater than 0 (cache is enabled) @@ -36,38 +38,71 @@ func (c *TimedCache[T]) Set(value T) { c.mutex.Lock() defer c.mutex.Unlock() - // cache is disabled if c.ttl == 0 { return } c.isSet = true c.value = value - c.expiration = time.Now().Add(c.ttl).UnixNano() + c.expiration = time.Now().Add(c.ttl) + + // Signal invalidation to Get() if cache is already expired + if time.Now().After(c.expiration) { + c.Invalidate() + } } func (c *TimedCache[T]) Get() (T, bool) { c.mutex.RLock() defer c.mutex.RUnlock() - if !c.isSet || time.Now().UnixNano() > c.expiration { + // If the invalidateTask() goroutine is currently invalidating the cache, + // the Get() method may return the stale cached value before the invalidation is complete. + // To avoid the stale cached value, we're requiring the Get() method to wait for the invalidation signal before returning the value. + select { + case <-c.stopChan: return c.value, false + default: + if !c.isSet || time.Now().After(c.expiration) { + return c.value, false + } + return c.value, true } - return c.value, true } func (c *TimedCache[T]) invalidateTask() { + ticker := time.NewTicker(c.ttl) + defer ticker.Stop() + for { - <-time.After(c.ttl) - if time.Now().UnixNano() > c.expiration { - c.Invalidate() + select { + case <-ticker.C: + c.mutex.Lock() + expired := time.Now().After(c.expiration) + c.mutex.Unlock() + if expired { + c.Invalidate() + } + case <-c.stopChan: + return + default: + // Check if TTL is still non-zero and return if true, to avoid possible memory leaks + if c.ttl == 0 { + return + } } } } +func (c *TimedCache[T]) Stop() { + close(c.stopChan) +} + func (c *TimedCache[T]) Invalidate() { c.mutex.Lock() defer c.mutex.Unlock() c.isSet = false + close(c.stopChan) + c.stopChan = make(chan struct{}) } diff --git a/core/pkg/policyhandler/cache_test.go b/core/pkg/policyhandler/cache_test.go index fc9d5482..3a3745b1 100644 --- a/core/pkg/policyhandler/cache_test.go +++ b/core/pkg/policyhandler/cache_test.go @@ -73,3 +73,69 @@ func TestTimedCache(t *testing.T) { }) } } + +func TestCache_SetAndGet(t *testing.T) { + cache := NewTimedCache[int](time.Second * 2) + + cache.Set(42) + + value, exists := cache.Get() + if !exists || value != 42 { + t.Errorf("Expected value: %v, Got: %v, Exists: %v", 42, value, exists) + } +} + +func TestCache_Expiration(t *testing.T) { + cache := NewTimedCache[int](time.Millisecond * 500) + + cache.Set(42) + + time.Sleep(time.Millisecond * 1000) // Wait for expiration + + value, exists := cache.Get() + if exists { + t.Errorf("Expected cache to be expired, but got value: %v", value) + } +} + +func TestCache_WithZeroTTL(t *testing.T) { + cache := NewTimedCache[string](0) + + cache.Set("hello") + + value, exists := cache.Get() + if exists { + t.Errorf("Expected cache to be disabled, but got value: %v", value) + } +} + +func TestCache_Invalidate(t *testing.T) { + cache := NewTimedCache[string](time.Second * 2) + + cache.Set("initial value") + + cache.Invalidate() + + value, exists := cache.Get() + if exists { + t.Errorf("Expected cache to be invalidated, but got value: %v", value) + } +} + +func TestCache_ConcurrentAccess(t *testing.T) { + cache := NewTimedCache[int](time.Second * 1) + + go func() { + cache.Set(42) + }() + + go func() { + time.Sleep(time.Millisecond * 500) + value, exists := cache.Get() + if !exists || value != 42 { + t.Errorf("Expected value: %v, Got: %v, Exists: %v", 42, value, exists) + } + }() + + time.Sleep(time.Second) +}