From 3e82e76b033684502da1d50e9b3b858a3545694d Mon Sep 17 00:00:00 2001
From: "Laisky.Cai" <github@laisky.com>
Date: Sun, 22 Dec 2024 12:23:32 +0000
Subject: [PATCH] feat: implement automatic removal of old request cost data

---
 common/helper/time.go |  1 +
 model/cost.go         | 32 ++++++++++++++++++++++++++++++--
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/common/helper/time.go b/common/helper/time.go
index 302746db..b74cc74a 100644
--- a/common/helper/time.go
+++ b/common/helper/time.go
@@ -5,6 +5,7 @@ import (
 	"time"
 )
 
+// GetTimestamp get current timestamp in seconds
 func GetTimestamp() int64 {
 	return time.Now().Unix()
 }
diff --git a/model/cost.go b/model/cost.go
index 3d5d18ff..82efe855 100644
--- a/model/cost.go
+++ b/model/cost.go
@@ -1,8 +1,13 @@
 package model
 
 import (
+	"fmt"
+	"math/rand"
+	"sync"
+
 	"github.com/pkg/errors"
 	"github.com/songquanpeng/one-api/common/helper"
+	"github.com/songquanpeng/one-api/common/logger"
 )
 
 type UserRequestCost struct {
@@ -25,8 +30,9 @@ func NewUserRequestCost(userID int, quotaID string, quota int64) *UserRequestCos
 }
 
 func (docu *UserRequestCost) Insert() error {
-	var err error
-	err = DB.Create(docu).Error
+	go removeOldRequestCost()
+
+	err := DB.Create(docu).Error
 	return errors.Wrap(err, "failed to insert UserRequestCost")
 }
 
@@ -45,3 +51,25 @@ func GetCostByRequestId(reqid string) (*UserRequestCost, error) {
 	docu.CostUSD = float64(docu.Quota) / 500000
 	return docu, nil
 }
+
+var muRemoveOldRequestCost sync.Mutex
+
+// removeOldRequestCost remove old request cost data,
+// this function will be executed every 1/1000 times.
+func removeOldRequestCost() {
+	if rand.Float32() > 0.001 {
+		return
+	}
+
+	if ok := muRemoveOldRequestCost.TryLock(); !ok {
+		return
+	}
+	defer muRemoveOldRequestCost.Unlock()
+
+	err := DB.
+		Where("created_time < ?", helper.GetTimestamp()-3600*24*7).
+		Delete(&UserRequestCost{}).Error
+	if err != nil {
+		logger.SysError(fmt.Sprintf("failed to remove old request cost: %s", err.Error()))
+	}
+}