From 8642c0b2fa97dcff012a6e7da972e85c732170ec Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 22 Apr 2020 10:05:44 +0100
Subject: [PATCH] Experimental LRU cache for room versions

---
 common/caching/cache.go       |  8 +++++++
 common/caching/inmemorylru.go | 41 +++++++++++++++++++++++++++++++++++
 go.mod                        |  1 +
 roomserver/query/query.go     |  9 +++++++-
 roomserver/roomserver.go      |  6 ++++-
 5 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 common/caching/cache.go
 create mode 100644 common/caching/inmemorylru.go

diff --git a/common/caching/cache.go b/common/caching/cache.go
new file mode 100644
index 000000000..09edb44f9
--- /dev/null
+++ b/common/caching/cache.go
@@ -0,0 +1,8 @@
+package caching
+
+import "github.com/matrix-org/gomatrixserverlib"
+
+type Cache interface {
+	GetRoomVersion(roomId string) (gomatrixserverlib.RoomVersion, bool)
+	StoreRoomVersion(roomId string, roomVersion gomatrixserverlib.RoomVersion)
+}
diff --git a/common/caching/inmemorylru.go b/common/caching/inmemorylru.go
new file mode 100644
index 000000000..94fac72f9
--- /dev/null
+++ b/common/caching/inmemorylru.go
@@ -0,0 +1,41 @@
+package caching
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/golang/groupcache/lru"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+type InMemoryLRUCache struct {
+	roomVersions      *lru.Cache
+	roomVersionsMutex sync.RWMutex
+}
+
+func NewInMemoryLRUCache() *InMemoryLRUCache {
+	return &InMemoryLRUCache{
+		roomVersions: lru.New(128),
+	}
+}
+
+func (c *InMemoryLRUCache) GetRoomVersion(roomID string) (gomatrixserverlib.RoomVersion, bool) {
+	fmt.Println("Cache hit for", roomID)
+	if c == nil {
+		return "", false
+	}
+	c.roomVersionsMutex.RLock()
+	defer c.roomVersionsMutex.RUnlock()
+	val, found := c.roomVersions.Get(roomID)
+	return val.(gomatrixserverlib.RoomVersion), found
+}
+
+func (c *InMemoryLRUCache) StoreRoomVersion(roomID string, roomVersion gomatrixserverlib.RoomVersion) {
+	fmt.Println("Cache store for", roomID)
+	if c == nil {
+		return
+	}
+	c.roomVersionsMutex.Lock()
+	defer c.roomVersionsMutex.Unlock()
+	c.roomVersions.Add(roomID, roomVersion)
+}
diff --git a/go.mod b/go.mod
index 8d91902d1..5d9c6e58d 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,7 @@
 module github.com/matrix-org/dendrite
 
 require (
+	github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6
 	github.com/gorilla/mux v1.7.3
 	github.com/hashicorp/golang-lru v0.5.4
 	github.com/kr/pretty v0.2.0 // indirect
diff --git a/roomserver/query/query.go b/roomserver/query/query.go
index 12d8436ef..db89766c0 100644
--- a/roomserver/query/query.go
+++ b/roomserver/query/query.go
@@ -22,6 +22,7 @@ import (
 	"net/http"
 
 	"github.com/matrix-org/dendrite/common"
+	"github.com/matrix-org/dendrite/common/caching"
 	"github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/matrix-org/dendrite/roomserver/auth"
 	"github.com/matrix-org/dendrite/roomserver/state"
@@ -97,7 +98,8 @@ type RoomserverQueryAPIDatabase interface {
 
 // RoomserverQueryAPI is an implementation of api.RoomserverQueryAPI
 type RoomserverQueryAPI struct {
-	DB RoomserverQueryAPIDatabase
+	DB    RoomserverQueryAPIDatabase
+	Cache caching.Cache
 }
 
 // QueryLatestEventsAndState implements api.RoomserverQueryAPI
@@ -896,11 +898,16 @@ func (r *RoomserverQueryAPI) QueryRoomVersionForRoom(
 	request *api.QueryRoomVersionForRoomRequest,
 	response *api.QueryRoomVersionForRoomResponse,
 ) error {
+	if roomVersion, ok := r.Cache.GetRoomVersion(request.RoomID); ok {
+		response.RoomVersion = roomVersion
+		return nil
+	}
 	roomVersion, err := r.DB.GetRoomVersionForRoom(ctx, request.RoomID)
 	if err != nil {
 		return err
 	}
 	response.RoomVersion = roomVersion
+	r.Cache.StoreRoomVersion(request.RoomID, response.RoomVersion)
 	return nil
 }
 
diff --git a/roomserver/roomserver.go b/roomserver/roomserver.go
index 2ffbf67de..3af40777e 100644
--- a/roomserver/roomserver.go
+++ b/roomserver/roomserver.go
@@ -21,6 +21,7 @@ import (
 
 	asQuery "github.com/matrix-org/dendrite/appservice/query"
 	"github.com/matrix-org/dendrite/common/basecomponent"
+	"github.com/matrix-org/dendrite/common/caching"
 	"github.com/matrix-org/dendrite/roomserver/alias"
 	"github.com/matrix-org/dendrite/roomserver/input"
 	"github.com/matrix-org/dendrite/roomserver/query"
@@ -48,7 +49,10 @@ func SetupRoomServerComponent(
 
 	inputAPI.SetupHTTP(http.DefaultServeMux)
 
-	queryAPI := query.RoomserverQueryAPI{DB: roomserverDB}
+	queryAPI := query.RoomserverQueryAPI{
+		DB:    roomserverDB,
+		Cache: caching.NewInMemoryLRUCache(),
+	}
 
 	queryAPI.SetupHTTP(http.DefaultServeMux)