More flexible caching (#1101)

6 files changed, 159 insertions, 112 deletions

diff --git a/internal/caching/cache_roomversions.go b/internal/caching/cache_roomversions.go
new file mode 100644
index 00000000..0b46d3d4
--- /dev/null
+++ b/internal/caching/cache_roomversions.go
@@ -0,0 +1,30 @@
+package caching
+
+import "github.com/matrix-org/gomatrixserverlib"
+
+const (
+	RoomVersionCacheName       = "room_versions"
+	RoomVersionCacheMaxEntries = 1024
+	RoomVersionCacheMutable    = false
+)
+
+// RoomVersionsCache contains the subset of functions needed for
+// a room version cache.
+type RoomVersionCache interface {
+	GetRoomVersion(roomID string) (roomVersion gomatrixserverlib.RoomVersion, ok bool)
+	StoreRoomVersion(roomID string, roomVersion gomatrixserverlib.RoomVersion)
+}
+
+func (c Caches) GetRoomVersion(roomID string) (gomatrixserverlib.RoomVersion, bool) {
+	val, found := c.RoomVersions.Get(roomID)
+	if found && val != nil {
+		if roomVersion, ok := val.(gomatrixserverlib.RoomVersion); ok {
+			return roomVersion, true
+		}
+	}
+	return "", false
+}
+
+func (c Caches) StoreRoomVersion(roomID string, roomVersion gomatrixserverlib.RoomVersion) {
+	c.RoomVersions.Set(roomID, roomVersion)
+}
diff --git a/internal/caching/cache_serverkeys.go b/internal/caching/cache_serverkeys.go
new file mode 100644
index 00000000..8c71ffbd
--- /dev/null
+++ b/internal/caching/cache_serverkeys.go
@@ -0,0 +1,41 @@
+package caching
+
+import (
+	"fmt"
+
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+const (
+	ServerKeyCacheName       = "server_key"
+	ServerKeyCacheMaxEntries = 4096
+	ServerKeyCacheMutable    = true
+)
+
+// ServerKeyCache contains the subset of functions needed for
+// a server key cache.
+type ServerKeyCache interface {
+	GetServerKey(request gomatrixserverlib.PublicKeyLookupRequest) (response gomatrixserverlib.PublicKeyLookupResult, ok bool)
+	StoreServerKey(request gomatrixserverlib.PublicKeyLookupRequest, response gomatrixserverlib.PublicKeyLookupResult)
+}
+
+func (c Caches) GetServerKey(
+	request gomatrixserverlib.PublicKeyLookupRequest,
+) (gomatrixserverlib.PublicKeyLookupResult, bool) {
+	key := fmt.Sprintf("%s/%s", request.ServerName, request.KeyID)
+	val, found := c.ServerKeys.Get(key)
+	if found && val != nil {
+		if keyLookupResult, ok := val.(gomatrixserverlib.PublicKeyLookupResult); ok {
+			return keyLookupResult, true
+		}
+	}
+	return gomatrixserverlib.PublicKeyLookupResult{}, false
+}
+
+func (c Caches) StoreServerKey(
+	request gomatrixserverlib.PublicKeyLookupRequest,
+	response gomatrixserverlib.PublicKeyLookupResult,
+) {
+	key := fmt.Sprintf("%s/%s", request.ServerName, request.KeyID)
+	c.ServerKeys.Set(key, response)
+}
diff --git a/internal/caching/caches.go b/internal/caching/caches.go
new file mode 100644
index 00000000..70f380ba
--- /dev/null
+++ b/internal/caching/caches.go
@@ -0,0 +1,15 @@
+package caching
+
+// Caches contains a set of references to caches. They may be
+// different implementations as long as they satisfy the Cache
+// interface.
+type Caches struct {
+	RoomVersions Cache // implements RoomVersionCache
+	ServerKeys   Cache // implements ServerKeyCache
+}
+
+// Cache is the interface that an implementation must satisfy.
+type Cache interface {
+	Get(key string) (value interface{}, ok bool)
+	Set(key string, value interface{})
+}
diff --git a/internal/caching/immutablecache.go b/internal/caching/immutablecache.go
deleted file mode 100644
index fea05dd1..00000000
--- a/internal/caching/immutablecache.go
+++ /dev/null
@@ -1,17 +0,0 @@
-package caching
-
-import (
-	"github.com/matrix-org/gomatrixserverlib"
-)
-
-const (
-	RoomVersionMaxCacheEntries = 1024
-	ServerKeysMaxCacheEntries  = 1024
-)
-
-type ImmutableCache interface {
-	GetRoomVersion(roomId string) (gomatrixserverlib.RoomVersion, bool)
-	StoreRoomVersion(roomId string, roomVersion gomatrixserverlib.RoomVersion)
-	GetServerKey(request gomatrixserverlib.PublicKeyLookupRequest) (gomatrixserverlib.PublicKeyLookupResult, bool)
-	StoreServerKey(request gomatrixserverlib.PublicKeyLookupRequest, response gomatrixserverlib.PublicKeyLookupResult)
-}
diff --git a/internal/caching/immutableinmemorylru.go b/internal/caching/immutableinmemorylru.go
deleted file mode 100644
index 36cd56dc..00000000
--- a/internal/caching/immutableinmemorylru.go
+++ /dev/null
@@ -1,95 +0,0 @@
-package caching
-
-import (
-	"fmt"
-
-	lru "github.com/hashicorp/golang-lru"
-	"github.com/matrix-org/gomatrixserverlib"
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/promauto"
-)
-
-type ImmutableInMemoryLRUCache struct {
-	roomVersions *lru.Cache
-	serverKeys   *lru.Cache
-}
-
-func NewImmutableInMemoryLRUCache() (*ImmutableInMemoryLRUCache, error) {
-	roomVersionCache, rvErr := lru.New(RoomVersionMaxCacheEntries)
-	if rvErr != nil {
-		return nil, rvErr
-	}
-	serverKeysCache, rvErr := lru.New(ServerKeysMaxCacheEntries)
-	if rvErr != nil {
-		return nil, rvErr
-	}
-	cache := &ImmutableInMemoryLRUCache{
-		roomVersions: roomVersionCache,
-		serverKeys:   serverKeysCache,
-	}
-	cache.configureMetrics()
-	return cache, nil
-}
-
-func (c *ImmutableInMemoryLRUCache) configureMetrics() {
-	promauto.NewGaugeFunc(prometheus.GaugeOpts{
-		Namespace: "dendrite",
-		Subsystem: "caching",
-		Name:      "number_room_version_entries",
-		Help:      "The number of room version entries cached.",
-	}, func() float64 {
-		return float64(c.roomVersions.Len())
-	})
-
-	promauto.NewGaugeFunc(prometheus.GaugeOpts{
-		Namespace: "dendrite",
-		Subsystem: "caching",
-		Name:      "number_server_key_entries",
-		Help:      "The number of server key entries cached.",
-	}, func() float64 {
-		return float64(c.serverKeys.Len())
-	})
-}
-
-func checkForInvalidMutation(cache *lru.Cache, key string, value interface{}) {
-	if peek, ok := cache.Peek(key); ok && peek != value {
-		panic(fmt.Sprintf("invalid use of immutable cache tries to mutate existing value of %q", key))
-	}
-}
-
-func (c *ImmutableInMemoryLRUCache) GetRoomVersion(roomID string) (gomatrixserverlib.RoomVersion, bool) {
-	val, found := c.roomVersions.Get(roomID)
-	if found && val != nil {
-		if roomVersion, ok := val.(gomatrixserverlib.RoomVersion); ok {
-			return roomVersion, true
-		}
-	}
-	return "", false
-}
-
-func (c *ImmutableInMemoryLRUCache) StoreRoomVersion(roomID string, roomVersion gomatrixserverlib.RoomVersion) {
-	checkForInvalidMutation(c.roomVersions, roomID, roomVersion)
-	c.roomVersions.Add(roomID, roomVersion)
-}
-
-func (c *ImmutableInMemoryLRUCache) GetServerKey(
-	request gomatrixserverlib.PublicKeyLookupRequest,
-) (gomatrixserverlib.PublicKeyLookupResult, bool) {
-	key := fmt.Sprintf("%s/%s", request.ServerName, request.KeyID)
-	val, found := c.serverKeys.Get(key)
-	if found && val != nil {
-		if keyLookupResult, ok := val.(gomatrixserverlib.PublicKeyLookupResult); ok {
-			return keyLookupResult, true
-		}
-	}
-	return gomatrixserverlib.PublicKeyLookupResult{}, false
-}
-
-func (c *ImmutableInMemoryLRUCache) StoreServerKey(
-	request gomatrixserverlib.PublicKeyLookupRequest,
-	response gomatrixserverlib.PublicKeyLookupResult,
-) {
-	key := fmt.Sprintf("%s/%s", request.ServerName, request.KeyID)
-	checkForInvalidMutation(c.roomVersions, key, response)
-	c.serverKeys.Add(request, response)
-}
diff --git a/internal/caching/impl_inmemorylru.go b/internal/caching/impl_inmemorylru.go
new file mode 100644
index 00000000..f7901d2e
--- /dev/null
+++ b/internal/caching/impl_inmemorylru.go
@@ -0,0 +1,73 @@
+package caching
+
+import (
+	"fmt"
+
+	lru "github.com/hashicorp/golang-lru"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+func NewInMemoryLRUCache() (*Caches, error) {
+	roomVersions, err := NewInMemoryLRUCachePartition(
+		RoomVersionCacheName,
+		RoomVersionCacheMutable,
+		RoomVersionCacheMaxEntries,
+	)
+	if err != nil {
+		return nil, err
+	}
+	serverKeys, err := NewInMemoryLRUCachePartition(
+		ServerKeyCacheName,
+		ServerKeyCacheMutable,
+		ServerKeyCacheMaxEntries,
+	)
+	if err != nil {
+		return nil, err
+	}
+	return &Caches{
+		RoomVersions: roomVersions,
+		ServerKeys:   serverKeys,
+	}, nil
+}
+
+type InMemoryLRUCachePartition struct {
+	name       string
+	mutable    bool
+	maxEntries int
+	lru        *lru.Cache
+}
+
+func NewInMemoryLRUCachePartition(name string, mutable bool, maxEntries int) (*InMemoryLRUCachePartition, error) {
+	var err error
+	cache := InMemoryLRUCachePartition{
+		name:       name,
+		mutable:    mutable,
+		maxEntries: maxEntries,
+	}
+	cache.lru, err = lru.New(maxEntries)
+	if err != nil {
+		return nil, err
+	}
+	promauto.NewGaugeFunc(prometheus.GaugeOpts{
+		Namespace: "dendrite",
+		Subsystem: "caching_in_memory_lru",
+		Name:      name,
+	}, func() float64 {
+		return float64(cache.lru.Len())
+	})
+	return &cache, nil
+}
+
+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))
+		}
+	}
+	c.lru.Add(key, value)
+}
+
+func (c *InMemoryLRUCachePartition) Get(key string) (value interface{}, ok bool) {
+	return c.lru.Get(key)
+}