Add LRU cache for some DB operations

2 files changed, 83 insertions, 27 deletions

diff --git a/db.go b/db.go
index a004496..da84c08 100644
--- a/db.go
+++ b/db.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/hashicorp/golang-lru"
 	"github.com/jinzhu/gorm"
 	_ "github.com/jinzhu/gorm/dialects/mysql"
 	_ "github.com/jinzhu/gorm/dialects/postgres"
@@ -16,6 +17,7 @@ import (
 )
 
 var db *gorm.DB
+var dbCache *lru.TwoQueueCache
 
 func InitDb() error {
 	var err error
@@ -41,6 +43,11 @@ func InitDb() error {
 	db.AutoMigrate(&DbJoinedRoom{})
 	db.Model(&DbJoinedRoom{}).AddIndex("idx_user_protocol_account", "mx_user_id", "protocol", "account_name")
 
+	dbCache, err = lru.New2Q(10000)
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -54,12 +61,17 @@ type DbAccountConfig struct {
 	Config   string
 }
 
-// Long-term cache entries
-type DbKv struct {
+// List of joined channels to be re-joined on reconnect
+type DbJoinedRoom struct {
 	gorm.Model
 
-	Key   string `gorm:"unique_index"`
-	Value string
+	// User id and account name
+	MxUserID    string
+	Protocol    string
+	AccountName string
+
+	// Room ID
+	RoomID connector.RoomID
 }
 
 // User mapping between protocol user IDs and puppeted matrix ids
@@ -101,20 +113,18 @@ type DbPmRoomMap struct {
 	MxRoomID string `gorm:"index:mxroomoid"`
 }
 
-// List of joined channels to be re-joined on reconnect
-type DbJoinedRoom struct {
+// Key-value store for various things
+type DbKv struct {
 	gorm.Model
 
-	// User id and account name
-	MxUserID    string
-	Protocol    string
-	AccountName string
-
-	// Room ID
-	RoomID connector.RoomID
+	Key   string `gorm:"unique_index"`
+	Value string
 }
 
 // ---- Simple locking mechanism
+// Slot keys are strings that identify the object we are acting upon
+// They define which lock to lock for a certain operation
+// They are also used as keys in the LRU cache
 
 var dbLocks [256]sync.Mutex
 
@@ -128,9 +138,19 @@ func dbUnlockSlot(key string) {
 	dbLocks[slot].Unlock()
 }
 
-// ----
+// ---- Key-value store supporting atomic test-and-set
+
+func dbKvSlotKey(key string) string {
+	return "kv:" + key
+}
 
 func dbKvGet(key string) string {
+	slot_key := dbKvSlotKey(key)
+
+	if ent, ok := dbCache.Get(slot_key); ok {
+		return ent.(string)
+	}
+
 	var entry DbKv
 	if db.Where(&DbKv{Key: key}).First(&entry).RecordNotFound() {
 		return ""
@@ -140,35 +160,53 @@ func dbKvGet(key string) string {
 }
 
 func dbKvPut(key string, value string) {
+	slot_key := dbKvSlotKey(key)
+
+	dbLockSlot(slot_key)
+	defer dbUnlockSlot(slot_key)
+
 	var entry DbKv
 	db.Where(&DbKv{Key: key}).Assign(&DbKv{Value: value}).FirstOrCreate(&entry)
+	dbCache.Add(slot_key, value)
 }
 
 func dbKvTestAndSet(key string, value string) bool {
-	dbLockSlot(key)
-	defer dbUnlockSlot(key)
+	slot_key := dbKvSlotKey(key)
+
+	dbLockSlot(slot_key)
+	defer dbUnlockSlot(slot_key)
 
 	// True if value was changed, false if was already set
-	if dbKvGet(key) != value {
-		dbKvPut(key, value)
-		return true
+	if dbKvGet(key) == value {
+		return false
 	}
-	return false
+
+	var entry DbKv
+	db.Where(&DbKv{Key: key}).Assign(&DbKv{Value: value}).FirstOrCreate(&entry)
+	dbCache.Add(slot_key, value)
+	return true
 }
 
+// ----
+
 func dbGetMxRoom(protocol string, roomId connector.RoomID) (string, error) {
-	slot_key := fmt.Sprintf("room: %s / %s", protocol, roomId)
+	slot_key := fmt.Sprintf("room:%s/%s", protocol, roomId)
+
 	dbLockSlot(slot_key)
 	defer dbUnlockSlot(slot_key)
 
-	var room DbRoomMap
+	if cached, ok := dbCache.Get(slot_key); ok {
+		return cached.(string), nil
+	}
 
 	// Check if room exists in our mapping,
 	// If not create it
+	var room DbRoomMap
 	must_create := db.First(&room, DbRoomMap{
 		Protocol: protocol,
 		RoomID:   roomId,
 	}).RecordNotFound()
+
 	if must_create {
 		alias := roomAlias(protocol, roomId)
 		// Lookup alias
@@ -192,24 +230,31 @@ func dbGetMxRoom(protocol string, roomId connector.RoomID) (string, error) {
 		}
 		db.Create(&room)
 	}
-	log.Tracef("Got room id for %s %s: %s", protocol, roomId, room.MxRoomID)
+
+	log.Tracef("%s -> %s", slot_key, room.MxRoomID)
+	dbCache.Add(slot_key, room.MxRoomID)
 
 	return room.MxRoomID, nil
 }
 
 func dbGetMxPmRoom(protocol string, them connector.UserID, themMxId string, usMxId string, usAccount string) (string, error) {
-	slot_key := fmt.Sprintf("pmroom: %s / %s / %s / %s", protocol, usMxId, usAccount, them)
+	slot_key := fmt.Sprintf("pmroom:%s/%s/%s/%s", protocol, usMxId, usAccount, them)
+
 	dbLockSlot(slot_key)
 	defer dbUnlockSlot(slot_key)
 
-	var room DbPmRoomMap
+	if cached, ok := dbCache.Get(slot_key); ok {
+		return cached.(string), nil
+	}
 
+	var room DbPmRoomMap
 	must_create := db.First(&room, DbPmRoomMap{
 		MxUserID:    usMxId,
 		Protocol:    protocol,
 		AccountName: usAccount,
 		UserID:      them,
 	}).RecordNotFound()
+
 	if must_create {
 		name := fmt.Sprintf("%s (%s)", them, protocol)
 
@@ -234,16 +279,23 @@ func dbGetMxPmRoom(protocol string, them connector.UserID, themMxId string, usMx
 		}
 		db.Create(&room)
 	}
-	log.Tracef("Got PM room id for %s %s %s %s: %s", usMxId, protocol, usAccount, them, room.MxRoomID)
+
+	log.Tracef("%s -> %s", slot_key, room.MxRoomID)
+	dbCache.Add(slot_key, room.MxRoomID)
 
 	return room.MxRoomID, nil
 }
 
 func dbGetMxUser(protocol string, userId connector.UserID) (string, error) {
-	slot_key := fmt.Sprintf("user: %s / %s", protocol, userId)
+	slot_key := fmt.Sprintf("user:%s/%s", protocol, userId)
+
 	dbLockSlot(slot_key)
 	defer dbUnlockSlot(slot_key)
 
+	if cached, ok := dbCache.Get(slot_key); ok {
+		return cached.(string), nil
+	}
+
 	var user DbUserMap
 
 	must_create := db.First(&user, DbUserMap{
@@ -272,6 +324,9 @@ func dbGetMxUser(protocol string, userId connector.UserID) (string, error) {
 		db.Create(&user)
 	}
 
+	log.Tracef("%s -> %s", slot_key, user.MxUserID)
+	dbCache.Add(slot_key, user.MxUserID)
+
 	return user.MxUserID, nil
 }
 
diff --git a/go.mod b/go.mod
index e44bd99..452fe93 100644
--- a/go.mod
+++ b/go.mod
@@ -10,6 +10,7 @@ require (
 	github.com/go-ldap/ldap/v3 v3.1.7
 	github.com/gorilla/mux v1.7.4
 	github.com/gorilla/sessions v1.2.0
+	github.com/hashicorp/golang-lru v0.5.3
 	github.com/jinzhu/gorm v1.9.12
 	github.com/lrstanley/girc v0.0.0-20190801035559-4fc93959e1a7
 	github.com/matterbridge/go-xmpp v0.0.0-20180529212104-cd19799fba91