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" "time"
) )
// GetTimestamp get current timestamp in seconds
func GetTimestamp() int64 { func GetTimestamp() int64 {
return time.Now().Unix() return time.Now().Unix()
} }

View File

@ -1,8 +1,13 @@
package model package model
import ( import (
"fmt"
"math/rand"
"sync"
"github.com/pkg/errors" "github.com/pkg/errors"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger"
) )
type UserRequestCost struct { type UserRequestCost struct {
@ -25,8 +30,9 @@ func NewUserRequestCost(userID int, quotaID string, quota int64) *UserRequestCos
} }
func (docu *UserRequestCost) Insert() error { func (docu *UserRequestCost) Insert() error {
var err error go removeOldRequestCost()
err = DB.Create(docu).Error
err := DB.Create(docu).Error
return errors.Wrap(err, "failed to insert UserRequestCost") return errors.Wrap(err, "failed to insert UserRequestCost")
} }
@ -45,3 +51,25 @@ func GetCostByRequestId(reqid string) (*UserRequestCost, error) {
docu.CostUSD = float64(docu.Quota) / 500000 docu.CostUSD = float64(docu.Quota) / 500000
return docu, nil 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()))
}
}