Add LoginSession cache

Author: JojiiOfficial

models/LoginSessions.go

@@ -1,6 +1,8 @@
 package models
 
 import (
+	"sync"
+
 	"github.com/JojiiOfficial/gaw"
 	log "github.com/sirupsen/logrus"
 	"gorm.io/gorm"
@@ -14,30 +16,58 @@ type LoginSession struct {
 	Token     string
 	Requests  int64
 	MachineID string
+
+	mx sync.Mutex `gorm:"-"`
 }
 
 // SessionTokenLength length of session token
 const SessionTokenLength = 64
 
 //GetUserFromSession return user from session
 func GetUserFromSession(db *gorm.DB, token string) (*User, error) {
-	var session LoginSession
-	err := db.Model(&LoginSession{}).Where(&LoginSession{
-		Token: token,
-	}).Preload("User").Preload("User.Role").Find(&session).Error
+	var err error
 
-	if err != nil {
-		return nil, err
+	// Try to get session from cache
+	session := sessionCache.getSession(token)
+	if session == nil {
+		// Load from DB if not cached
+		session, err = loadSession(token, db)
+		if err != nil {
+			return nil, err
+		}
+
+		// Add token to cache
+		sessionCache.addSession(session, db)
 	}
 
 	// Increase request counter
-	session.Requests++
-	err = db.Save(&session).Error
+	go func() {
+		session.mx.Lock()
+		defer session.mx.Unlock()
+
+		session.Requests++
+		err = db.Model(&LoginSession{}).Save(session).Error
+		if err != nil {
+			log.Error(err)
+		}
+	}()
+
+	return session.User, nil
+}
+
+func loadSession(token string, db *gorm.DB) (*LoginSession, error) {
+	var ss LoginSession
+
+	// load session from db
+	err := db.Model(&LoginSession{}).Where(&LoginSession{
+		Token: token,
+	}).Preload("User").Preload("User.Role").Find(&ss).Error
+
 	if err != nil {
-		log.Error(err)
+		return nil, err
 	}
 
-	return session.User, nil
+	return &ss, nil
 }
 
 // NewSession create new login session

models/SessionCache.go

@@ -0,0 +1,123 @@
+package models
+
+import (
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+)
+
+// Cache sessions
+var sessionCache SessionCache
+
+// SessionCache cache sessions in ram
+type SessionCache struct {
+	cache map[string]*sessionCacheEntry
+
+	mx sync.Mutex
+}
+
+type sessionCacheEntry struct {
+	db      *gorm.DB
+	session *LoginSession
+	valid   bool
+
+	lastCheck              int64
+	requestsSinceLastCheck uint32
+
+	mx sync.Mutex
+}
+
+// init the Cache
+func (sc *SessionCache) init() {
+	sc.mx.Lock()
+	defer sc.mx.Unlock()
+
+	if sc.cache == nil {
+		sc.cache = make(map[string]*sessionCacheEntry)
+	}
+}
+
+// Get Session from Cache
+func (sc *SessionCache) getSession(token string) *LoginSession {
+	sc.init()
+
+	sc.mx.Lock()
+	defer sc.mx.Unlock()
+
+	// Get cache entry
+	session, ok := sc.cache[token]
+	if !ok {
+		return nil
+	}
+
+	// Update request counter
+	session.requestsSinceLastCheck++
+
+	// Remove session from
+	// cache if invalid
+	if !session.valid {
+		delete(sc.cache, token)
+		return nil
+	}
+
+	// Update if neccessary
+	if session.needUpdate() {
+		if !session.update() {
+			return nil
+		}
+	}
+
+	return session.session
+}
+
+// Add LoginSession to cache
+func (sc *SessionCache) addSession(session *LoginSession, db *gorm.DB) {
+	sc.init()
+
+	sc.mx.Lock()
+	defer sc.mx.Unlock()
+
+	sc.cache[session.Token] = &sessionCacheEntry{
+		db:        db,
+		valid:     true,
+		session:   session,
+		lastCheck: time.Now().Unix(),
+	}
+}
+
+func (sce *sessionCacheEntry) update() bool {
+	sce.mx.Lock()
+	defer sce.mx.Unlock()
+
+	// Reload LoginSession frm DB
+	session, err := loadSession(sce.session.Token, sce.db)
+	if err != nil {
+		log.Warn("Session invalid: ", err)
+		sce.valid = false
+		return false
+	}
+
+	// Update cache entry
+	sce.session = session
+	sce.valid = true
+	sce.lastCheck = time.Now().Unix()
+	sce.requestsSinceLastCheck = 0
+
+	return true
+}
+
+// Allow 60s cache
+const (
+	maxCacheValidTime     = int64(10)
+	maxRequestsUntilCheck = uint32(10)
+)
+
+// Return true if session needs update from DB
+func (sce *sessionCacheEntry) needUpdate() bool {
+	sce.mx.Lock()
+	defer sce.mx.Unlock()
+
+	return time.Now().Unix()-maxCacheValidTime > sce.lastCheck || sce.requestsSinceLastCheck > maxRequestsUntilCheck
+}

storage/Database.go

@@ -21,6 +21,8 @@ func ConnectToDatabase(config *models.Config) (*gorm.DB, error) {
 		return nil, err
 	}
 
+	//	db = db.Debug()
+
 	//Automigration
 	db.AutoMigrate(
 		&models.Role{},