feat: implement automatic removal of old request cost data

This commit is contained in:
Laisky.Cai 2024-12-22 12:23:32 +00:00
parent a9604fd59c
commit 3e82e76b03
2 changed files with 31 additions and 2 deletions

View File

@ -5,6 +5,7 @@ import (
"time"
)
// GetTimestamp get current timestamp in seconds
func GetTimestamp() int64 {
return time.Now().Unix()
}

View File

@ -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()))
}
}