Use gomatrixserverlib key caching

cmd/dendrite-client-api-server/main.go

@@ -33,7 +33,7 @@ func main() {
 	deviceDB := base.CreateDeviceDB()
 	keyDB := base.CreateKeyDB()
 	federation := base.CreateFederationClient()
-	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, cfg.Matrix.KeyPerspectives)
+	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, base.ImmutableCache, cfg.Matrix.KeyPerspectives)
 
 	asQuery := base.CreateHTTPAppServiceAPIs()
 	rsAPI := base.CreateHTTPRoomserverAPIs()

cmd/dendrite-demo-libp2p/main.go

@@ -147,7 +147,8 @@ func main() {
 	deviceDB := base.Base.CreateDeviceDB()
 	keyDB := createKeyDB(base)
 	federation := createFederationClient(base)
-	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, cfg.Matrix.KeyPerspectives)
+	// TODO: cache here
+	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, nil, cfg.Matrix.KeyPerspectives)
 
 	rsAPI := roomserver.SetupRoomServerComponent(
 		&base.Base, keyRing, federation,

cmd/dendrite-federation-api-server/main.go

@@ -33,7 +33,7 @@ func main() {
 	keyDB := base.CreateKeyDB()
 	federation := base.CreateFederationClient()
 	fsAPI := base.CreateHTTPFederationSenderAPIs()
-	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, cfg.Matrix.KeyPerspectives)
+	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, base.ImmutableCache, cfg.Matrix.KeyPerspectives)
 
 	rsAPI := base.CreateHTTPRoomserverAPIs()
 	asAPI := base.CreateHTTPAppServiceAPIs()

cmd/dendrite-federation-sender-server/main.go

@@ -27,7 +27,7 @@ func main() {
 
 	federation := base.CreateFederationClient()
 	keyDB := base.CreateKeyDB()
-	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, cfg.Matrix.KeyPerspectives)
+	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, base.ImmutableCache, cfg.Matrix.KeyPerspectives)
 	rsAPI := base.CreateHTTPRoomserverAPIs()
 	fsAPI := federationsender.SetupFederationSenderComponent(
 		base, federation, rsAPI, &keyRing,

cmd/dendrite-monolith-server/main.go

@@ -56,7 +56,7 @@ func main() {
 	deviceDB := base.CreateDeviceDB()
 	keyDB := base.CreateKeyDB()
 	federation := base.CreateFederationClient()
-	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, cfg.Matrix.KeyPerspectives)
+	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, base.ImmutableCache, cfg.Matrix.KeyPerspectives)
 
 	rsAPI := roomserver.SetupRoomServerComponent(
 		base, keyRing, federation,

cmd/dendrite-room-server/main.go

@@ -26,7 +26,7 @@ func main() {
 	defer base.Close() // nolint: errcheck
 	keyDB := base.CreateKeyDB()
 	federation := base.CreateFederationClient()
-	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, cfg.Matrix.KeyPerspectives)
+	keyRing := keydb.CreateKeyRing(federation.Client, keyDB, base.ImmutableCache, cfg.Matrix.KeyPerspectives)
 
 	fsAPI := base.CreateHTTPFederationSenderAPIs()
 	rsAPI := roomserver.SetupRoomServerComponent(base, keyRing, federation)

common/caching/immutablecache.go

@@ -4,9 +4,11 @@ import "github.com/matrix-org/gomatrixserverlib"
 
 const (
 	RoomVersionMaxCacheEntries = 128
+	ServerKeysMaxCacheEntries  = 128
 )
 
 type ImmutableCache interface {
+	gomatrixserverlib.KeyCache
 	GetRoomVersion(roomId string) (gomatrixserverlib.RoomVersion, bool)
 	StoreRoomVersion(roomId string, roomVersion gomatrixserverlib.RoomVersion)
 }

common/caching/immutableinmemorylru.go

@@ -9,6 +9,7 @@ import (
 
 type ImmutableInMemoryLRUCache struct {
 	roomVersions *lru.Cache
+	serverKeys   *lru.Cache
 }
 
 func NewImmutableInMemoryLRUCache() (*ImmutableInMemoryLRUCache, error) {
@@ -16,8 +17,13 @@ func NewImmutableInMemoryLRUCache() (*ImmutableInMemoryLRUCache, error) {
 	if rvErr != nil {
 		return nil, rvErr
 	}
+	serverKeysCache, rvErr := lru.New(ServerKeysMaxCacheEntries)
+	if rvErr != nil {
+		return nil, rvErr
+	}
 	return &ImmutableInMemoryLRUCache{
 		roomVersions: roomVersionCache,
+		serverKeys:   serverKeysCache,
 	}, nil
 }
 
@@ -41,3 +47,25 @@ func (c *ImmutableInMemoryLRUCache) StoreRoomVersion(roomID string, 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)
+}

common/keydb/keyring.go

@@ -17,6 +17,7 @@ package keydb
 import (
 	"encoding/base64"
 
+	"github.com/matrix-org/dendrite/common/caching"
 	"github.com/matrix-org/dendrite/common/config"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/sirupsen/logrus"
@@ -29,6 +30,7 @@ import (
 // backed by the given KeyDatabase.
 func CreateKeyRing(client gomatrixserverlib.Client,
 	keyDB gomatrixserverlib.KeyDatabase,
+	cache caching.ImmutableCache,
 	cfg config.KeyPerspectives) gomatrixserverlib.KeyRing {
 
 	fetchers := gomatrixserverlib.KeyRing{
@@ -38,6 +40,7 @@ func CreateKeyRing(client gomatrixserverlib.Client,
 			},
 		},
 		KeyDatabase: keyDB,
+		KeyCache:    cache,
 	}
 
 	logrus.Info("Enabled direct key fetcher")

common/keydb/sqlite3/server_key_table.go

@@ -20,10 +20,8 @@ import (
 	"database/sql"
 	"strings"
 
-	lru "github.com/hashicorp/golang-lru"
 	"github.com/matrix-org/dendrite/common"
 	"github.com/matrix-org/gomatrixserverlib"
-	"github.com/matrix-org/util"
 )
 
 const serverKeysSchema = `
@@ -66,16 +64,10 @@ type serverKeyStatements struct {
 	db                       *sql.DB
 	bulkSelectServerKeysStmt *sql.Stmt
 	upsertServerKeysStmt     *sql.Stmt
-
-	cache *lru.Cache // nameAndKeyID => gomatrixserverlib.PublicKeyLookupResult
 }
 
 func (s *serverKeyStatements) prepare(db *sql.DB) (err error) {
 	s.db = db
-	s.cache, err = lru.New(64)
-	if err != nil {
-		return
-	}
 	_, err = db.Exec(serverKeysSchema)
 	if err != nil {
 		return
@@ -98,21 +90,6 @@ func (s *serverKeyStatements) bulkSelectServerKeys(
 		nameAndKeyIDs = append(nameAndKeyIDs, nameAndKeyID(request))
 	}
 
-	// If we can satisfy all of the requests from the cache, do so. TODO: Allow partial matches with merges.
-	cacheResults := map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult{}
-	for request := range requests {
-		r, ok := s.cache.Get(nameAndKeyID(request))
-		if !ok {
-			break
-		}
-		cacheResult := r.(gomatrixserverlib.PublicKeyLookupResult)
-		cacheResults[request] = cacheResult
-	}
-	if len(cacheResults) == len(requests) {
-		util.GetLogger(ctx).Infof("KeyDB cache hit for %d keys", len(cacheResults))
-		return cacheResults, nil
-	}
-
 	query := strings.Replace(bulkSelectServerKeysSQL, "($1)", common.QueryVariadic(len(nameAndKeyIDs)), 1)
 
 	iKeyIDs := make([]interface{}, len(nameAndKeyIDs))
@@ -158,7 +135,6 @@ func (s *serverKeyStatements) upsertServerKeys(
 	request gomatrixserverlib.PublicKeyLookupRequest,
 	key gomatrixserverlib.PublicKeyLookupResult,
 ) error {
-	s.cache.Add(nameAndKeyID(request), key)
 	_, err := s.upsertServerKeysStmt.ExecContext(
 		ctx,
 		string(request.ServerName),

go.mod

@@ -17,7 +17,7 @@ require (
 	github.com/matrix-org/go-http-js-libp2p v0.0.0-20200318135427-31631a9ef51f
 	github.com/matrix-org/go-sqlite3-js v0.0.0-20200325174927-327088cdef10
 	github.com/matrix-org/gomatrix v0.0.0-20190528120928-7df988a63f26
-	github.com/matrix-org/gomatrixserverlib v0.0.0-20200511154227-5cc71d36632b
+	github.com/matrix-org/gomatrixserverlib v0.0.0-20200514163156-a69606a3855c
 	github.com/matrix-org/naffka v0.0.0-20200422140631-181f1ee7401f
 	github.com/matrix-org/util v0.0.0-20190711121626-527ce5ddefc7
 	github.com/mattn/go-sqlite3 v2.0.2+incompatible

go.sum

@@ -369,6 +369,8 @@ github.com/matrix-org/gomatrixserverlib v0.0.0-20200124100636-0c2ec91d1df5 h1:km
 github.com/matrix-org/gomatrixserverlib v0.0.0-20200124100636-0c2ec91d1df5/go.mod h1:FsKa2pWE/bpQql9H7U4boOPXFoJX/QcqaZZ6ijLkaZI=
 github.com/matrix-org/gomatrixserverlib v0.0.0-20200511154227-5cc71d36632b h1:nAmSc1KvQOumoRTz/LD68KyrB6Q5/6q7CmQ5Bswc2nM=
 github.com/matrix-org/gomatrixserverlib v0.0.0-20200511154227-5cc71d36632b/go.mod h1:JsAzE1Ll3+gDWS9JSUHPJiiyAksvOOnGWF2nXdg4ZzU=
+github.com/matrix-org/gomatrixserverlib v0.0.0-20200514163156-a69606a3855c h1:S0+PWkxoUl1/wA1o6OcCXQGpjc64E/sQB+QU/6VCGMc=
+github.com/matrix-org/gomatrixserverlib v0.0.0-20200514163156-a69606a3855c/go.mod h1:JsAzE1Ll3+gDWS9JSUHPJiiyAksvOOnGWF2nXdg4ZzU=
 github.com/matrix-org/naffka v0.0.0-20200127221512-0716baaabaf1 h1:osLoFdOy+ChQqVUn2PeTDETFftVkl4w9t/OW18g3lnk=
 github.com/matrix-org/naffka v0.0.0-20200127221512-0716baaabaf1/go.mod h1:cXoYQIENbdWIQHt1SyCo6Bl3C3raHwJ0wgVrXHSqf+A=
 github.com/matrix-org/naffka v0.0.0-20200422140631-181f1ee7401f h1:pRz4VTiRCO4zPlEMc3ESdUOcW4PXHH4Kj+YDz1XyE+Y=