improve cache handling

database/migrations/20231102165658_tables.sql

@@ -0,0 +1,20 @@
+-- +goose Up
+-- +goose StatementBegin
+
+ALTER TABLE teldrive.users DROP COLUMN IF EXISTS tg_session;
+
+CREATE TABLE teldrive.sessions (
+    session text NOT NULL,
+    user_id bigint NOT NULL,
+    hash text NOT NULL,
+    created_at timestamp null default timezone('utc'::text,now()),
+    PRIMARY KEY(session, hash),
+    FOREIGN KEY (user_id) REFERENCES teldrive.users(user_id)
+);
+
+-- +goose StatementEnd
+
+-- +goose Down
+-- +goose StatementBegin
+DROP TABLE IF EXISTS teldrive.sessions;
+-- +goose StatementEnd

go.mod

@@ -3,6 +3,7 @@ module github.com/divyam234/teldrive
 go 1.21
 
 require (
+	github.com/coocood/freecache v1.2.4
 	github.com/divyam234/cors v1.4.2
 	github.com/gin-gonic/gin v1.9.1
 	github.com/go-co-op/gocron v1.35.3
@@ -24,6 +25,7 @@ require (
 )
 
 require (
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/chenzhuoyu/iasm v0.9.0 // indirect
 	github.com/google/uuid v1.4.0 // indirect
 	github.com/robfig/cron/v3 v3.0.1 // indirect

go.sum

@@ -4,12 +4,17 @@ github.com/bytedance/sonic v1.10.0 h1:qtNZduETEIWJVIyDl01BeNxur2rW9OwTQ/yBqFRkKE
 github.com/bytedance/sonic v1.10.0/go.mod h1:iZcSUejdk5aukTND/Eu/ivjQuEL0Cu9/rf50Hi0u/g4=
 github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqylYbM=
 github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
+github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
+github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY=
 github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583jCggY9gE99b6G5LEC39OIiVsWj+R97kbl5odCEk=
 github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d h1:77cEq6EriyTZ0g/qfRdp61a3Uu/AWrgIq2s0ClJV1g0=
 github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d/go.mod h1:8EPpVsBuRksnlj1mLy4AWzRNQYxauNi62uWcE3to6eA=
 github.com/chenzhuoyu/iasm v0.9.0 h1:9fhXjVzq5hUy2gkhhgHl95zG2cEAhw9OSGs8toWWAwo=
 github.com/chenzhuoyu/iasm v0.9.0/go.mod h1:Xjy2NpN3h7aUqeqM+woSuuvxmIe6+DDsiNLIrkAmYog=
+github.com/coocood/freecache v1.2.4 h1:UdR6Yz/X1HW4fZOuH0Z94KwG851GWOSknua5VUbb/5M=
+github.com/coocood/freecache v1.2.4/go.mod h1:RBUWa/Cy+OHdfTGFEhEuE1pMCMX51Ncizj7rthiQ3vk=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=

main.go

@@ -12,6 +12,7 @@ import (
 	"github.com/divyam234/teldrive/utils"
 
 	"github.com/divyam234/cors"
+	"github.com/divyam234/teldrive/utils/cache"
 	"github.com/divyam234/teldrive/utils/cron"
 	"github.com/gin-gonic/gin"
 	"github.com/go-co-op/gocron"
@@ -29,6 +30,8 @@ func main() {
 
 	database.InitDB()
 
+	cache.InitCache()
+
 	scheduler := gocron.NewScheduler(time.UTC)
 
 	scheduler.Every(1).Hour().Do(cron.FilesDeleteJob)

models/session.model.go

@@ -0,0 +1,12 @@
+package models
+
+import (
+	"time"
+)
+
+type Session struct {
+	UserId    int64     `gorm:"type:bigint;primaryKey"`
+	Hash      string    `gorm:"type:text"`
+	Session   string    `gorm:"type:text"`
+	CreatedAt time.Time `gorm:"default:timezone('utc'::text, now())"`
+}

routes/user.go

@@ -96,15 +96,4 @@ func addUserRoutes(rg *gin.RouterGroup) {
 
 		c.JSON(http.StatusOK, res)
 	})
-
-	r.DELETE("/cache", func(c *gin.Context) {
-		res, err := userService.ClearCache(c)
-
-		if err != nil {
-			c.AbortWithError(err.Code, err.Error)
-			return
-		}
-
-		c.JSON(http.StatusOK, res)
-	})
 }

services/auth.service.go

@@ -16,13 +16,11 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/divyam234/teldrive/database"
 	"github.com/divyam234/teldrive/models"
 	"github.com/divyam234/teldrive/schemas"
 	"github.com/divyam234/teldrive/types"
 	"github.com/divyam234/teldrive/utils"
 	"github.com/divyam234/teldrive/utils/auth"
-	"github.com/divyam234/teldrive/utils/kv"
 	"github.com/divyam234/teldrive/utils/tgc"
 	"github.com/gin-gonic/gin"
 	"github.com/go-jose/go-jose/v3/jwt"
@@ -145,9 +143,8 @@ func (as *AuthService) LogIn(c *gin.Context) (*schemas.Message, *types.AppError)
 		IsPremium: session.IsPremium,
 	}
 
-	tokenBytes, _ := json.Marshal(jwtClaims)
-	md5hash := md5.Sum(tokenBytes)
-	hexToken := hex.EncodeToString(md5hash[:])
+	tokenhash := md5.Sum([]byte(session.Sesssion))
+	hexToken := hex.EncodeToString(tokenhash[:])
 	jwtClaims.Hash = hexToken
 
 	jweToken, err := auth.Encode(jwtClaims)
@@ -168,7 +165,7 @@ func (as *AuthService) LogIn(c *gin.Context) (*schemas.Message, *types.AppError)
 
 	if err := as.Db.Model(&models.User{}).Where("user_id = ?", session.UserID).
 		Find(&result).Error; err != nil {
-		return nil, &types.AppError{Error: errors.New("failed to create or update user"),
+		return nil, &types.AppError{Error: errors.New("failed to find user"),
 			Code: http.StatusInternalServerError}
 	}
 	if len(result) == 0 {
@@ -192,17 +189,16 @@ func (as *AuthService) LogIn(c *gin.Context) (*schemas.Message, *types.AppError)
 			return nil, &types.AppError{Error: errors.New("failed to create or update user"),
 				Code: http.StatusInternalServerError}
 		}
-	} else {
-		if err := as.Db.Model(&models.User{}).Where("user_id = ?", session.UserID).
-			Update("tg_session", session.Sesssion).Error; err != nil {
-			return nil, &types.AppError{Error: errors.New("failed to create or update user"),
-				Code: http.StatusInternalServerError}
-		}
 	}
 
 	setCookie(c, as.SessionCookieName, jweToken, as.SessionMaxAge)
 
-	database.KV.Set(kv.Key("sessions", hexToken), tokenBytes)
+	//create session
+
+	if err := as.Db.Create(&models.Session{UserId: session.UserID, Hash: hexToken, Session: session.Sesssion}).Error; err != nil {
+		return nil, &types.AppError{Error: errors.New("failed to create  user session"),
+			Code: http.StatusInternalServerError}
+	}
 
 	return &schemas.Message{Status: true, Message: "login success"}, nil
 }
@@ -254,7 +250,7 @@ func (as *AuthService) Logout(c *gin.Context) (*schemas.Message, *types.AppError
 	})
 
 	setCookie(c, as.SessionCookieName, "", -1)
-	database.KV.Delete(kv.Key("sessions", jwtUser.Hash))
+	as.Db.Where("session = ?", jwtUser.TgSession).Delete(&models.Session{})
 	return &schemas.Message{Status: true, Message: "logout success"}, nil
 }
 

services/common.go

@@ -12,7 +12,7 @@ import (
 	"github.com/divyam234/teldrive/schemas"
 	"github.com/divyam234/teldrive/types"
 	"github.com/divyam234/teldrive/utils"
-	"github.com/divyam234/teldrive/utils/kv"
+	"github.com/divyam234/teldrive/utils/cache"
 	"github.com/divyam234/teldrive/utils/tgc"
 	"github.com/gin-gonic/gin"
 	"github.com/gotd/td/telegram"
@@ -21,6 +21,12 @@ import (
 	"github.com/thoas/go-funk"
 )
 
+func CopyList[T any](list []T) []T {
+	newList := make([]T, len(list))
+	copy(newList, list)
+	return newList
+}
+
 func getChunk(ctx context.Context, tgClient *telegram.Client, location tg.InputFileLocationClass, offset int64, limit int64) ([]byte, error) {
 
 	req := &tg.UploadGetFileRequest{
@@ -84,6 +90,16 @@ func getBotInfo(ctx context.Context, token string) (*BotInfo, error) {
 
 func getParts(ctx context.Context, client *telegram.Client, file *schemas.FileOutFull, userID string) ([]types.Part, error) {
 
+	parts := []types.Part{}
+
+	key := fmt.Sprintf("messages:%s:%s", file.ID, userID)
+
+	err := cache.GetCache().Get(key, &parts)
+
+	if err == nil {
+		return parts, nil
+	}
+
 	ids := funk.Map(*file.Parts, func(part models.Part) tg.InputMessageClass {
 		return tg.InputMessageClass(&tg.InputMessageID{ID: int(part.ID)})
 	})
@@ -104,15 +120,14 @@ func getParts(ctx context.Context, client *telegram.Client, file *schemas.FileOu
 
 	messages := res.(*tg.MessagesChannelMessages)
 
-	parts := []types.Part{}
-
 	for _, message := range messages.Messages {
 		item := message.(*tg.Message)
 		media := item.Media.(*tg.MessageMediaDocument)
 		document := media.Document.(*tg.Document)
 		location := document.AsInputDocumentFileLocation()
 		parts = append(parts, types.Part{Location: location, Start: 0, End: document.Size - 1, Size: document.Size})
 	}
+	cache.GetCache().Set(key, &parts, 3600)
 	return parts, nil
 }
 
@@ -124,7 +139,7 @@ func rangedParts(parts []types.Part, start, end int64) []types.Part {
 
 	endPartNumber := int64(math.Ceil(float64(end) / float64(chunkSize)))
 
-	partsToDownload := parts[startPartNumber:endPartNumber]
+	partsToDownload := CopyList(parts[startPartNumber:endPartNumber])
 	partsToDownload[0].Start = start % chunkSize
 	partsToDownload[len(partsToDownload)-1].End = end % chunkSize
 
@@ -159,19 +174,21 @@ func GetDefaultChannel(ctx context.Context, userID int64) (int64, error) {
 
 	var channelID int64
 
-	key := kv.Key("users", strconv.FormatInt(userID, 10), "channel")
+	key := fmt.Sprintf("users:channel:%d", userID)
 
-	err := kv.GetValue(database.KV, key, &channelID)
+	err := cache.GetCache().Get(key, &channelID)
 
-	if err != nil {
-		var channelIds []int64
-		database.DB.Model(&models.Channel{}).Where("user_id = ?", userID).Where("selected = ?", true).
-			Pluck("channel_id", &channelIds)
+	if err == nil {
+		return channelID, nil
+	}
 
-		if len(channelIds) == 1 {
-			channelID = channelIds[0]
-			kv.SetValue(database.KV, key, &channelID)
-		}
+	var channelIds []int64
+	database.DB.Model(&models.Channel{}).Where("user_id = ?", userID).Where("selected = ?", true).
+		Pluck("channel_id", &channelIds)
+
+	if len(channelIds) == 1 {
+		channelID = channelIds[0]
+		cache.GetCache().Set(key, channelID, 0)
 	}
 
 	if channelID == 0 {
@@ -190,18 +207,42 @@ func GetBotsToken(ctx context.Context, userID int64) ([]string, error) {
 		return nil, err
 	}
 
-	key := kv.Key("users", strconv.FormatInt(userID, 10), strconv.FormatInt(channelId, 10), "bots")
+	key := fmt.Sprintf("users:bots:%d:%d", userID, channelId)
 
-	err = kv.GetValue(database.KV, key, &bots)
+	err = cache.GetCache().Get(key, &bots)
 
-	if err != nil {
-		if err := database.DB.Model(&models.Bot{}).Where("user_id = ?", userID).
-			Where("channel_id = ?", channelId).Pluck("token", &bots).Error; err != nil {
-			return nil, err
-		}
-		kv.SetValue(database.KV, key, &bots)
+	if err == nil {
+		return bots, nil
 	}
 
+	if err := database.DB.Model(&models.Bot{}).Where("user_id = ?", userID).
+		Where("channel_id = ?", channelId).Pluck("token", &bots).Error; err != nil {
+		return nil, err
+	}
+
+	cache.GetCache().Set(key, &bots, 0)
 	return bots, nil
 
 }
+
+func GetSessionByHash(hash string) (*models.Session, error) {
+
+	var session models.Session
+
+	key := fmt.Sprintf("sessions:%s", hash)
+
+	err := cache.GetCache().Get(key, &session)
+
+	if err == nil {
+		return &session, nil
+	}
+
+	if err := database.DB.Model(&models.Session{}).Where("hash = ?", hash).First(&session).Error; err != nil {
+		return nil, err
+	}
+
+	cache.GetCache().Set(key, &session, 0)
+
+	return &session, nil
+
+}