From 172634ed546c6677001c423d13082b8526e39db6 Mon Sep 17 00:00:00 2001
From: Ananya Kumar Mallik <ananya.mallik@paymentsense.com>
Date: Fri, 6 Dec 2024 15:32:54 +0000
Subject: [PATCH] Fix caching logic and TTL for cache

Signed-off-by: Ananya Kumar Mallik <ananya.mallik@paymentsense.com>
---
 collectors/cache.go     | 36 ++++++++++++++++++++++++++++++++++--
 stackdriver_exporter.go | 13 ++++++-------
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/collectors/cache.go b/collectors/cache.go
index 49eb4f7..1761fdc 100644
--- a/collectors/cache.go
+++ b/collectors/cache.go
@@ -88,21 +88,33 @@ type collectorCacheEntry struct {
 
 // NewCollectorCache returns a new CollectorCache with the given TTL
 func NewCollectorCache(ttl time.Duration) *CollectorCache {
-	return &CollectorCache{
+	c := &CollectorCache{
 		cache: make(map[string]*collectorCacheEntry),
 		ttl:   ttl,
 	}
+
+	go c.cleanup()
+	return c
 }
 
 // Get returns a MonitoringCollector if the key is found and not expired
+// If key is found it resets the TTL for the collector
 func (c *CollectorCache) Get(key string) (*MonitoringCollector, bool) {
 	c.lock.RLock()
 	defer c.lock.RUnlock()
 
 	entry, ok := c.cache[key]
-	if !ok || time.Now().After(entry.expiry) {
+
+	if !ok {
+		return nil, false
+	}
+
+	if time.Now().After(entry.expiry) {
+		delete(c.cache, key)
 		return nil, false
 	}
+
+	entry.expiry = time.Now().Add(c.ttl)
 	return entry.collector, true
 }
 
@@ -116,3 +128,23 @@ func (c *CollectorCache) Store(key string, collector *MonitoringCollector) {
 	defer c.lock.Unlock()
 	c.cache[key] = entry
 }
+
+func (c *CollectorCache) cleanup() {
+	ticker := time.NewTicker(5 * time.Minute)
+	defer ticker.Stop()
+	for range ticker.C {
+		c.removeExpired()
+	}
+}
+
+func (c *CollectorCache) removeExpired() {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	now := time.Now()
+	for key, entry := range c.cache {
+		if now.After(entry.expiry) {
+			delete(c.cache, key)
+		}
+	}
+}
diff --git a/stackdriver_exporter.go b/stackdriver_exporter.go
index b60f644..b2ffcb3 100644
--- a/stackdriver_exporter.go
+++ b/stackdriver_exporter.go
@@ -205,16 +205,15 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 }
 
 func newHandler(projectIDs []string, metricPrefixes []string, metricExtraFilters []collectors.MetricFilter, m *monitoring.Service, logger *slog.Logger, additionalGatherer prometheus.Gatherer) *handler {
-	ttl := 2 * time.Hour
+	var ttl time.Duration
 	// Add collector caching TTL as max of deltas aggregation or descriptor caching
 	if *monitoringMetricsAggregateDeltas || *monitoringDescriptorCacheTTL > 0 {
-		featureTTL := *monitoringMetricsDeltasTTL
-		if *monitoringDescriptorCacheTTL > featureTTL {
-			featureTTL = *monitoringDescriptorCacheTTL
-		}
-		if featureTTL > ttl {
-			ttl = featureTTL
+		ttl = *monitoringMetricsDeltasTTL
+		if *monitoringDescriptorCacheTTL > ttl {
+			ttl = *monitoringDescriptorCacheTTL
 		}
+	} else {
+		ttl = 2 * time.Hour
 	}
 
 	logger.Info("Creating collector cache", "ttl", ttl)