Allow specifying max age for caches (#2239)

* Allow specifying max age for caches

* Evict cache entry if it's found to be stale when we call Get

* Fix bugs

7 files changed, 48 insertions, 5 deletions

diff --git a/internal/caching/cache_federationevents.go b/internal/caching/cache_federationevents.go
index d10b333a..b79cc809 100644
--- a/internal/caching/cache_federationevents.go
+++ b/internal/caching/cache_federationevents.go
@@ -10,6 +10,7 @@ const (
 	FederationEventCacheName       = "federation_event"
 	FederationEventCacheMaxEntries = 256
 	FederationEventCacheMutable    = true // to allow use of Unset only
+	FederationEventCacheMaxAge     = CacheNoMaxAge
 )
 
 // FederationCache contains the subset of functions needed for
diff --git a/internal/caching/cache_roominfo.go b/internal/caching/cache_roominfo.go
index f32d6ba9..60d22128 100644
--- a/internal/caching/cache_roominfo.go
+++ b/internal/caching/cache_roominfo.go
@@ -1,6 +1,8 @@
 package caching
 
 import (
+	"time"
+
 	"github.com/matrix-org/dendrite/roomserver/types"
 )
 
@@ -16,6 +18,7 @@ const (
 	RoomInfoCacheName       = "roominfo"
 	RoomInfoCacheMaxEntries = 1024
 	RoomInfoCacheMutable    = true
+	RoomInfoCacheMaxAge     = time.Minute * 5
 )
 
 // RoomInfosCache contains the subset of functions needed for
diff --git a/internal/caching/cache_roomservernids.go b/internal/caching/cache_roomservernids.go
index 6d413093..1918a2f1 100644
--- a/internal/caching/cache_roomservernids.go
+++ b/internal/caching/cache_roomservernids.go
@@ -10,6 +10,7 @@ const (
 	RoomServerRoomIDsCacheName       = "roomserver_room_ids"
 	RoomServerRoomIDsCacheMaxEntries = 1024
 	RoomServerRoomIDsCacheMutable    = false
+	RoomServerRoomIDsCacheMaxAge     = CacheNoMaxAge
 )
 
 type RoomServerCaches interface {
diff --git a/internal/caching/cache_roomversions.go b/internal/caching/cache_roomversions.go
index 0b46d3d4..92d2eab0 100644
--- a/internal/caching/cache_roomversions.go
+++ b/internal/caching/cache_roomversions.go
@@ -6,6 +6,7 @@ const (
 	RoomVersionCacheName       = "room_versions"
 	RoomVersionCacheMaxEntries = 1024
 	RoomVersionCacheMutable    = false
+	RoomVersionCacheMaxAge     = CacheNoMaxAge
 )
 
 // RoomVersionsCache contains the subset of functions needed for
diff --git a/internal/caching/cache_serverkeys.go b/internal/caching/cache_serverkeys.go
index 4697fb4d..4eb10fe6 100644
--- a/internal/caching/cache_serverkeys.go
+++ b/internal/caching/cache_serverkeys.go
@@ -10,6 +10,7 @@ const (
 	ServerKeyCacheName       = "server_key"
 	ServerKeyCacheMaxEntries = 4096
 	ServerKeyCacheMutable    = true
+	ServerKeyCacheMaxAge     = CacheNoMaxAge
 )
 
 // ServerKeyCache contains the subset of functions needed for
diff --git a/internal/caching/caches.go b/internal/caching/caches.go
index 39926d73..722405de 100644
--- a/internal/caching/caches.go
+++ b/internal/caching/caches.go
@@ -1,5 +1,7 @@
 package caching
 
+import "time"
+
 // Caches contains a set of references to caches. They may be
 // different implementations as long as they satisfy the Cache
 // interface.
@@ -19,3 +21,5 @@ type Cache interface {
 	Set(key string, value interface{})
 	Unset(key string)
 }
+
+const CacheNoMaxAge = time.Duration(0)
diff --git a/internal/caching/impl_inmemorylru.go b/internal/caching/impl_inmemorylru.go
index 7b21f136..c9c8fd08 100644
--- a/internal/caching/impl_inmemorylru.go
+++ b/internal/caching/impl_inmemorylru.go
@@ -14,6 +14,7 @@ func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 		RoomVersionCacheName,
 		RoomVersionCacheMutable,
 		RoomVersionCacheMaxEntries,
+		RoomVersionCacheMaxAge,
 		enablePrometheus,
 	)
 	if err != nil {
@@ -23,6 +24,7 @@ func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 		ServerKeyCacheName,
 		ServerKeyCacheMutable,
 		ServerKeyCacheMaxEntries,
+		ServerKeyCacheMaxAge,
 		enablePrometheus,
 	)
 	if err != nil {
@@ -32,6 +34,7 @@ func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 		RoomServerRoomIDsCacheName,
 		RoomServerRoomIDsCacheMutable,
 		RoomServerRoomIDsCacheMaxEntries,
+		RoomServerRoomIDsCacheMaxAge,
 		enablePrometheus,
 	)
 	if err != nil {
@@ -41,6 +44,7 @@ func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 		RoomInfoCacheName,
 		RoomInfoCacheMutable,
 		RoomInfoCacheMaxEntries,
+		RoomInfoCacheMaxAge,
 		enablePrometheus,
 	)
 	if err != nil {
@@ -50,6 +54,7 @@ func NewInMemoryLRUCache(enablePrometheus bool) (*Caches, error) {
 		FederationEventCacheName,
 		FederationEventCacheMutable,
 		FederationEventCacheMaxEntries,
+		FederationEventCacheMaxAge,
 		enablePrometheus,
 	)
 	if err != nil {
@@ -96,15 +101,22 @@ type InMemoryLRUCachePartition struct {
 	name       string
 	mutable    bool
 	maxEntries int
+	maxAge     time.Duration
 	lru        *lru.Cache
 }
 
-func NewInMemoryLRUCachePartition(name string, mutable bool, maxEntries int, enablePrometheus bool) (*InMemoryLRUCachePartition, error) {
+type inMemoryLRUCacheEntry struct {
+	value   interface{}
+	created time.Time
+}
+
+func NewInMemoryLRUCachePartition(name string, mutable bool, maxEntries int, maxAge time.Duration, enablePrometheus bool) (*InMemoryLRUCachePartition, error) {
 	var err error
 	cache := InMemoryLRUCachePartition{
 		name:       name,
 		mutable:    mutable,
 		maxEntries: maxEntries,
+		maxAge:     maxAge,
 	}
 	cache.lru, err = lru.New(maxEntries)
 	if err != nil {
@@ -124,11 +136,16 @@ func NewInMemoryLRUCachePartition(name string, mutable bool, maxEntries int, ena
 
 func (c *InMemoryLRUCachePartition) Set(key string, value interface{}) {
 	if !c.mutable {
-		if peek, ok := c.lru.Peek(key); ok && peek != value {
-			panic(fmt.Sprintf("invalid use of immutable cache tries to mutate existing value of %q", key))
+		if peek, ok := c.lru.Peek(key); ok {
+			if entry, ok := peek.(*inMemoryLRUCacheEntry); ok && entry.value != value {
+				panic(fmt.Sprintf("invalid use of immutable cache tries to mutate existing value of %q", key))
+			}
 		}
 	}
-	c.lru.Add(key, value)
+	c.lru.Add(key, &inMemoryLRUCacheEntry{
+		value:   value,
+		created: time.Now(),
+	})
 }
 
 func (c *InMemoryLRUCachePartition) Unset(key string) {
@@ -139,5 +156,20 @@ func (c *InMemoryLRUCachePartition) Unset(key string) {
 }
 
 func (c *InMemoryLRUCachePartition) Get(key string) (value interface{}, ok bool) {
-	return c.lru.Get(key)
+	v, ok := c.lru.Get(key)
+	if !ok {
+		return nil, false
+	}
+	entry, ok := v.(*inMemoryLRUCacheEntry)
+	switch {
+	case ok && c.maxAge == CacheNoMaxAge:
+		return entry.value, ok // There's no maximum age policy
+	case ok && time.Since(entry.created) < c.maxAge:
+		return entry.value, ok // The value for the key isn't stale
+	default:
+		// Either the key was found and it was stale, or the key
+		// wasn't found at all
+		c.lru.Remove(key)
+		return nil, false
+	}
 }