Merge pull request #1510 from sulemaanhamza/pkg-policyhandler-cache

Enhancement of TimedCache with Bug Fixes and Test Suite

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{})
 }

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)
+}